Coverage for /builds/kinetik161/ase/ase/calculators/plumed.py: 89.83%
118 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +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.units import fs, kJ, mol, nm
11def restart_from_trajectory(prev_traj, *args, prev_steps=None, atoms=None,
12 **kwargs):
13 """This function helps the user to restart a plumed simulation
14 from a trajectory file.
16 Parameters
17 ----------
18 prev_traj : Trajectory object
19 previous simulated trajectory
21 prev_steps : int. Default steps in prev_traj.
22 number of previous steps
24 others :
25 Same parameters of :mod:`~ase.calculators.plumed` calculator
27 Returns
28 -------
29 Plumed calculator
32 .. note:: prev_steps is crucial when trajectory does not contain
33 all the previous steps.
34 """
35 atoms.calc = Plumed(*args, atoms=atoms, restart=True, **kwargs)
37 with Trajectory(prev_traj) as traj:
38 if prev_steps is None:
39 atoms.calc.istep = len(traj) - 1
40 else:
41 atoms.calc.istep = prev_steps
42 atoms.set_positions(traj[-1].get_positions())
43 atoms.set_momenta(traj[-1].get_momenta())
44 return atoms.calc
47class Plumed(Calculator):
48 implemented_properties = ['energy', 'forces']
50 def __init__(self, calc, input, timestep, atoms=None, kT=1., log='',
51 restart=False, use_charge=False, update_charge=False):
52 """
53 Plumed calculator is used for simulations of enhanced sampling methods
54 with the open-source code PLUMED (plumed.org).
56 [1] The PLUMED consortium, Nat. Methods 16, 670 (2019)
57 [2] Tribello, Bonomi, Branduardi, Camilloni, and Bussi,
58 Comput. Phys. Commun. 185, 604 (2014)
60 Parameters
61 ----------
62 calc: Calculator object
63 It computes the unbiased forces
65 input: List of strings
66 It contains the setup of plumed actions
68 timestep: float
69 Time step of the simulated dynamics
71 atoms: Atoms
72 Atoms object to be attached
74 kT: float. Default 1.
75 Value of the thermal energy in eV units. It is important for
76 some methods of plumed like Well-Tempered Metadynamics.
78 log: string
79 Log file of the plumed calculations
81 restart: boolean. Default False
82 True if the simulation is restarted.
84 use_charge: boolean. Default False
85 True if you use some collective variable which needs charges. If
86 use_charges is True and update_charge is False, you have to define
87 initial charges and then this charge will be used during all
88 simulation.
90 update_charge: boolean. Default False
91 True if you want the charges to be updated each time step. This
92 will fail in case that calc does not have 'charges' in its
93 properties.
96 .. note:: For this case, the calculator is defined strictly with the
97 object atoms inside. This is necessary for initializing the
98 Plumed object. For conserving ASE convention, it can be initialized
99 as atoms.calc = (..., atoms=atoms, ...)
102 .. note:: In order to guarantee a proper restart, the user has to fix
103 momenta, positions and Plumed.istep, where the positions and
104 momenta corresponds to the last configuration in the previous
105 simulation, while Plumed.istep is the number of timesteps
106 performed previously. This can be done using
107 ase.calculators.plumed.restart_from_trajectory.
108 """
110 from plumed import Plumed as pl
112 if atoms is None:
113 raise TypeError('plumed calculator has to be defined with the \
114 object atoms inside.')
116 self.istep = 0
117 Calculator.__init__(self, atoms=atoms)
119 self.input = input
120 self.calc = calc
121 self.use_charge = use_charge
122 self.update_charge = update_charge
124 if world.rank == 0:
125 natoms = len(atoms.get_positions())
126 self.plumed = pl()
128 ''' Units setup
129 warning: inputs and outputs of plumed will still be in
130 plumed units.
132 The change of Plumed units to ASE units is:
133 kjoule/mol to eV
134 nm to Angstrom
135 ps to ASE time units
136 ASE and plumed - charge unit is in e units
137 ASE and plumed - mass unit is in a.m.u units '''
139 ps = 1000 * fs
140 self.plumed.cmd("setMDEnergyUnits", mol / kJ)
141 self.plumed.cmd("setMDLengthUnits", 1 / nm)
142 self.plumed.cmd("setMDTimeUnits", 1 / ps)
143 self.plumed.cmd("setMDChargeUnits", 1.)
144 self.plumed.cmd("setMDMassUnits", 1.)
146 self.plumed.cmd("setNatoms", natoms)
147 self.plumed.cmd("setMDEngine", "ASE")
148 self.plumed.cmd("setLogFile", log)
149 self.plumed.cmd("setTimestep", float(timestep))
150 self.plumed.cmd("setRestart", restart)
151 self.plumed.cmd("setKbT", float(kT))
152 self.plumed.cmd("init")
153 for line in input:
154 self.plumed.cmd("readInputLine", line)
155 self.atoms = atoms
157 def _get_name(self):
158 return f'{self.calc.name}+Plumed'
160 def calculate(self, atoms=None, properties=['energy', 'forces'],
161 system_changes=all_changes):
162 Calculator.calculate(self, atoms, properties, system_changes)
164 comp = self.compute_energy_and_forces(self.atoms.get_positions(),
165 self.istep)
166 energy, forces = comp
167 self.istep += 1
168 self.results['energy'], self. results['forces'] = energy, forces
170 def compute_energy_and_forces(self, pos, istep):
171 unbiased_energy = self.calc.get_potential_energy(self.atoms)
172 unbiased_forces = self.calc.get_forces(self.atoms)
174 if world.rank == 0:
175 ener_forc = self.compute_bias(pos, istep, unbiased_energy)
176 else:
177 ener_forc = None
178 energy_bias, forces_bias = broadcast(ener_forc)
179 energy = unbiased_energy + energy_bias
180 forces = unbiased_forces + forces_bias
181 return energy, forces
183 def compute_bias(self, pos, istep, unbiased_energy):
184 self.plumed.cmd("setStep", istep)
186 if self.use_charge:
187 if 'charges' in self.calc.implemented_properties and \
188 self.update_charge:
189 charges = self.calc.get_charges(atoms=self.atoms.copy())
191 elif self.atoms.has('initial_charges') and not self.update_charge:
192 charges = self.atoms.get_initial_charges()
194 else:
195 assert not self.update_charge, "Charges cannot be updated"
196 assert self.update_charge, "Not initial charges in Atoms"
198 self.plumed.cmd("setCharges", charges)
200 # Box for functions with PBC in plumed
201 if self.atoms.cell:
202 cell = np.asarray(self.atoms.get_cell())
203 self.plumed.cmd("setBox", cell)
205 self.plumed.cmd("setPositions", pos)
206 self.plumed.cmd("setEnergy", unbiased_energy)
207 self.plumed.cmd("setMasses", self.atoms.get_masses())
208 forces_bias = np.zeros((self.atoms.get_positions()).shape)
209 self.plumed.cmd("setForces", forces_bias)
210 virial = np.zeros((3, 3))
211 self.plumed.cmd("setVirial", virial)
212 self.plumed.cmd("prepareCalc")
213 self.plumed.cmd("performCalc")
214 energy_bias = np.zeros((1,))
215 self.plumed.cmd("getBias", energy_bias)
216 return [energy_bias, forces_bias]
218 def write_plumed_files(self, images):
219 """ This function computes what is required in
220 plumed input for some trajectory.
222 The outputs are saved in the typical files of
223 plumed such as COLVAR, HILLS """
224 for i, image in enumerate(images):
225 pos = image.get_positions()
226 self.compute_energy_and_forces(pos, i)
227 return self.read_plumed_files()
229 def read_plumed_files(self, file_name=None):
230 read_files = {}
231 if file_name is not None:
232 read_files[file_name] = np.loadtxt(file_name, unpack=True)
233 else:
234 for line in self.input:
235 if line.find('FILE') != -1:
236 ini = line.find('FILE')
237 end = line.find(' ', ini)
238 if end == -1:
239 file_name = line[ini + 5:]
240 else:
241 file_name = line[ini + 5:end]
242 read_files[file_name] = np.loadtxt(file_name, unpack=True)
244 if len(read_files) == 0:
245 if exists('COLVAR'):
246 read_files['COLVAR'] = np.loadtxt('COLVAR', unpack=True)
247 if exists('HILLS'):
248 read_files['HILLS'] = np.loadtxt('HILLS', unpack=True)
249 assert not len(read_files) == 0, "There are not files for reading"
250 return read_files
252 def __enter__(self):
253 return self
255 def __exit__(self, *args):
256 self.plumed.finalize()