from scm import plams
from tcutility.data import molecules
from tcutility.job.generic import Job
from tcutility import log
import os
j = os.path.join
[docs]
class CRESTJob(Job):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.crest_path = 'crest'
self.xtb_path = 'xtb'
self._charge = 0
self._spinpol = 0
self._temp = 400
self._mdlen = 'x1'
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 = [
'coords.xyz',
f'-xnam "{self.xtb_path}"',
'--noreftopo',
f'-c {self._charge}',
f'-u {self._spinpol}',
f'-tnmd {self._temp}',
f'-mdlen {self._mdlen}',
]
options = ' '.join(options)
with open(self.runfile_path, 'w+') as runf:
runf.write('#!/bin/sh\n\n') # the shebang is not written by default by ADF
runf.write('\n'.join(self._preambles) + '\n\n')
runf.write(f'{self.crest_path} {options}\n')
runf.write('\n'.join(self._postambles))
return True
[docs]
def spin_polarization(self, val: int):
'''
Set the spin-polarization of the system.
'''
self._spinpol = 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._spinpol = (val - 1)//2
[docs]
def charge(self, val: int):
'''
Set the charge of the system.
'''
self._charge = val
[docs]
def md_temperature(self, val: float):
'''
Set the temperature of the molecular dynamics steps. Defaults to 400K.
'''
self._temp = 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._mdlen = f'x{val}'
@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: int = 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.crest_path = 'crest'
self._nsolv = 10
self._fix_solute = False
self._ensemble_generation_mode = 'NCI-MTD'
self._solvent = None
self._nofix = True
[docs]
def solvent(self, mol, nsolv=None):
if self._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: int = 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() 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)
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