import os
from typing import Union
from scm import plams
import tcutility.log as log
from tcutility import spell_check
from tcutility.data import molecules
from tcutility.job.generic import Job
__all__ = ["CRESTJob", "QCGJob"]
j = os.path.join
[docs]
class CRESTJob(Job):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._crest_path = "crest"
self.options = {}
self._set_default_options()
def _set_default_options(self):
self.charge(0)
self.spin_polarization(0)
self.md_temperature(400)
self.md_length(1)
self.crest_path = "crest"
def _setup_job(self):
self.add_postamble("mkdir rotamers")
self.add_postamble("mkdir conformers")
self.add_postamble(f"split -d -l {len(self._molecule.atoms) + 2} -a 5 crest_conformers.xyz conformers/")
self.add_postamble(f"split -d -l {len(self._molecule.atoms) + 2} -a 5 crest_rotamers.xyz rotamers/")
self.add_postamble("for file in conformers/* rotamers/*")
self.add_postamble("do")
self.add_postamble(' mv "$file" "$file.xyz"')
self.add_postamble("done")
os.makedirs(self.workdir, exist_ok=True)
self._molecule.write(j(self.workdir, "coords.xyz"))
options = " ".join([f"{k} {v}" for k, v in self.options.items()])
options += " --noreftopo"
with open(self.runfile_path, "w+") as runf:
runf.write("#!/bin/sh\n\n")
runf.write("\n".join(self._preambles) + "\n\n")
runf.write(f"{self._crest_path} coords.xyz {options}\n")
runf.write("\n".join(self._postambles))
return True
@property
def crest_path(self) -> str:
return self._crest_path
@crest_path.setter
def crest_path(self, val: str):
self._crest_path = val
[docs]
def spin_polarization(self, val: int):
"""
Set the spin-polarization of the system.
"""
self.options["-uhf"] = val
[docs]
def multiplicity(self, val: int):
"""
Set the multiplicity of the system. If the value is not one the calculation will also be unrestricted.
We use the following values:
1) singlet
2) doublet
3) triplet
4) ...
The multiplicity is equal to 2*S+1 for spin-polarization of S.
"""
self.options["-uhf"] = (val - 1) // 2
[docs]
def charge(self, val: int):
"""
Set the charge of the system.
"""
self.options["-chrg"] = val
[docs]
def md_temperature(self, val: float):
"""
Set the temperature of the molecular dynamics steps. Defaults to 400K.
"""
self.options["-tnmd"] = val
[docs]
def md_length(self, val: float):
"""
Set the length of the molecular dynamics steps. The default length will be multiplied by this value, e.g. the default value is 1.
"""
self.options["-mdlen"] = f"x{val}"
[docs]
def solvent(self, name: Union[str, None] = None, model: str = "alpb"):
"""
Model solvation using the ALPB or GBSA model.
For available solvents see the `XTB documentation <https://xtb-docs.readthedocs.io/en/latest/gbsa.html#parameterized-solvents>`_.
Args:
name: the name of the solvent you want to use. If ``None`` turns off solvation.
model: the name of the model to use. Must be ``alpb`` or ``gbsa``. Defaults to ``alpb``.
"""
if name is None:
self.options.pop("-g", None)
self.options.pop("-alpb", None)
return
spell_check.check(model, ["alpb", "gbsa"], ignore_case=True)
_alpb_solvents = [
"acetone",
"acetonitrile",
"aniline",
"benzaldehyde",
"benzene",
"ch2cl2",
"chcl3",
"cs2",
"dioxane",
"dmf",
"dmso",
"ether",
"ethylacetate",
"furane",
"hexandecane",
"hexane",
"methanol",
"nitromethane",
"octanol",
"woctanol",
"phenol",
"toluene",
"thf",
"water",
]
_gbsa_solvents = [
"acetone",
"acetonitrile",
"CH2Cl2",
"CHCl3",
"CS2",
"DMF",
"DMSO",
"ether",
"H2O",
"methanol",
"n-hexane",
"THF",
"toluene",
]
if model.lower() == "alpb":
spell_check.check(name, _alpb_solvents, ignore_case=True)
self.options.pop("-g", None)
self.options["-alpb"] = name
else:
spell_check.check(name, _gbsa_solvents, ignore_case=True)
self.options.pop("-alpb", None)
self.options["-g"] = name
@property
def best_conformer_path(self):
return j(self.workdir, "crest_best.xyz")
@property
def conformer_directory(self):
return j(self.workdir, "conformers")
@property
def rotamer_directory(self):
return j(self.workdir, "rotamers")
[docs]
def get_rotamer_xyz(self, number: Union[int, None] = None):
"""
Return paths to rotamer xyz files for this job.
Args:
number: the number of files to return, defaults to 10. If the directory already exists, for example if the job was already run, we will return up to `number` files.
"""
if os.path.exists(self.rotamer_directory):
return [j(self.rotamer_directory, file) for i, file in enumerate(os.listdir(self.rotamer_directory))]
for i in range(number or 10):
yield j(self.rotamer_directory, f"{str(i).zfill(5)}.xyz")
[docs]
class QCGJob(CRESTJob):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._nsolv = 10
self._fix_solute = False
self._ensemble_generation_mode = "NCI-MTD"
self._solvent = None
self._nofix = True
[docs]
def solvent(self, name: Union[str, None] = None, model: str = "alpb", nsolv: Union[int, None] = None, mol: Union[str, plams.Molecule, plams.Atom, list, None] = None):
if nsolv:
self._nsolv = nsolv
if isinstance(mol, plams.Molecule):
self._solvent = mol
elif isinstance(mol, str) and os.path.exists(mol):
self._solvent = plams.Molecule(mol)
elif isinstance(mol, str):
# here we can get the molecule name using tcutility.data.molecules
self._solvent = molecules.get(mol)
# we can check if the alpb solvent name is defined in the xyz file
alpb = self._solvent.flags.get("alpb", None)
if alpb:
self.alpb(alpb)
elif isinstance(mol, list) and isinstance(mol[0], plams.Atom):
self._solvent = plams.Molecule()
[self._solvent.add_atom(atom) for atom in mol]
elif isinstance(mol, plams.Atom):
self._solvent = plams.Molecule()
self._solvent.add_atom(mol)
[docs]
def nsolv(self, nsolv):
self._nsolv = nsolv
[docs]
def alpb(self, solvent):
self._alpb = solvent
[docs]
def ensemble_mode(self, mode):
self._ensemble_generation_mode = mode
[docs]
def nofix(self, enable=True):
self._nofix = enable
def _setup_job(self):
self.add_postamble("mkdir ensemble/ensemble")
self.add_postamble(f"split -d -l {len(self._molecule.atoms) + len(self._solvent.atoms) * self._nsolv + 2} -a 5 ensemble/final_ensemble.xyz ensemble/ensemble")
self.add_postamble("for file in ensemble/ensemble/*")
self.add_postamble("do")
self.add_postamble(' mv "$file" "$file.xyz"')
self.add_postamble("done")
if not self._solvent:
log.error(f"Did not provide a solvent molecule for this job. Call the {self.__class__.__name__}.solvent method to add one.")
os.makedirs(self.workdir, exist_ok=True)
self._molecule.write(j(self.workdir, "coords.xyz"))
self._solvent.write(j(self.workdir, "solvent.xyz"))
ensemble_mode_option = {
"NCI-MTD": "--ncimtd",
"MD": "--md",
"MTD": "--mtd",
}.get(self._ensemble_generation_mode, self._ensemble_generation_mode)
options = [
"coords.xyz",
f'-xnam "{self.xtb_path}"',
"--qcg solvent.xyz",
"--ensemble",
f"--nsolv {self._nsolv}",
f"--chrg {self._charge}",
f"--uhf {self._spinpol}",
f"--tnmd {self._temp}",
f"--mdlen {50 * float(self._mdlen[1:])}",
ensemble_mode_option,
]
if self._alpb:
options.append(f"--alpb {self._alpb}")
if self._nofix:
options.append("--nofix")
options = " ".join(options)
with open(self.runfile_path, "w+") as runf:
runf.write("#!/bin/sh\n\n")
runf.write("\n".join(self._preambles) + "\n\n")
runf.write(f"{self._crest_path} {options}\n")
runf.write("\n".join(self._postambles))
return True
@property
def ensemble_directory(self):
return j(self.workdir, "ensemble", "ensemble")
[docs]
def get_ensemble_xyz(self, number: Union[int, None] = None):
"""
Return paths to conformer xyz files for this job.
Args:
number: the number of files to return, defaults to 10. If the directory already exists, for example if the job was already run, we will return up to `number` files.
"""
if os.path.exists(self.ensemble_directory):
return [j(self.ensemble_directory, file) for i, file in enumerate(sorted(os.listdir(self.ensemble_directory)))]
for i in range(number or 10):
yield j(self.ensemble_directory, f"{str(i).zfill(5)}.xyz")
@property
def best_ensemble_path(self):
return j(self.workdir, "ensemble", "ensemble", "crest_best.xyz")
if __name__ == "__main__":
with CRESTJob(test_mode=True) as job:
job.rundir = "tmp/SN2"
job.name = "CREST"
# job.molecule('../../../test/fixtures/xyz/transitionstate_radical_addition.xyz')
# job.solvent('water')
job.crest_path = "crest"
# job.sbatch(p='tc', ntasks_per_node=32)
# with QCGJob(test_mode=True) as job:
# job.rundir = 'calculations/Ammonia'
# job.name = 'QCG'
# job.molecule('ammonia.xyz')
# job.solvent('water', 10)
# print(job._solvent)
# job.sbatch(p='tc', n=32)
# for i in range(40):
# with CRESTJob(test_mode=False, overwrite=True) as job:
# job.rundir = 'tmp/SN2'
# job.name = 'CREST'
# job.molecule('../../../test/fixtures/xyz/transitionstate_radical_addition.xyz')
# job.sbatch(p='tc', ntasks_per_node=32)
TCutility