from pyiron import Project
from pyiron.base.generic.parameters import GenericParameters
from pyiron.atomistics.structure.atoms import Atoms
from pyiron.atomistics.job.atomistic import AtomisticGenericJob
from pyiron.base.settings.generic import Settings
from collections import OrderedDict
from yaff import System, log as Yafflog, ForceField
Yafflog.set_level(Yafflog.silent)
from quickff import read_abinitio
from quickff.tools import set_ffatypes
from quickff.settings import key_checks
from molmod.units import *
from molmod.constants import *
from molmod.io.chk import load_chk, dump_chk
from molmod.periodic import periodic as pt
import os, posixpath, numpy as np, h5py, matplotlib.pyplot as pp
[docs]def write_chk(input_dict, working_directory='.'):
# collect data and initialize Yaff system
if 'cell' in input_dict.keys() and input_dict['cell'] is not None:
system = System(input_dict['numbers'], input_dict['pos']*angstrom, rvecs=input_dict['cell']*angstrom, ffatypes=input_dict['ffatypes_man'], ffatype_ids=input_dict['ffatype_ids_man'])
else:
system = System(input_dict['numbers'], input_dict['pos']*angstrom, ffatypes=input_dict['ffatypes_man'], ffatype_ids=input_dict['ffatype_ids_man'])
# determine masses, bonds and ffaypes from ffatype_rules
system.detect_bonds()
system.set_standard_masses()
# write dictionnairy to MolMod CHK file
system.to_file(posixpath.join(working_directory,'input.chk'))
# Reload input.chk as dictionairy and add AI input data
d = load_chk(posixpath.join(working_directory,'input.chk'))
assert isinstance(input_dict['aiener'], float), "AI energy not defined in input, use job.read_abintio(...)"
assert isinstance(input_dict['aigrad'], np.ndarray), "AI gradient not defined in input, use job.read_abintio(...)"
assert isinstance(input_dict['aihess'], np.ndarray), "AI hessian not defined in input, use job.read_abintio(...)"
d['energy'] = input_dict['aiener']
d['grad'] = input_dict['aigrad']
d['hess'] = input_dict['aihess']
dump_chk(posixpath.join(working_directory,'input.chk'), d)
[docs]def write_pars(pars,fn,working_directory='.'):
with open(posixpath.join(working_directory,fn), 'w') as f:
for line in pars:
f.write(line)
[docs]def write_config(input_dict,working_directory='.'):
with open(posixpath.join(working_directory,'config.txt'), 'w') as f:
for key in key_checks.keys():
if key in input_dict.keys():
value = str(input_dict[key])
if key=='ffatypes': assert value == 'None'
print('%s: %s' %(key+' '*(30-len(key)), value), file=f)
[docs]def collect_output(fn_pars, fn_sys):
# this routine reads the output parameter file containing the covalent pars
output_dict = {'generic/bond': [], 'generic/bend': [], 'generic/torsion': [], 'generic/oopdist': [], 'generic/cross': []}
kinds = ['bond', 'bend', 'torsion', 'oopdist', 'cross']
with open(fn_pars, 'r') as f:
for line in f.readlines():
for key in kinds:
if key in line.lower():
output_dict['generic/%s' %key].append(line)
system = System.from_file(fn_sys)
output_dict['system/numbers'] = system.numbers
output_dict['system/pos'] = system.pos/angstrom
if system.cell is not None:
output_dict['system/rvecs'] = system.cell.rvecs/angstrom
output_dict['system/bonds'] = system.bonds
output_dict['system/ffatypes'] = np.asarray(system.ffatypes,'S22')
output_dict['system/ffatype_ids'] = system.ffatype_ids
return output_dict
[docs]class QuickFF(AtomisticGenericJob):
def __init__(self, project, job_name):
super(QuickFF, self).__init__(project, job_name)
self.__name__ = "QuickFF"
self._executable_activate(enforce=True)
self.input = QuickFFInput()
self.ffatypes = None
self.ffatype_ids = None
self.aiener = None
self.aigrad = None
self.aihess = None
self.fn_ei = None
self.fn_vdw = None
[docs] def read_abinitio(self, fn):
numbers, coords, energy, grad, hess, masses, rvecs, pbc = read_abinitio(fn)
coords /= angstrom
if rvecs is not None:
rvecs /= angstrom
self.structure = Atoms(numbers=numbers, positions=coords, cell=rvecs)
self.aiener = energy
self.aigrad = grad
self.aihess = hess
[docs] def detect_ffatypes(self, ffatypes=None, ffatype_rules=None, ffatype_level=None):
'''
Define atom types by explicitely giving them through the
ffatypes keyword, defining atype rules using the ATSELECT
language implemented in Yaff (see the Yaff documentation at
http://molmod.github.io/yaff/ug_atselect.html) or by specifying
the ffatype_level employing the built-in routine in QuickFF.
'''
numbers = np.array([pt[symbol].number for symbol in self.structure.get_chemical_symbols()])
if self.structure.cell is None:
system = System(numbers, self.structure.positions.copy()*angstrom)
else:
system = System(numbers, self.structure.positions.copy()*angstrom, rvecs=self.structure.cell*angstrom)
system.detect_bonds()
if not sum([ffatypes is None, ffatype_rules is None, ffatype_level is None]) == 2:
raise IOError('Exactly one of ffatypes, ffatype_rules and ffatype_level should be defined')
if ffatypes is not None:
assert ffatype_rules is None, 'ffatypes and ffatype_rules cannot be defined both'
system.ffatypes = ffatypes
system.ffatype_ids = None
system._init_derived_ffatypes()
if ffatype_rules is not None:
system.detect_ffatypes(ffatype_rules)
if ffatype_level is not None:
set_ffatypes(system, ffatype_level)
self.ffatypes = system.ffatypes.copy()
self.ffatype_ids = system.ffatype_ids.copy()
[docs] def set_ei(self, fn):
self.input['ei'] = fn.split('/')[-1]
self.fn_ei = fn
[docs] def set_vdw(self, fn):
self.input['vdw'] = fn.split('/')[-1]
self.fn_vdw = fn
[docs] def collect_output(self):
output_dict = collect_output(posixpath.join(self.working_directory, self.input['fn_yaff']), posixpath.join(self.working_directory, self.input['fn_sys']))
with self.project_hdf5.open("output") as hdf5_output:
for k, v in output_dict.items():
hdf5_output[k] = v
[docs] def to_hdf(self, hdf=None, group_name=None):
super(QuickFF, self).to_hdf(hdf=hdf, group_name=group_name)
with self.project_hdf5.open("input") as hdf5_input:
self.structure.to_hdf(hdf5_input)
self.input.to_hdf(hdf5_input)
[docs] def from_hdf(self, hdf=None, group_name=None):
super(QuickFF, self).from_hdf(hdf=hdf, group_name=group_name)
with self.project_hdf5.open("input") as hdf5_input:
self.input.from_hdf(hdf5_input)
self.structure = Atoms().from_hdf(hdf5_input)
[docs] def get_structure(self, iteration_step=-1, wrap_atoms=True):
"""
Overwrite the get_structure routine from AtomisticGenericJob because we want to avoid
defining a unit cell when one does not exist
"""
raise NotImplementedError
[docs] def log(self):
with open(posixpath.join(self.working_directory, 'quickff.log')) as f:
print(f.read())
[docs] def get_yaff_system(self):
system = System.from_file(posixpath.join(self.working_directory, self.input['fn_sys']))
return system
[docs] def get_yaff_ff(self, rcut=15*angstrom, alpha_scale=3.2, gcut_scale=1.5, smooth_ei=True):
system = self.get_yaff_system()
fn_pars = posixpath.join(self.working_directory, self.input['fn_yaff'])
if not os.path.isfile(fn_pars):
raise IOError('No pars.txt file find in job working directory. Have you already run the job?')
ff = ForceField.generate(
system, fn_pars, rcut=rcut, alpha_scale=alpha_scale,
gcut_scale=gcut_scale, smooth_ei=smooth_ei
)
return ff