Coverage for ase / calculators / plumed.py: 89.39%
132 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1from os.path import exists
3import numpy as np
5from ase.calculators.calculator import Calculator, all_changes
6from ase.io.trajectory import Trajectory
7from ase.parallel import broadcast, world
8from ase.stress import full_3x3_to_voigt_6_stress
9from ase.units import fs, kJ, mol, nm
12def restart_from_trajectory(
13 prev_traj, *args, prev_steps=None, atoms=None, **kwargs
14):
15 """This function helps the user to restart a plumed simulation
16 from a trajectory file.
18 Parameters
19 ----------
20 prev_traj : Trajectory object
21 previous simulated trajectory
23 prev_steps : int. Default steps in prev_traj.
24 number of previous steps
26 others :
27 Same parameters of :mod:`~ase.calculators.plumed` calculator
29 Returns
30 -------
31 Plumed calculator
34 .. note:: prev_steps is crucial when trajectory does not contain
35 all the previous steps.
36 """
37 atoms.calc = Plumed(*args, atoms=atoms, restart=True, **kwargs)
39 with Trajectory(prev_traj) as traj:
40 if prev_steps is None:
41 atoms.calc.istep = len(traj) - 1
42 else:
43 atoms.calc.istep = prev_steps
44 atoms.set_positions(traj[-1].get_positions())
45 atoms.set_momenta(traj[-1].get_momenta())
46 return atoms.calc
49class Plumed(Calculator):
50 implemented_properties = ['energy', 'forces', 'stress']
52 def __init__(
53 self,
54 calc,
55 input,
56 timestep,
57 atoms=None,
58 kT=1.0,
59 log='',
60 restart=False,
61 use_charge=False,
62 update_charge=False,
63 ):
64 """
65 Plumed calculator is used for simulations of enhanced sampling methods
66 with the open-source code PLUMED (plumed.org).
68 [1] The PLUMED consortium, Nat. Methods 16, 670 (2019)
69 [2] Tribello, Bonomi, Branduardi, Camilloni, and Bussi,
70 Comput. Phys. Commun. 185, 604 (2014)
72 Parameters
73 ----------
74 calc: Calculator object
75 It computes the unbiased forces
77 input: List of strings
78 It contains the setup of plumed actions
80 timestep: float
81 Time step of the simulated dynamics
83 atoms: Atoms
84 Atoms object to be attached
86 kT: float. Default 1.
87 Value of the thermal energy in eV units. It is important for
88 some methods of plumed like Well-Tempered Metadynamics.
90 log: string
91 Log file of the plumed calculations
93 restart: boolean. Default False
94 True if the simulation is restarted.
96 use_charge: boolean. Default False
97 True if you use some collective variable which needs charges. If
98 use_charges is True and update_charge is False, you have to define
99 initial charges and then this charge will be used during all
100 simulation.
102 update_charge: boolean. Default False
103 True if you want the charges to be updated each time step. This
104 will fail in case that calc does not have 'charges' in its
105 properties.
108 .. note:: For this case, the calculator is defined strictly with the
109 object atoms inside. This is necessary for initializing the
110 Plumed object. For conserving ASE convention, it can be initialized
111 as atoms.calc = (..., atoms=atoms, ...)
114 .. note:: In order to guarantee a proper restart, the user has to fix
115 momenta, positions and Plumed.istep, where the positions and
116 momenta corresponds to the last configuration in the previous
117 simulation, while Plumed.istep is the number of timesteps
118 performed previously. This can be done using
119 ase.calculators.plumed.restart_from_trajectory.
120 """
122 from plumed import Plumed as pl
124 if atoms is None:
125 raise TypeError(
126 'plumed calculator has to be defined with the \
127 object atoms inside.'
128 )
130 self.istep = 0
131 Calculator.__init__(self, atoms=atoms)
133 self.input = input
134 self.calc = calc
135 self.use_charge = use_charge
136 self.update_charge = update_charge
138 if world.rank == 0:
139 natoms = len(atoms.get_positions())
140 self.plumed = pl()
142 """ Units setup
143 warning: inputs and outputs of plumed will still be in
144 plumed units.
146 The change of Plumed units to ASE units is:
147 kjoule/mol to eV
148 nm to Angstrom
149 ps to ASE time units
150 ASE and plumed - charge unit is in e units
151 ASE and plumed - mass unit is in a.m.u units """
153 ps = 1000 * fs
154 self.plumed.cmd('setMDEnergyUnits', mol / kJ)
155 self.plumed.cmd('setMDLengthUnits', 1 / nm)
156 self.plumed.cmd('setMDTimeUnits', 1 / ps)
157 self.plumed.cmd('setMDChargeUnits', 1.0)
158 self.plumed.cmd('setMDMassUnits', 1.0)
160 self.plumed.cmd('setNatoms', natoms)
161 self.plumed.cmd('setMDEngine', 'ASE')
162 self.plumed.cmd('setLogFile', log)
163 self.plumed.cmd('setTimestep', float(timestep))
164 self.plumed.cmd('setRestart', restart)
165 self.plumed.cmd('setKbT', float(kT))
166 self.plumed.cmd('init')
167 for line in input:
168 self.plumed.cmd('readInputLine', line)
169 self.atoms = atoms
171 def _get_name(self):
172 return f'{self.calc.name}+Plumed'
174 def calculate(
175 self,
176 atoms=None,
177 properties=['energy', 'forces'],
178 system_changes=all_changes,
179 ):
180 Calculator.calculate(self, atoms, properties, system_changes)
182 comp = self._compute_properties(self.atoms.get_positions(), self.istep)
183 energy, forces, stress = comp
184 self.istep += 1
185 self.results['energy'] = float(energy)
186 self.results['forces'] = forces
187 if self.atoms.cell.rank == 3:
188 self.results['stress'] = stress
190 def _compute_properties(self, pos, istep):
191 unbiased_energy = self.calc.get_potential_energy(self.atoms)
192 unbiased_forces = self.calc.get_forces(self.atoms)
193 if world.rank == 0:
194 ener_forc_strs = self.compute_bias(pos, istep, unbiased_energy)
195 else:
196 ener_forc_strs = None
197 energy_bias, forces_bias, virial_bias = broadcast(ener_forc_strs)
198 energy = unbiased_energy + energy_bias
199 forces = unbiased_forces + forces_bias
200 if self.atoms.cell.rank == 3:
201 try:
202 unbiased_stress = self.calc.get_stress(self.atoms)
203 except Exception:
204 unbiased_stress = np.zeros(6)
205 volume = self.atoms.get_volume()
206 biased_stress_voigt = (
207 full_3x3_to_voigt_6_stress(virial_bias) / volume
208 )
209 stress = unbiased_stress + biased_stress_voigt
210 else:
211 stress = None
212 return energy, forces, stress
214 def compute_bias(self, pos, istep, unbiased_energy):
215 self.plumed.cmd('setStep', istep)
217 if self.use_charge:
218 if (
219 'charges' in self.calc.implemented_properties
220 and self.update_charge
221 ):
222 charges = self.calc.get_charges(atoms=self.atoms.copy())
224 elif self.atoms.has('initial_charges') and not self.update_charge:
225 charges = self.atoms.get_initial_charges()
227 else:
228 assert not self.update_charge, 'Charges cannot be updated'
229 assert self.update_charge, 'Not initial charges in Atoms'
231 self.plumed.cmd('setCharges', charges)
233 # Box for functions with PBC in plumed
234 if self.atoms.cell:
235 cell = np.asarray(self.atoms.get_cell())
236 self.plumed.cmd('setBox', cell)
238 self.plumed.cmd('setPositions', pos)
239 self.plumed.cmd('setEnergy', unbiased_energy)
240 self.plumed.cmd('setMasses', self.atoms.get_masses())
241 forces_bias = np.zeros((self.atoms.get_positions()).shape)
242 self.plumed.cmd('setForces', forces_bias)
243 virial_bias = np.zeros((3, 3))
244 self.plumed.cmd('setVirial', virial_bias)
245 self.plumed.cmd('prepareCalc')
246 self.plumed.cmd('performCalc')
247 energy_bias = np.zeros((1,))
248 self.plumed.cmd('getBias', energy_bias)
249 energy_bias = energy_bias[0]
250 return [energy_bias, forces_bias, virial_bias]
252 def write_plumed_files(self, images):
253 """This function computes what is required in
254 plumed input for some trajectory.
256 The outputs are saved in the typical files of
257 plumed such as COLVAR, HILLS"""
258 for i, image in enumerate(images):
259 pos = image.get_positions()
260 self._compute_properties(pos, i)
261 return self.read_plumed_files()
263 def read_plumed_files(self, file_name=None):
264 read_files = {}
265 if file_name is not None:
266 read_files[file_name] = np.loadtxt(file_name, unpack=True)
267 else:
268 for line in self.input:
269 if line.find('FILE') != -1:
270 ini = line.find('FILE')
271 end = line.find(' ', ini)
272 if end == -1:
273 file_name = line[ini + 5 :]
274 else:
275 file_name = line[ini + 5 : end]
276 read_files[file_name] = np.loadtxt(file_name, unpack=True)
278 if len(read_files) == 0:
279 if exists('COLVAR'):
280 read_files['COLVAR'] = np.loadtxt('COLVAR', unpack=True)
281 if exists('HILLS'):
282 read_files['HILLS'] = np.loadtxt('HILLS', unpack=True)
283 assert len(read_files) != 0, 'There are not files for reading'
284 return read_files
286 def __enter__(self):
287 return self
289 def __exit__(self, *args):
290 self.plumed.finalize()