Coverage for /builds/ase/ase/ase/calculators/plumed.py : 89.83%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from ase.calculators.calculator import Calculator, all_changes
2from ase.io.trajectory import Trajectory
3from ase.parallel import broadcast
4from ase.parallel import world
5import numpy as np
6from os.path import exists
7from ase.units import fs, mol, kJ, nm
10def restart_from_trajectory(prev_traj, *args, prev_steps=None, atoms=None,
11 **kwargs):
12 """This function helps the user to restart a plumed simulation
13 from a trajectory file.
15 Parameters
16 ----------
17 prev_traj : Trajectory object
18 previous simulated trajectory
20 prev_steps : int. Default steps in prev_traj.
21 number of previous steps
23 others :
24 Same parameters of :mod:`~ase.calculators.plumed` calculator
26 Returns
27 -------
28 Plumed calculator
31 .. note:: prev_steps is crucial when trajectory does not contain
32 all the previous steps.
33 """
34 atoms.calc = Plumed(*args, atoms=atoms, restart=True, **kwargs)
36 with Trajectory(prev_traj) as traj:
37 if prev_steps is None:
38 atoms.calc.istep = len(traj) - 1
39 else:
40 atoms.calc.istep = prev_steps
41 atoms.set_positions(traj[-1].get_positions())
42 atoms.set_momenta(traj[-1].get_momenta())
43 return atoms.calc
46class Plumed(Calculator):
47 implemented_properties = ['energy', 'forces']
49 def __init__(self, calc, input, timestep, atoms=None, kT=1., log='',
50 restart=False, use_charge=False, update_charge=False):
51 """
52 Plumed calculator is used for simulations of enhanced sampling methods
53 with the open-source code PLUMED (plumed.org).
55 [1] The PLUMED consortium, Nat. Methods 16, 670 (2019)
56 [2] Tribello, Bonomi, Branduardi, Camilloni, and Bussi,
57 Comput. Phys. Commun. 185, 604 (2014)
59 Parameters
60 ----------
61 calc: Calculator object
62 It computes the unbiased forces
64 input: List of strings
65 It contains the setup of plumed actions
67 timestep: float
68 Time step of the simulated dynamics
70 atoms: Atoms
71 Atoms object to be attached
73 kT: float. Default 1.
74 Value of the thermal energy in eV units. It is important for
75 some methods of plumed like Well-Tempered Metadynamics.
77 log: string
78 Log file of the plumed calculations
80 restart: boolean. Default False
81 True if the simulation is restarted.
83 use_charge: boolean. Default False
84 True if you use some collective variable which needs charges. If
85 use_charges is True and update_charge is False, you have to define
86 initial charges and then this charge will be used during all
87 simulation.
89 update_charge: boolean. Default False
90 True if you want the charges to be updated each time step. This
91 will fail in case that calc does not have 'charges' in its
92 properties.
95 .. note:: For this case, the calculator is defined strictly with the
96 object atoms inside. This is necessary for initializing the
97 Plumed object. For conserving ASE convention, it can be initialized
98 as atoms.calc = (..., atoms=atoms, ...)
101 .. note:: In order to guarantee a proper restart, the user has to fix
102 momenta, positions and Plumed.istep, where the positions and
103 momenta corresponds to the last configuration in the previous
104 simulation, while Plumed.istep is the number of timesteps
105 performed previously. This can be done using
106 ase.calculators.plumed.restart_from_trajectory.
107 """
109 from plumed import Plumed as pl
111 if atoms is None:
112 raise TypeError('plumed calculator has to be defined with the \
113 object atoms inside.')
115 self.istep = 0
116 Calculator.__init__(self, atoms=atoms)
118 self.input = input
119 self.calc = calc
120 self.use_charge = use_charge
121 self.update_charge = update_charge
123 if world.rank == 0:
124 natoms = len(atoms.get_positions())
125 self.plumed = pl()
127 ''' Units setup
128 warning: inputs and outputs of plumed will still be in
129 plumed units.
131 The change of Plumed units to ASE units is:
132 kjoule/mol to eV
133 nm to Angstrom
134 ps to ASE time units
135 ASE and plumed - charge unit is in e units
136 ASE and plumed - mass unit is in a.m.u units '''
138 ps = 1000 * fs
139 self.plumed.cmd("setMDEnergyUnits", mol / kJ)
140 self.plumed.cmd("setMDLengthUnits", 1 / nm)
141 self.plumed.cmd("setMDTimeUnits", 1 / ps)
142 self.plumed.cmd("setMDChargeUnits", 1.)
143 self.plumed.cmd("setMDMassUnits", 1.)
145 self.plumed.cmd("setNatoms", natoms)
146 self.plumed.cmd("setMDEngine", "ASE")
147 self.plumed.cmd("setLogFile", log)
148 self.plumed.cmd("setTimestep", float(timestep))
149 self.plumed.cmd("setRestart", restart)
150 self.plumed.cmd("setKbT", float(kT))
151 self.plumed.cmd("init")
152 for line in input:
153 self.plumed.cmd("readInputLine", line)
154 self.atoms = atoms
156 def _get_name(self):
157 return f'{self.calc.name}+Plumed'
159 def calculate(self, atoms=None, properties=['energy', 'forces'],
160 system_changes=all_changes):
161 Calculator.calculate(self, atoms, properties, system_changes)
163 comp = self.compute_energy_and_forces(self.atoms.get_positions(),
164 self.istep)
165 energy, forces = comp
166 self.istep += 1
167 self.results['energy'], self. results['forces'] = energy, forces
169 def compute_energy_and_forces(self, pos, istep):
170 unbiased_energy = self.calc.get_potential_energy(self.atoms)
171 unbiased_forces = self.calc.get_forces(self.atoms)
173 if world.rank == 0:
174 ener_forc = self.compute_bias(pos, istep, unbiased_energy)
175 else:
176 ener_forc = None
177 energy_bias, forces_bias = broadcast(ener_forc)
178 energy = unbiased_energy + energy_bias
179 forces = unbiased_forces + forces_bias
180 return energy, forces
182 def compute_bias(self, pos, istep, unbiased_energy):
183 self.plumed.cmd("setStep", istep)
185 if self.use_charge:
186 if 'charges' in self.calc.implemented_properties and \
187 self.update_charge:
188 charges = self.calc.get_charges(atoms=self.atoms.copy())
190 elif self.atoms.has('initial_charges') and not self.update_charge:
191 charges = self.atoms.get_initial_charges()
193 else:
194 assert not self.update_charge, "Charges cannot be updated"
195 assert self.update_charge, "Not initial charges in Atoms"
197 self.plumed.cmd("setCharges", charges)
199 # Box for functions with PBC in plumed
200 if self.atoms.cell:
201 cell = np.asarray(self.atoms.get_cell())
202 self.plumed.cmd("setBox", cell)
204 self.plumed.cmd("setPositions", pos)
205 self.plumed.cmd("setEnergy", unbiased_energy)
206 self.plumed.cmd("setMasses", self.atoms.get_masses())
207 forces_bias = np.zeros((self.atoms.get_positions()).shape)
208 self.plumed.cmd("setForces", forces_bias)
209 virial = np.zeros((3, 3))
210 self.plumed.cmd("setVirial", virial)
211 self.plumed.cmd("prepareCalc")
212 self.plumed.cmd("performCalc")
213 energy_bias = np.zeros((1,))
214 self.plumed.cmd("getBias", energy_bias)
215 return [energy_bias, forces_bias]
217 def write_plumed_files(self, images):
218 """ This function computes what is required in
219 plumed input for some trajectory.
221 The outputs are saved in the typical files of
222 plumed such as COLVAR, HILLS """
223 for i, image in enumerate(images):
224 pos = image.get_positions()
225 self.compute_energy_and_forces(pos, i)
226 return self.read_plumed_files()
228 def read_plumed_files(self, file_name=None):
229 read_files = {}
230 if file_name is not None:
231 read_files[file_name] = np.loadtxt(file_name, unpack=True)
232 else:
233 for line in self.input:
234 if line.find('FILE') != -1:
235 ini = line.find('FILE')
236 end = line.find(' ', ini)
237 if end == -1:
238 file_name = line[ini + 5:]
239 else:
240 file_name = line[ini + 5:end]
241 read_files[file_name] = np.loadtxt(file_name, unpack=True)
243 if len(read_files) == 0:
244 if exists('COLVAR'):
245 read_files['COLVAR'] = np.loadtxt('COLVAR', unpack=True)
246 if exists('HILLS'):
247 read_files['HILLS'] = np.loadtxt('HILLS', unpack=True)
248 assert not len(read_files) == 0, "There are not files for reading"
249 return read_files
251 def __enter__(self):
252 return self
254 def __exit__(self, *args):
255 self.plumed.finalize()