def init_md(self, name, time_step=0.5, temp_init=300, temp_bath=None, reset=False, interval=1):
"""
Initialize an ase molecular dynamics trajectory. The logfile needs to be specifies, so that old trajectories
are not overwritten. This functionality can be used to subsequently carry out equilibration and production.
Args:
name (str): Basic name of logfile and trajectory
time_step (float): Time step in fs (default=0.5)
temp_init (float): Initial temperature of the system in K (default is 300)
temp_bath (float): Carry out Langevin NVT dynamics at the specified temperature. If set to None, NVE
dynamics are performed instead (default=None)
reset (bool): Whether dynamics should be restarted with new initial conditions (default=False)
interval (int): Data is stored every interval steps (default=1)
"""
# If a previous dynamics run has been performed, don't reinitialize velocities unless explicitely requested
# via restart=True
if not self.dynamics or reset:
self._init_velocities(temp_init=temp_init)
# Set up dynamics
if temp_bath is None:
self.dynamics = VelocityVerlet(self.molecule, time_step * units.fs)
else:
self.dynamics = Langevin(self.molecule, time_step * units.fs, temp_bath * units.kB, 0.01)
# Create monitors for logfile and a trajectory file
logfile = os.path.join(self.working_dir, "%s.log" % name)
trajfile = os.path.join(self.working_dir, "%s.traj" % name)
logger = MDLogger(self.dynamics, self.molecule, logfile, stress=False, peratom=False, header=True, mode='a')
trajectory = Trajectory(trajfile, 'w', self.molecule)
# Attach monitors to trajectory
self.dynamics.attach(logger, interval=interval)
self.dynamics.attach(trajectory.write, interval=interval)
def __call__(self, atoms):
self.set_atoms(atoms)
self.backup()
# Start med en temperatur!
dyn = Langevin(atoms, self.timestep, self.temp, self.friction)
dyn.run(self.steps)
def create_system(self,
name,
time_step=1.0,
temp=300,
temp_init=None,
restart=False,
store=1,
nvt=False,
friction=0.001):
"""
Parameters
-----------
name : str
Name for output files.
time_step : float, optional
Time step in fs for simulation.
temp : float, optional
Temperature in K for NVT simulation.
temp_init : float, optional
Optional different temperature for initialization than thermostate set at.
restart : bool, optional
Determines whether simulation is restarted or not,
determines whether new velocities are initialized.
store : int, optional
Frequency at which output is written to log files.
nvt : bool, optional
Determines whether to run NVT simulation, default is False.
friction : float, optional
friction coefficient in fs^-1 for Langevin integrator
"""
if temp_init is None: temp_init = temp
if not self.md or restart:
MaxwellBoltzmannDistribution(self.mol, temp_init * units.kB)
if not nvt:
self.md = VelocityVerlet(self.mol, time_step * units.fs)
else:
self.md = Langevin(self.mol, time_step * units.fs, temp * units.kB,
friction / units.fs)
logfile = os.path.join(self.tmp, "{}.log".format(name))
trajfile = os.path.join(self.tmp, "{}.traj".format(name))
logger = MDLogger(self.md,
self.mol,
logfile,
stress=False,
peratom=False,
header=True,
mode="a")
trajectory = Trajectory(trajfile, "w", self.mol)
self.md.attach(logger, interval=store)
self.md.attach(trajectory.write, interval=store)
def test_CO2linear_Au111_langevin(testdir):
"""Test Langevin with constraints for rigid linear
triatomic molecules"""
rng = np.random.RandomState(0)
eref = 3.133526
zpos = cos(134.3 / 2.0 * pi / 180.0) * 1.197
xpos = sin(134.3 / 2.0 * pi / 180.0) * 1.19
co2 = Atoms('COO',
positions=[(-xpos + 1.2, 0, -zpos), (-xpos + 1.2, -1.1, -zpos),
(-xpos + 1.2, 1.1, -zpos)])
slab = fcc111('Au', size=(2, 2, 4), vacuum=2 * 5, orthogonal=True)
slab.center()
add_adsorbate(slab, co2, 1.5, 'bridge')
slab.set_pbc((True, True, False))
d0 = co2.get_distance(-3, -2)
d1 = co2.get_distance(-3, -1)
d2 = co2.get_distance(-2, -1)
calc = EMT()
slab.calc = calc
constraint = FixLinearTriatomic(triples=[(-2, -3, -1)])
slab.set_constraint(constraint)
fr = 0.1
dyn = Langevin(slab,
2.0 * units.fs,
temperature_K=300,
friction=fr,
trajectory='langevin_%.1f.traj' % fr,
logfile='langevin_%.1f.log' % fr,
loginterval=20,
rng=rng)
dyn.run(100)
# Check that the temperature is within a reasonable range
T = slab.get_temperature()
assert T > 100
assert T < 500
# Check that the constraints work
assert abs(slab.get_distance(-3, -2, mic=1) - d0) < 1e-9
assert abs(slab.get_distance(-3, -1, mic=1) - d1) < 1e-9
assert abs(slab.get_distance(-2, -1, mic=1) - d2) < 1e-9
# If the energy differs from the reference energy
# it is most probable that the redistribution of
# random forces in Langevin is not working properly
assert abs(slab.get_potential_energy() - eref) < 1e-4
def run(self, calc, filename):
slab = self.starting_geometry.copy()
slab.set_calculator(calc)
np.random.seed(1)
MaxwellBoltzmannDistribution(slab, self.temp * units.kB)
if self.ensemble == "NVE":
dyn = VelocityVerlet(slab, self.dt * units.fs)
elif self.ensemble == "nvtberendsen":
dyn = nvtberendsen.NVTBerendsen(slab,
self.dt * units.fs,
self.temp,
taut=300 * units.fs)
elif self.ensemble == "langevin":
dyn = Langevin(slab, self.dt * units.fs, self.temp * units.kB,
0.002)
traj = ase.io.Trajectory(filename + ".traj", "w", slab)
dyn.attach(traj.write, interval=1)
try:
fixed_atoms = len(slab.constraints[0].get_indices())
except:
fixed_atoms = 0
pass
def printenergy(a=slab):
"""Function to print( the potential, kinetic, and total energy)"""
epot = a.get_potential_energy() / len(a)
ekin = a.get_kinetic_energy() / (len(a) - fixed_atoms)
print("Energy per atom: Epot = %.3feV Ekin = %.3feV (T=%3.0fK) "
"Etot = %.3feV" % (epot, ekin, ekin /
(1.5 * units.kB), epot + ekin))
if printenergy:
dyn.attach(printenergy, interval=10)
dyn.run(self.count)
def generate_data(count, filename, temp, hook, cons_t=False):
"""Generates test or training data with a simple MD simulation."""
traj = ase.io.Trajectory(filename, "w")
slab = fcc100("Cu", size=(3, 3, 3))
ads = molecule("CO")
add_adsorbate(slab, ads, 5, offset=(1, 1))
cons = FixAtoms(indices=[
atom.index for atom in slab if (atom.tag == 2 or atom.tag == 3)
])
if hook:
cons2 = Hookean(a1=28, a2=27, rt=1.58, k=10.0)
slab.set_constraint([cons, cons2])
else:
slab.set_constraint(cons)
slab.center(vacuum=13., axis=2)
slab.set_pbc(True)
slab.wrap(pbc=[True] * 3)
slab.set_calculator(EMT())
slab.get_forces()
dyn = QuasiNewton(slab)
dyn.run(fmax=0.05)
traj.write(slab)
if cons_t is True:
dyn = Langevin(slab, 1.0 * units.fs, temp * units.kB, 0.002)
else:
dyn = VelocityVerlet(slab, dt=1.0 * units.fs)
for step in range(count):
dyn.run(20)
traj.write(slab)
def md_run(images, count, calc, filename, dir, temp, cons_t=False):
"""Generates test or training data with a simple MD simulation."""
log = Logger("results/logs/" + filename + ".txt")
traj = ase.io.Trajectory("".join([dir, filename, ".traj"]), "w")
slab = images[0].copy()
slab.set_calculator(calc)
slab.get_forces()
traj.write(slab)
if cons_t is True:
dyn = Langevin(slab, 1.0 * units.fs, temp * units.kB, 0.002)
else:
dyn = VelocityVerlet(slab, dt=1.0 * units.fs)
time_start = time.time()
for step in range(count):
dyn.run(20)
traj.write(slab)
time_elapsed = time.time() - time_start
log("MD Simulation Dynamics: %s" % dyn)
log("MD Simulation Time: %s \n" % time_elapsed)
def test_fixrotation_asap(asap3):
rng = np.random.RandomState(123)
with seterr(all='raise'):
atoms = bulk('Au', cubic=True).repeat((3, 3, 10))
atoms.pbc = False
atoms.center(vacuum=5.0 + np.max(atoms.cell) / 2)
print(atoms)
atoms.calc = asap3.EMT()
MaxwellBoltzmannDistribution(atoms, temperature_K=300, force_temp=True,
rng=rng)
Stationary(atoms)
check_inertia(atoms)
md = Langevin(
atoms,
timestep=20 * fs,
temperature_K=300,
friction=1e-3,
logfile='-',
loginterval=500,
rng=rng)
fx = FixRotation(atoms)
md.attach(fx)
md.run(steps=1000)
check_inertia(atoms)
def generate_data(count, filename, temp, hook, cons_t=False):
"""Generates test or training data with a simple MD simulation."""
traj = ase.io.Trajectory(filename, "w")
pair = molecule("CO")
cons = Hookean(a1=0, a2=1, rt=1.58, k=10.0)
# pair.set_constraint(cons)
pair.set_calculator(EMT())
pair.get_potential_energy()
dyn = QuasiNewton(pair, trajectory=(filename[:-5] + "_relax.traj"))
dyn.run(fmax=0.05)
traj.write(pair)
MaxwellBoltzmannDistribution(pair, temp * units.kB)
if cons_t is True:
dyn = Langevin(pair, 5 * units.fs, temp * units.kB, 0.002)
else:
dyn = VelocityVerlet(pair, dt=1.0 * units.fs)
for step in range(count - 1):
dyn.run(50)
traj.write(pair)
def test_langevin_asap(asap3):
with seterr(all='raise'):
rng = np.random.RandomState(0)
a = bulk('Au', cubic=True).repeat((4, 4, 4))
a.pbc = (False, False, False)
a.center(vacuum=2.0)
print(a)
a.calc = asap3.EMT()
# Set temperature to 10 K
MaxwellBoltzmannDistribution(a,
temperature_K=10,
force_temp=True,
rng=rng)
Stationary(a)
assert abs(a.get_temperature() - 10) < 0.0001
# Langevin dynamics should raise this to 300 K
T = 300
md = Langevin(a,
timestep=4 * fs,
temperature_K=T,
friction=0.01,
logfile='-',
loginterval=500,
rng=rng)
md.run(steps=3000)
# Now gather the temperature over 10000 timesteps, collecting it
# every 5 steps
temp = []
for i in range(1500):
md.run(steps=5)
temp.append(a.get_temperature())
temp = np.array(temp)
avgtemp = np.mean(temp)
fluct = np.std(temp)
print("Temperature is {:.2f} K +/- {:.2f} K".format(avgtemp, fluct))
assert abs(avgtemp - T) < 10.0
atoms.center()
atoms = atoms.repeat((3, 3, 3))
atoms.set_pbc(True)
# Set constraints for rigid triatomic molecules
nm = 27
atoms.constraints = FixLinearTriatomic(triples=[(3 * i, 3 * i + 1, 3 * i + 2)
for i in range(nm)])
tag = 'acn_27mol_300K'
atoms.calc = ACN(rc=np.min(np.diag(atoms.cell)) / 2)
# Create Langevin object
md = Langevin(atoms,
1 * units.fs,
temperature=300 * units.kB,
friction=0.01,
logfile=tag + '.log')
traj = Trajectory(tag + '.traj', 'w', atoms)
md.attach(traj.write, interval=1)
md.run(5000)
# Repeat box and equilibrate further
atoms.set_constraint()
atoms = atoms.repeat((2, 2, 2))
nm = 216
atoms.constraints = FixLinearTriatomic(triples=[(3 * i, 3 * i + 1, 3 * i + 2)
for i in range(nm)])
tag = 'acn_216mol_300K'
a.set_calculator(calc)
FIRE(a).run(fmax=fmax, steps=10000)
# Langevin quench to T1
print 'Langevin quench to {0}...'.format(T1)
Langevin(a, dt*fs, T1*kB, 1.0/(500*fs),
logfile='-', loginterval=int(100/dt)).run(int(time/dt))
# Langevin quench to T2
print 'Langevin quench to {0}...'.format(T2)
Langevin(a, dt*fs, T2*kB, 1.0/(500*fs),
logfile='-', loginterval=int(100/dt)).run(int(time/dt))
# Langevin quench to T3
print 'Langevin quench to {0}...'.format(T3)
dyn = Langevin(a, dt*fs, T3*kB, 1.0/(500*fs), logfile='-',
loginterval=int(100/dt))
dyn.run(int(time/dt))
# Collect pair distribution function
print 'Collect pair distribution function...'
p = PairAndAngleDistribution(a.get_calculator(), g2_cutoff, coord_cutoff,
npairbins=nbins, nanglebins=nbins)
dyn.attach(p, interval=int(100/dt))
dyn.run(int(time/dt))
print 'Writing files...'
# Write snapshot
write('rho_{0}.traj'.format(density), a)
# Write pair distribution function
add_adsorbate(slab, co2, 1.5, 'bridge')
slab.set_pbc((True, True, False))
d0 = co2.get_distance(-3, -2)
d1 = co2.get_distance(-3, -1)
d2 = co2.get_distance(-2, -1)
calc = EMT()
slab.set_calculator(calc)
constraint = FixLinearTriatomic(triples=[(-2, -3, -1)])
slab.set_constraint(constraint)
fr = 0.1
dyn = Langevin(slab,
2.0 * units.fs,
300 * units.kB,
fr,
trajectory='langevin_%.1f.traj' % fr,
logfile='langevin_%.1f.log' % fr,
loginterval=20,
rng=rng)
dyn.run(100)
# Check that the temperature is within a reasonable range
T = slab.get_temperature()
assert T > 100
assert T < 500
# Check that the constraints work
assert abs(slab.get_distance(-3, -2, mic=1) - d0) < 1e-9
assert abs(slab.get_distance(-3, -1, mic=1) - d1) < 1e-9
assert abs(slab.get_distance(-2, -1, mic=1) - d2) < 1e-9
# Global
# Stationary(atoms)
## Dynamics
# from ase.md.verlet import VelocityVerlet
# dyn = VelocityVerlet(
# atoms,
# d_t,
# trajectory = label+'.traj',
# logfile = 'log_'+label+'.txt',
# )
from ase.md import Langevin
dyn = Langevin(
atoms = atoms,
timestep = 10 *units.fs,
temperature = temp,
friction = friction,
trajectory = label+'.traj',
logfile = 'log-'+label+'.txt',
)
# from ase.md.npt import NPT
# dyn = NPT(
# atoms = atoms,
# timestep = d_t,
# temperature = temp,
# externalstress = 0.,
# ttime = 75 * units.fs,
# pfactor = (75. *units.fs)**2 * 100. *units.GPa,
# trajectory = label+'.traj',
# logfile = 'log_'+label+'.txt',
# )
### relax option
def run():
rng = np.random.RandomState(0)
a = Atoms('4X',
masses=[1, 2, 3, 4],
positions=[(0, 0, 0), (1, 0, 0), (0, 1, 0), (0.1, 0.2, 0.7)],
calculator=TstPotential())
print(a.get_forces())
# Langevin should reproduce Verlet if friction is 0.
md = Langevin(a,
0.5 * fs,
temperature_K=300,
friction=0.0,
logfile='-',
loginterval=500)
traj = Trajectory('4N.traj', 'w', a)
md.attach(traj, 100)
e0 = a.get_total_energy()
md.run(steps=2000)
del traj
assert abs(read('4N.traj').get_total_energy() - e0) < 0.0001
# Try again with nonzero friction.
md = Langevin(a,
0.5 * fs,
temperature_K=300,
friction=0.001,
logfile='-',
loginterval=500,
rng=rng)
traj = Trajectory('4NA.traj', 'w', a)
md.attach(traj, 100)
md.run(steps=2000)
# We cannot test the temperature without a lot of statistics.
# Asap does that. But if temperature is quite unreasonable,
# something is very wrong.
T = a.get_temperature()
assert T > 50
assert T < 1000
qn = QuasiNewton(a)
qn.run(0.001)
assert abs(a.get_potential_energy() - 1.0) < 0.000002
# Relax with the unscreened potential
print 'Relax with quick potential...'
a.set_calculator(quick_calc)
FIRE(a).run(fmax=fmax, steps=10000)
# Relax with the screened potential
print 'Relax with proper potential...'
a.set_calculator(calc)
FIRE(a).run(fmax=fmax, steps=10000)
# Langevin quench to T1
print 'Langevin quench to {0}...'.format(T1)
Langevin(a,
dt * fs,
T1 * kB,
1.0 / (500 * fs),
logfile='-',
loginterval=int(100 / dt)).run(int(time / dt))
# Langevin quench to T2
print 'Langevin quench to {0}...'.format(T2)
Langevin(a,
dt * fs,
T2 * kB,
1.0 / (500 * fs),
logfile='-',
loginterval=int(100 / dt)).run(int(time / dt))
# Langevin quench to T3
print 'Langevin quench to {0}...'.format(T3)
dyn = Langevin(a,
class ASE_MD:
"""
Setups and runs the MD simulation. Serves as an interface to the EVB Hamiltonian class and ASE.
"""
def __init__(self,
ase_atoms,
tmp,
calc_omm_d1,
calc_omm_d2,
calc_nn_d1,
calc_nn_d2,
off_diag,
shift=0):
"""
Parameters
-----------
ase_atoms : str
Location of input structure, gets created to ASE Atoms object.
tmp : str
Location for tmp directory.
calc_omm_d1 : Object
Contains OpenMM force field for diabat 1.
calc_omm_d2 : Object
Contains OpenMM force field for diabat 2.
calc_nn_d1 : Object
Diabat_NN instance for diabat 1.
calc_nn_d2 : Object
Diabat NN instance for diabat 2.
off_diag : PyTorch model
Model for predicting H12 energy and forces.
shift : float
Diabat 2 energy shift to diabat 1 reference.
"""
self.tmp = tmp
if not os.path.isdir(self.tmp):
os.makedirs(self.tmp)
self.mol = read(ase_atoms)
calculator = EVB_Hamiltonian(calc_omm_d1, calc_omm_d2, calc_nn_d1,
calc_nn_d2, off_diag, shift)
self.mol.set_calculator(calculator)
self.md = None
def create_system(self,
name,
time_step=1.0,
temp=300,
temp_init=None,
restart=False,
store=1,
nvt=False,
friction=0.001):
"""
Parameters
-----------
name : str
Name for output files.
time_step : float, optional
Time step in fs for simulation.
temp : float, optional
Temperature in K for NVT simulation.
temp_init : float, optional
Optional different temperature for initialization than thermostate set at.
restart : bool, optional
Determines whether simulation is restarted or not,
determines whether new velocities are initialized.
store : int, optional
Frequency at which output is written to log files.
nvt : bool, optional
Determines whether to run NVT simulation, default is False.
friction : float, optional
friction coefficient in fs^-1 for Langevin integrator
"""
if temp_init is None: temp_init = temp
if not self.md or restart:
MaxwellBoltzmannDistribution(self.mol, temp_init * units.kB)
if not nvt:
self.md = VelocityVerlet(self.mol, time_step * units.fs)
else:
self.md = Langevin(self.mol, time_step * units.fs, temp * units.kB,
friction / units.fs)
logfile = os.path.join(self.tmp, "{}.log".format(name))
trajfile = os.path.join(self.tmp, "{}.traj".format(name))
logger = MDLogger(self.md,
self.mol,
logfile,
stress=False,
peratom=False,
header=True,
mode="a")
trajectory = Trajectory(trajfile, "w", self.mol)
self.md.attach(logger, interval=store)
self.md.attach(trajectory.write, interval=store)
def write_mol(self, name, ftype="xyz", append=False):
"""
Write out current molecule structure.
Parameters
-----------
name : str
Name of the output file.
ftype : str, optional
Determines output file format, default xyz.
append : bool, optional
Append to existing output file or not.
"""
path = os.path.join(self.tmp, "{}.{}".format(name, ftype))
write(path, self.mol, format=ftype, append=append)
def run_md(self, steps):
"""
Run MD simulation.
Parameters
-----------
steps : int
Number of MD steps
"""
self.md.run(steps)
n = a.get_atomic_numbers().copy()
# Relax with the quick potential
a.set_atomic_numbers([6]*len(a))
a.set_calculator(quick_calc)
FIRE(a, downhill_check=True).run(fmax=1.0, steps=nsteps_relax)
a.set_atomic_numbers(n)
write(initial_fn, a)
else:
print('... reading %s ...' % liquid_fn)
a = read(liquid_fn)
# Thermalize with the slow (but correct) potential
a.set_calculator(calc)
traj = Trajectory(liquid_fn, 'a', a)
dyn = Langevin(a, dt1, T1, 1.0/tau1,
logfile='-', loginterval=int(dtdump/dt1))
dyn.attach(traj.write, interval=int(dtdump/dt1)) # every 100 fs
nsteps = int(teq/dt1)-len(traj)*int(dtdump/dt1)
print('Need to run for further {} steps to reach total of {} steps.'.format(nsteps, int(teq/dt1)))
if nsteps <= 0:
nsteps = 1
dyn.run(nsteps)
traj.close()
# Write snapshot
write(liquid_final_fn, a)
else:
print('... reading %s ...' % liquid_final_fn)
a = read(liquid_final_fn)
a.set_calculator(calc)
def save(smiles, species, coordinates):
print(len(species))
print(smiles)
for s, c in zip(species, coordinates):
print(s, *c)
smiles = parser.smiles
m = Chem.MolFromSmiles(smiles)
m = Chem.AddHs(m)
AllChem.EmbedMolecule(m, useRandomCoords=True)
AllChem.UFFOptimizeMolecule(m)
pos = m.GetConformer().GetPositions()
natoms = m.GetNumAtoms()
species = [m.GetAtomWithIdx(j).GetSymbol() for j in range(natoms)]
atoms = Atoms(species, positions=pos)
atoms.calc = EMT()
md = Langevin(atoms,
1 * units.fs,
temperature=parser.temperature * units.kB,
friction=0.01)
for _ in range(parser.conformations):
md.run(1)
positions = atoms.get_positions()
save(smiles, species, positions)
n = a.get_atomic_numbers().copy()
# Relax with the quick potential
a.set_atomic_numbers([6]*len(a))
a.set_calculator(quick_calc)
FIRE(a, downhill_check=True).run(fmax=1.0, steps=10000)
a.set_atomic_numbers(n)
write(initial_fn, a)
else:
print('... reading %s ...' % liquid_fn)
a = read(liquid_fn)
# Thermalize with the slow (but correct) potential
a.set_calculator(calc)
traj = Trajectory(liquid_fn, 'a', a)
dyn = Langevin(a, dt1, T1, 1.0/tau1,
logfile='-', loginterval=int(dtdump/dt1))
dyn.attach(traj.write, interval=int(dtdump/dt1)) # every 100 fs
nsteps = int(teq/dt1)-len(traj)*int(dtdump/dt1)
print('Need to run for further {} steps to reach total of {} steps.'.format(nsteps, int(teq/dt1)))
if nsteps <= 0:
nsteps = 1
dyn.run(nsteps)
traj.close()
# Write snapshot
write(liquid_final_fn, a)
else:
print('... reading %s ...' % liquid_final_fn)
a = read(liquid_final_fn)
a.set_calculator(calc)
from ase.md import Langevin
from ase.io import Trajectory, read
from ase.optimize import QuasiNewton
from ase.utils import seterr
with seterr(all='raise'):
a = Atoms('4X',
masses=[1, 2, 3, 4],
positions=[(0, 0, 0),
(1, 0, 0),
(0, 1, 0),
(0.1, 0.2, 0.7)],
calculator=TestPotential())
print(a.get_forces())
# Langevin should reproduce Verlet if friction is 0.
md = Langevin(a, 0.5 * fs, 300 * kB, 0.0, logfile='-', loginterval=500)
traj = Trajectory('4N.traj', 'w', a)
md.attach(traj, 100)
e0 = a.get_total_energy()
md.run(steps=10000)
del traj
assert abs(read('4N.traj').get_total_energy() - e0) < 0.0001
# Try again with nonzero friction.
md = Langevin(a, 0.5 * fs, 300 * kB, 0.001, logfile='-', loginterval=500)
traj = Trajectory('4NA.traj', 'w', a)
md.attach(traj, 100)
md.run(steps=10000)
# We cannot test the temperature without a lot of statistics.
# Asap does that. But if temperature is quite unreasonable,
class AseInterface:
"""
Interface for ASE calculations (optimization and molecular dynamics)
Args:
molecule_path (str): Path to initial geometry
ml_model (object): Trained model
working_dir (str): Path to directory where files should be stored
device (str): cpu or cuda
"""
def __init__(
self,
molecule_path,
ml_model,
working_dir,
device="cpu",
energy="energy",
forces="forces",
energy_units="eV",
forces_units="eV/Angstrom",
environment_provider=SimpleEnvironmentProvider(),
):
# Setup directory
self.working_dir = working_dir
if not os.path.exists(self.working_dir):
os.makedirs(self.working_dir)
# Load the molecule
self.molecule = None
self._load_molecule(molecule_path)
# Set up calculator
calculator = SpkCalculator(
ml_model,
device=device,
energy=energy,
forces=forces,
energy_units=energy_units,
forces_units=forces_units,
environment_provider=environment_provider,
)
self.molecule.set_calculator(calculator)
# Unless initialized, set dynamics to False
self.dynamics = False
def _load_molecule(self, molecule_path):
"""
Load molecule from file (can handle all ase formats).
Args:
molecule_path (str): Path to molecular geometry
"""
file_format = os.path.splitext(molecule_path)[-1]
if file_format == "xyz":
self.molecule = read_xyz(molecule_path)
else:
self.molecule = read(molecule_path)
def save_molecule(self, name, file_format="xyz", append=False):
"""
Save the current molecular geometry.
Args:
name (str): Name of save-file.
file_format (str): Format to store geometry (default xyz).
append (bool): If set to true, geometry is added to end of file
(default False).
"""
molecule_path = os.path.join(self.working_dir,
"%s.%s" % (name, file_format))
if file_format == "xyz":
# For extended xyz format, plain is needed since ase can not parse
# the extxyz it writes
write_xyz(molecule_path, self.molecule, plain=True)
else:
write(molecule_path,
self.molecule,
format=file_format,
append=append)
def calculate_single_point(self):
"""
Perform a single point computation of the energies and forces and
store them to the working directory. The format used is the extended
xyz format. This functionality is mainly intended to be used for
interfaces.
"""
energy = self.molecule.get_potential_energy()
forces = self.molecule.get_forces()
self.molecule.energy = energy
self.molecule.forces = forces
self.save_molecule("single_point", file_format="extxyz")
def init_md(
self,
name,
time_step=0.5,
temp_init=300,
temp_bath=None,
reset=False,
interval=1,
):
"""
Initialize an ase molecular dynamics trajectory. The logfile needs to
be specifies, so that old trajectories are not overwritten. This
functionality can be used to subsequently carry out equilibration and
production.
Args:
name (str): Basic name of logfile and trajectory
time_step (float): Time step in fs (default=0.5)
temp_init (float): Initial temperature of the system in K
(default is 300)
temp_bath (float): Carry out Langevin NVT dynamics at the specified
temperature. If set to None, NVE dynamics are performed
instead (default=None)
reset (bool): Whether dynamics should be restarted with new initial
conditions (default=False)
interval (int): Data is stored every interval steps (default=1)
"""
# If a previous dynamics run has been performed, don't reinitialize
# velocities unless explicitly requested via restart=True
if not self.dynamics or reset:
self._init_velocities(temp_init=temp_init)
# Set up dynamics
if temp_bath is None:
self.dynamics = VelocityVerlet(self.molecule, time_step * units.fs)
else:
self.dynamics = Langevin(
self.molecule,
time_step * units.fs,
temp_bath * units.kB,
1.0 / (100.0 * units.fs),
)
# Create monitors for logfile and a trajectory file
logfile = os.path.join(self.working_dir, "%s.log" % name)
trajfile = os.path.join(self.working_dir, "%s.traj" % name)
logger = MDLogger(
self.dynamics,
self.molecule,
logfile,
stress=False,
peratom=False,
header=True,
mode="a",
)
trajectory = Trajectory(trajfile, "w", self.molecule)
# Attach monitors to trajectory
self.dynamics.attach(logger, interval=interval)
self.dynamics.attach(trajectory.write, interval=interval)
def _init_velocities(self,
temp_init=300,
remove_translation=True,
remove_rotation=True):
"""
Initialize velocities for molecular dynamics
Args:
temp_init (float): Initial temperature in Kelvin (default 300)
remove_translation (bool): Remove translation components of
velocity (default True)
remove_rotation (bool): Remove rotation components of velocity
(default True)
"""
MaxwellBoltzmannDistribution(self.molecule, temp_init * units.kB)
if remove_translation:
Stationary(self.molecule)
if remove_rotation:
ZeroRotation(self.molecule)
def run_md(self, steps):
"""
Perform a molecular dynamics simulation using the settings specified
upon initializing the class.
Args:
steps (int): Number of simulation steps performed
"""
if not self.dynamics:
raise AttributeError("Dynamics need to be initialized using the"
" 'setup_md' function")
self.dynamics.run(steps)
def optimize(self, fmax=1.0e-2, steps=1000):
"""
Optimize a molecular geometry using the Quasi Newton optimizer in ase
(BFGS + line search)
Args:
fmax (float): Maximum residual force change (default 1.e-2)
steps (int): Maximum number of steps (default 1000)
"""
name = "optimization"
optimize_file = os.path.join(self.working_dir, name)
optimizer = QuasiNewton(
self.molecule,
trajectory="%s.traj" % optimize_file,
restart="%s.pkl" % optimize_file,
)
optimizer.run(fmax, steps)
# Save final geometry in xyz format
self.save_molecule(name)
def compute_normal_modes(self, write_jmol=True):
"""
Use ase calculator to compute numerical frequencies for the molecule
Args:
write_jmol (bool): Write frequencies to input file for
visualization in jmol (default=True)
"""
freq_file = os.path.join(self.working_dir, "normal_modes")
# Compute frequencies
frequencies = Vibrations(self.molecule, name=freq_file)
frequencies.run()
# Print a summary
frequencies.summary()
# Write jmol file if requested
if write_jmol:
frequencies.write_jmol()
from ase.calculators.test import TestPotential
from ase.md import Langevin
from ase.io import Trajectory, read
from ase.optimize import QuasiNewton
from ase.utils import seterr
rng = np.random.RandomState(0)
with seterr(all='raise'):
a = Atoms('4X',
masses=[1, 2, 3, 4],
positions=[(0, 0, 0), (1, 0, 0), (0, 1, 0), (0.1, 0.2, 0.7)],
calculator=TestPotential())
print(a.get_forces())
# Langevin should reproduce Verlet if friction is 0.
md = Langevin(a, 0.5 * fs, 300 * kB, 0.0, logfile='-', loginterval=500)
traj = Trajectory('4N.traj', 'w', a)
md.attach(traj, 100)
e0 = a.get_total_energy()
md.run(steps=10000)
del traj
assert abs(read('4N.traj').get_total_energy() - e0) < 0.0001
# Try again with nonzero friction.
md = Langevin(a,
0.5 * fs,
300 * kB,
0.001,
logfile='-',
loginterval=500,
rng=rng)