Coverage for /builds/ase/ase/ase/calculators/plumed.py: 89.83%
118 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3from os.path import exists
5import numpy as np
7from ase.calculators.calculator import Calculator, all_changes
8from ase.io.trajectory import Trajectory
9from ase.parallel import broadcast, world
10from ase.units import fs, kJ, mol, nm
13def restart_from_trajectory(prev_traj, *args, prev_steps=None, atoms=None,
14 **kwargs):
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']
52 def __init__(self, calc, input, timestep, atoms=None, kT=1., log='',
53 restart=False, use_charge=False, update_charge=False):
54 """
55 Plumed calculator is used for simulations of enhanced sampling methods
56 with the open-source code PLUMED (plumed.org).
58 [1] The PLUMED consortium, Nat. Methods 16, 670 (2019)
59 [2] Tribello, Bonomi, Branduardi, Camilloni, and Bussi,
60 Comput. Phys. Commun. 185, 604 (2014)
62 Parameters
63 ----------
64 calc: Calculator object
65 It computes the unbiased forces
67 input: List of strings
68 It contains the setup of plumed actions
70 timestep: float
71 Time step of the simulated dynamics
73 atoms: Atoms
74 Atoms object to be attached
76 kT: float. Default 1.
77 Value of the thermal energy in eV units. It is important for
78 some methods of plumed like Well-Tempered Metadynamics.
80 log: string
81 Log file of the plumed calculations
83 restart: boolean. Default False
84 True if the simulation is restarted.
86 use_charge: boolean. Default False
87 True if you use some collective variable which needs charges. If
88 use_charges is True and update_charge is False, you have to define
89 initial charges and then this charge will be used during all
90 simulation.
92 update_charge: boolean. Default False
93 True if you want the charges to be updated each time step. This
94 will fail in case that calc does not have 'charges' in its
95 properties.
98 .. note:: For this case, the calculator is defined strictly with the
99 object atoms inside. This is necessary for initializing the
100 Plumed object. For conserving ASE convention, it can be initialized
101 as atoms.calc = (..., atoms=atoms, ...)
104 .. note:: In order to guarantee a proper restart, the user has to fix
105 momenta, positions and Plumed.istep, where the positions and
106 momenta corresponds to the last configuration in the previous
107 simulation, while Plumed.istep is the number of timesteps
108 performed previously. This can be done using
109 ase.calculators.plumed.restart_from_trajectory.
110 """
112 from plumed import Plumed as pl
114 if atoms is None:
115 raise TypeError('plumed calculator has to be defined with the \
116 object atoms inside.')
118 self.istep = 0
119 Calculator.__init__(self, atoms=atoms)
121 self.input = input
122 self.calc = calc
123 self.use_charge = use_charge
124 self.update_charge = update_charge
126 if world.rank == 0:
127 natoms = len(atoms.get_positions())
128 self.plumed = pl()
130 ''' Units setup
131 warning: inputs and outputs of plumed will still be in
132 plumed units.
134 The change of Plumed units to ASE units is:
135 kjoule/mol to eV
136 nm to Angstrom
137 ps to ASE time units
138 ASE and plumed - charge unit is in e units
139 ASE and plumed - mass unit is in a.m.u units '''
141 ps = 1000 * fs
142 self.plumed.cmd("setMDEnergyUnits", mol / kJ)
143 self.plumed.cmd("setMDLengthUnits", 1 / nm)
144 self.plumed.cmd("setMDTimeUnits", 1 / ps)
145 self.plumed.cmd("setMDChargeUnits", 1.)
146 self.plumed.cmd("setMDMassUnits", 1.)
148 self.plumed.cmd("setNatoms", natoms)
149 self.plumed.cmd("setMDEngine", "ASE")
150 self.plumed.cmd("setLogFile", log)
151 self.plumed.cmd("setTimestep", float(timestep))
152 self.plumed.cmd("setRestart", restart)
153 self.plumed.cmd("setKbT", float(kT))
154 self.plumed.cmd("init")
155 for line in input:
156 self.plumed.cmd("readInputLine", line)
157 self.atoms = atoms
159 def _get_name(self):
160 return f'{self.calc.name}+Plumed'
162 def calculate(self, atoms=None, properties=['energy', 'forces'],
163 system_changes=all_changes):
164 Calculator.calculate(self, atoms, properties, system_changes)
166 comp = self.compute_energy_and_forces(self.atoms.get_positions(),
167 self.istep)
168 energy, forces = comp
169 self.istep += 1
170 self.results['energy'], self. results['forces'] = energy, forces
172 def compute_energy_and_forces(self, pos, istep):
173 unbiased_energy = self.calc.get_potential_energy(self.atoms)
174 unbiased_forces = self.calc.get_forces(self.atoms)
176 if world.rank == 0:
177 ener_forc = self.compute_bias(pos, istep, unbiased_energy)
178 else:
179 ener_forc = None
180 energy_bias, forces_bias = broadcast(ener_forc)
181 energy = unbiased_energy + energy_bias
182 forces = unbiased_forces + forces_bias
183 return energy, forces
185 def compute_bias(self, pos, istep, unbiased_energy):
186 self.plumed.cmd("setStep", istep)
188 if self.use_charge:
189 if 'charges' in self.calc.implemented_properties and \
190 self.update_charge:
191 charges = self.calc.get_charges(atoms=self.atoms.copy())
193 elif self.atoms.has('initial_charges') and not self.update_charge:
194 charges = self.atoms.get_initial_charges()
196 else:
197 assert not self.update_charge, "Charges cannot be updated"
198 assert self.update_charge, "Not initial charges in Atoms"
200 self.plumed.cmd("setCharges", charges)
202 # Box for functions with PBC in plumed
203 if self.atoms.cell:
204 cell = np.asarray(self.atoms.get_cell())
205 self.plumed.cmd("setBox", cell)
207 self.plumed.cmd("setPositions", pos)
208 self.plumed.cmd("setEnergy", unbiased_energy)
209 self.plumed.cmd("setMasses", self.atoms.get_masses())
210 forces_bias = np.zeros((self.atoms.get_positions()).shape)
211 self.plumed.cmd("setForces", forces_bias)
212 virial = np.zeros((3, 3))
213 self.plumed.cmd("setVirial", virial)
214 self.plumed.cmd("prepareCalc")
215 self.plumed.cmd("performCalc")
216 energy_bias = np.zeros((1,))
217 self.plumed.cmd("getBias", energy_bias)
218 return [energy_bias, forces_bias]
220 def write_plumed_files(self, images):
221 """ This function computes what is required in
222 plumed input for some trajectory.
224 The outputs are saved in the typical files of
225 plumed such as COLVAR, HILLS """
226 for i, image in enumerate(images):
227 pos = image.get_positions()
228 self.compute_energy_and_forces(pos, i)
229 return self.read_plumed_files()
231 def read_plumed_files(self, file_name=None):
232 read_files = {}
233 if file_name is not None:
234 read_files[file_name] = np.loadtxt(file_name, unpack=True)
235 else:
236 for line in self.input:
237 if line.find('FILE') != -1:
238 ini = line.find('FILE')
239 end = line.find(' ', ini)
240 if end == -1:
241 file_name = line[ini + 5:]
242 else:
243 file_name = line[ini + 5:end]
244 read_files[file_name] = np.loadtxt(file_name, unpack=True)
246 if len(read_files) == 0:
247 if exists('COLVAR'):
248 read_files['COLVAR'] = np.loadtxt('COLVAR', unpack=True)
249 if exists('HILLS'):
250 read_files['HILLS'] = np.loadtxt('HILLS', unpack=True)
251 assert len(read_files) != 0, "There are not files for reading"
252 return read_files
254 def __enter__(self):
255 return self
257 def __exit__(self, *args):
258 self.plumed.finalize()