from os.path import exists
import numpy as np
from ase.calculators.calculator import Calculator, all_changes
from ase.io.trajectory import Trajectory
from ase.parallel import broadcast, world
from ase.units import fs, kJ, mol, nm
def restart_from_trajectory(prev_traj, *args, prev_steps=None, atoms=None,
**kwargs):
"""This function helps the user to restart a plumed simulation
from a trajectory file.
Parameters
----------
prev_traj : Trajectory object
previous simulated trajectory
prev_steps : int. Default steps in prev_traj.
number of previous steps
others :
Same parameters of :mod:`~ase.calculators.plumed` calculator
Returns
-------
Plumed calculator
.. note:: prev_steps is crucial when trajectory does not contain
all the previous steps.
"""
atoms.calc = Plumed(*args, atoms=atoms, restart=True, **kwargs)
with Trajectory(prev_traj) as traj:
if prev_steps is None:
atoms.calc.istep = len(traj) - 1
else:
atoms.calc.istep = prev_steps
atoms.set_positions(traj[-1].get_positions())
atoms.set_momenta(traj[-1].get_momenta())
return atoms.calc
[docs]class Plumed(Calculator):
implemented_properties = ['energy', 'forces']
def __init__(self, calc, input, timestep, atoms=None, kT=1., log='',
restart=False, use_charge=False, update_charge=False):
"""
Plumed calculator is used for simulations of enhanced sampling methods
with the open-source code PLUMED (plumed.org).
[1] The PLUMED consortium, Nat. Methods 16, 670 (2019)
[2] Tribello, Bonomi, Branduardi, Camilloni, and Bussi,
Comput. Phys. Commun. 185, 604 (2014)
Parameters
----------
calc: Calculator object
It computes the unbiased forces
input: List of strings
It contains the setup of plumed actions
timestep: float
Time step of the simulated dynamics
atoms: Atoms
Atoms object to be attached
kT: float. Default 1.
Value of the thermal energy in eV units. It is important for
some methods of plumed like Well-Tempered Metadynamics.
log: string
Log file of the plumed calculations
restart: boolean. Default False
True if the simulation is restarted.
use_charge: boolean. Default False
True if you use some collective variable which needs charges. If
use_charges is True and update_charge is False, you have to define
initial charges and then this charge will be used during all
simulation.
update_charge: boolean. Default False
True if you want the charges to be updated each time step. This
will fail in case that calc does not have 'charges' in its
properties.
.. note:: For this case, the calculator is defined strictly with the
object atoms inside. This is necessary for initializing the
Plumed object. For conserving ASE convention, it can be initialized
as atoms.calc = (..., atoms=atoms, ...)
.. note:: In order to guarantee a proper restart, the user has to fix
momenta, positions and Plumed.istep, where the positions and
momenta corresponds to the last configuration in the previous
simulation, while Plumed.istep is the number of timesteps
performed previously. This can be done using
ase.calculators.plumed.restart_from_trajectory.
"""
from plumed import Plumed as pl
if atoms is None:
raise TypeError('plumed calculator has to be defined with the \
object atoms inside.')
self.istep = 0
Calculator.__init__(self, atoms=atoms)
self.input = input
self.calc = calc
self.use_charge = use_charge
self.update_charge = update_charge
if world.rank == 0:
natoms = len(atoms.get_positions())
self.plumed = pl()
''' Units setup
warning: inputs and outputs of plumed will still be in
plumed units.
The change of Plumed units to ASE units is:
kjoule/mol to eV
nm to Angstrom
ps to ASE time units
ASE and plumed - charge unit is in e units
ASE and plumed - mass unit is in a.m.u units '''
ps = 1000 * fs
self.plumed.cmd("setMDEnergyUnits", mol / kJ)
self.plumed.cmd("setMDLengthUnits", 1 / nm)
self.plumed.cmd("setMDTimeUnits", 1 / ps)
self.plumed.cmd("setMDChargeUnits", 1.)
self.plumed.cmd("setMDMassUnits", 1.)
self.plumed.cmd("setNatoms", natoms)
self.plumed.cmd("setMDEngine", "ASE")
self.plumed.cmd("setLogFile", log)
self.plumed.cmd("setTimestep", float(timestep))
self.plumed.cmd("setRestart", restart)
self.plumed.cmd("setKbT", float(kT))
self.plumed.cmd("init")
for line in input:
self.plumed.cmd("readInputLine", line)
self.atoms = atoms
def _get_name(self):
return f'{self.calc.name}+Plumed'
def calculate(self, atoms=None, properties=['energy', 'forces'],
system_changes=all_changes):
Calculator.calculate(self, atoms, properties, system_changes)
comp = self.compute_energy_and_forces(self.atoms.get_positions(),
self.istep)
energy, forces = comp
self.istep += 1
self.results['energy'], self. results['forces'] = energy, forces
def compute_energy_and_forces(self, pos, istep):
unbiased_energy = self.calc.get_potential_energy(self.atoms)
unbiased_forces = self.calc.get_forces(self.atoms)
if world.rank == 0:
ener_forc = self.compute_bias(pos, istep, unbiased_energy)
else:
ener_forc = None
energy_bias, forces_bias = broadcast(ener_forc)
energy = unbiased_energy + energy_bias
forces = unbiased_forces + forces_bias
return energy, forces
def compute_bias(self, pos, istep, unbiased_energy):
self.plumed.cmd("setStep", istep)
if self.use_charge:
if 'charges' in self.calc.implemented_properties and \
self.update_charge:
charges = self.calc.get_charges(atoms=self.atoms.copy())
elif self.atoms.has('initial_charges') and not self.update_charge:
charges = self.atoms.get_initial_charges()
else:
assert not self.update_charge, "Charges cannot be updated"
assert self.update_charge, "Not initial charges in Atoms"
self.plumed.cmd("setCharges", charges)
# Box for functions with PBC in plumed
if self.atoms.cell:
cell = np.asarray(self.atoms.get_cell())
self.plumed.cmd("setBox", cell)
self.plumed.cmd("setPositions", pos)
self.plumed.cmd("setEnergy", unbiased_energy)
self.plumed.cmd("setMasses", self.atoms.get_masses())
forces_bias = np.zeros((self.atoms.get_positions()).shape)
self.plumed.cmd("setForces", forces_bias)
virial = np.zeros((3, 3))
self.plumed.cmd("setVirial", virial)
self.plumed.cmd("prepareCalc")
self.plumed.cmd("performCalc")
energy_bias = np.zeros((1,))
self.plumed.cmd("getBias", energy_bias)
return [energy_bias, forces_bias]
def write_plumed_files(self, images):
""" This function computes what is required in
plumed input for some trajectory.
The outputs are saved in the typical files of
plumed such as COLVAR, HILLS """
for i, image in enumerate(images):
pos = image.get_positions()
self.compute_energy_and_forces(pos, i)
return self.read_plumed_files()
def read_plumed_files(self, file_name=None):
read_files = {}
if file_name is not None:
read_files[file_name] = np.loadtxt(file_name, unpack=True)
else:
for line in self.input:
if line.find('FILE') != -1:
ini = line.find('FILE')
end = line.find(' ', ini)
if end == -1:
file_name = line[ini + 5:]
else:
file_name = line[ini + 5:end]
read_files[file_name] = np.loadtxt(file_name, unpack=True)
if len(read_files) == 0:
if exists('COLVAR'):
read_files['COLVAR'] = np.loadtxt('COLVAR', unpack=True)
if exists('HILLS'):
read_files['HILLS'] = np.loadtxt('HILLS', unpack=True)
assert len(read_files) != 0, "There are not files for reading"
return read_files
def __enter__(self):
return self
def __exit__(self, *args):
self.plumed.finalize()