diff --git a/architector/io_samplers.py b/architector/io_samplers.py index 9a572aa..57af6a1 100644 --- a/architector/io_samplers.py +++ b/architector/io_samplers.py @@ -6,10 +6,79 @@ from scipy.stats import bernoulli from openbabel import openbabel from architector.io_calc import CalcExecutor +from ase.md.langevin import Langevin +from ase.io.trajectory import Trajectory import architector.arch_context_manage as arch_context_manage from architector.io_molecule import convert_io_molecule from architector.io_obabel import (convert_mol2_obmol,convert_obmol_ase) +def md_sampler(relaxed_mol, temp=298.15, interval=50, n=50, warm_up=1000, + timestep=1.0, friction=0.02, debug=False, return_energies=False): + """md_sampler + Use langevin dynamics and specified temperature to sample structures. + + Parameters + ---------- + relaxed_mol : architector.io_molecule.Molecule + relaxed molcule + temp : float, optional + temperature, by default 298.15 + interval : int, optional + time interval to save sampled structures, by default 50 + n : int, optional + Number of samples to save, by default 50 + warm_up : int, optional + Warm-up interval from 0K, by default 1000 + timestep : float, optional + timestep in fs, by default 1.0 + friction : float, optional + langevin friction, by default 0.02 + debug : bool, optional + Print out stuff, by default False + return_energies : bool, optional + give back energies and rmsds, by default False + """ + # Friction increased to get to convergence faster. + mol2 = relaxed_mol.write_mol2('init.mol2', writestring=True) + init_ase = convert_io_molecule(mol2).ase_atoms + relaxed_atoms = relaxed_mol.ase_atoms + skip_n = int(warm_up/interval) + good = True + with arch_context_manage.make_temp_directory() as _: + dyn = Langevin(relaxed_atoms, timestep * ase.units.fs, temp* ase.units.kB, friction=friction) + def printenergy(a=relaxed_atoms): # store a reference to atoms in the definition. + """Function to print the potential, kinetic and total energy.""" + epot = a.get_potential_energy() / len(a) + ekin = a.get_kinetic_energy() / len(a) + print('Energy per atom: Epot = %.3feV Ekin = %.3feV (T=%3.0fK) ' + 'Etot = %.3feV' % (epot, ekin, ekin / (1.5 * ase.units.kB), epot + ekin)) + if debug: + dyn.attach(printenergy, interval=interval) + traj = Trajectory('moldyn3.traj', 'w', relaxed_atoms) + dyn.attach(traj.write, interval=interval) + # Now run the dynamics + dyn.run(warm_up + n*interval) + traj = Trajectory('moldyn3.traj') + trunc_traj = traj[skip_n:] + displaced_structures = [] + energies = [] + rmsds = [] + for image in trunc_traj: + tmpmol = convert_io_molecule(mol2) + tmpmol.ase_atoms = image + displaced_structures.append(tmpmol) + _,rmsd = reorder_align(init_ase,image,return_rmsd=True) + energies.append(image.calc.results['energy']) + rmsds.append(rmsd) + if good and return_energies: + return displaced_structures,energies,rmsds + elif return_energies: + return [],energies,rmsds + elif good: + return displaced_structures + else: + return [] + def bond_length_sampler(relaxed_mol, n=10, seed=42, max_dev=0.4, smallest_dist_cutoff=0.55, min_dist_cutoff=3,