Coverage for ase / md / langevinbaoab.py: 88.95%
172 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +0000
1import secrets
2import warnings
4import numpy as np
5from scipy.linalg import expm
7from ase import units
8from ase.md.md import MolecularDynamics
9from ase.stress import voigt_6_to_full_3x3_stress
12class LangevinBAOAB(MolecularDynamics):
13 """Time integrator using Langevin for positions and Langevin-Hoover
14 for cell (fixed cell, fixed cell shape + variable volume, or fully
15 variable cell) with BAOAB time propagation
17 BAOAB algorithm from Leimkuhler and Matthews "Robust and efficient
18 configurational molecular sampling via Langevin dynamics",
19 J. Chem. Phys. 138 174102 (2013).
20 https://doi.org/10.1063/1.4802990
22 There is some evidence that other time integration schemes, e.g. BAOA,
23 may be better (https://pubs.acs.org/doi/full/10.1021/acs.jctc.2c00585),
24 and it may be straightforward to add these, but none are not currently
25 supported .
27 Langevin-Hoover from Quigley and Probert "Langevin dynamics in constant
28 pressure extended systems", J. Chem. Phys 120 11432 (2004).
29 https://doi.org/10.1063/1.1755657
31 Parameters
32 ----------
33 atoms: Atoms
34 Atoms object for dynamics.
35 timestep: float
36 Timestep (ASE native units) for time propagation.
37 temperature_K: float, optional
38 Constant temperature to apply, in K. Enables constant temperature
39 dynamics with NVT or NPT, otherwise dynamics are NVE or NPH
40 (depending on ``externalstress``).
41 externalstress: float, ndarray(3), narray(6), narray((3, 3)), optional
42 Constant stress to apply, in ASE native units. Enables variable cell
43 dynamics with constant NPH or NPT (depending on ``temperature_K``).
44 Note that stress is negative of pressure, so *negative* values lead to
45 compression. Note also that barostat will keep mean stress **including
46 kinetic (i.e. ideal gas) contribution** equal to this value. Only
47 scalars are allowed if ``hydrostatic`` is True.
48 hydrostatic: bool, default False
49 Allow only hydrostatic strain (i.e. preserve cell *shape* but allow
50 overall scaling of volume).
51 T_tau: float, optional
52 Time constant for position degree of freedom Langevin. Defaults to 50 *
53 ``timestep`` if not specified.
54 P_tau: float, optional
55 Time constant for variable cell dynamics (cell fluctuation period
56 used to set P_mass heuristic for NPH, and both flucutation period and
57 Langevin timescale for NPT). Defaults to 20 * ``T_tau`` if T_tau is
58 provided, otherwise 1000 * ``timestep``.
59 P_mass: float, optional
60 Mass used for variable cell dynamics. Default is a heuristic value that
61 aims for fluctuation period of ``P_tau / 4``.
62 P_mass_factor: float, default 1.0
63 Factor to multiply heuristic ``P_mass`` if no user ``P_mass`` value is
64 explicitly provided
65 disable_cell_langevin: bool, default False
66 Turn off Langevin thermalization of cell DOF even if ``temperature_K``
67 is not ``None``. Variable cell will still be done if
68 ``externalstress`` is not ``None``, in which case cell equilibration
69 will rely on interaction between cell and position DOFs.
70 rng: np.random.Generator or argument to np.random.default_rng, default None
71 Random number generator for Langevin forces, or integer used as seed
72 for new Generator. If None, a random seed will be generated and
73 reported for future reproducibility of run
74 **kwargs: dict
75 Additional ase.md.md.MolecularDynamics kwargs.
76 """
78 def __init__(
79 self,
80 atoms,
81 timestep,
82 *,
83 temperature_K=None,
84 externalstress=None,
85 hydrostatic=False,
86 T_tau=None,
87 P_tau=None,
88 P_mass=None,
89 P_mass_factor=1.0,
90 disable_cell_langevin=False,
91 rng=None,
92 **kwargs,
93 ):
94 super().__init__(atoms, timestep, **kwargs)
96 self._set_externalstress_hydrostatic(externalstress, hydrostatic)
97 self.disable_cell_langevin = disable_cell_langevin
99 if temperature_K is not None:
100 # run constant T, need rng and T_tau
101 if rng is None:
102 # procedure recommended in
103 # https://blog.scientific-python.org/numpy/numpy-rng/#random-number-generation-with-numpy
104 seed = secrets.randbits(128)
105 self.rng = np.random.default_rng(seed)
106 warnings.warn(
107 f'No rng provided, generated one with seed={seed} from '
108 'secrets.randbits',
109 )
110 elif not isinstance(rng, np.random.Generator):
111 self.rng = np.random.default_rng(rng)
112 else:
113 # already a Generator
114 self.rng = rng
116 if T_tau is None:
117 T_tau = 50.0 * self.dt
118 warnings.warn(
119 'Got `temperature_K` but missing `T_tau`, '
120 f'defaulting to 50 * `timstep` = {T_tau}'
121 )
122 self.T_tau = T_tau
123 if self.T_tau is not None and self.T_tau <= 0:
124 raise ValueError(f'Invalid T_tau {self.T_tau} <= 0')
126 # default contribution to effective gamma used in _BAOAB_OU that comes
127 # from barostat from 2nd term in RHS of Quigley Eq. (5b)
128 self.gamma_mod = 0.0
130 if self.externalstress is not None:
131 # set various quantities required for variable cell dynamics
133 # pressure timescale
134 self._set_P_tau(P_tau)
135 if self.P_tau is not None and self.P_tau <= 0:
136 raise ValueError(f'Invalid P_tau {self.P_tau} <= 0')
138 # initial momentum of cell DOFs
139 self.p_eps = 0.0
141 # Hope that ASE Atoms.get_number_of_degrees_of_freedom() gives
142 # correct value. It's not, for example, completely obvious what
143 # should be # done about the 3 overall translation DOFs, since
144 # conventional # Langevin does not actually preserve those (i.e.
145 # violates conservation of momentum). See, e.g.,
146 # https://doi.org/10.1063/5.0286750
147 # for discussion of variants, e.g. DPD pairwise-force thermostat
148 if len(self.atoms.constraints) != 0:
149 warnings.warn(
150 'WARNING: LangevinBAOAB has not been '
151 'tested with constraints'
152 )
153 self.Ndof = self.atoms.get_number_of_degrees_of_freedom()
155 self._set_barostat_mass(P_mass, P_mass_factor)
157 # must call _after_ var cell quantities are set, because actual
158 # parameters used in dynamics depend on temperature, so this function
159 # needs to know P_tau, etc
160 self.set_temperature(temperature_K, from_init=True)
162 # initialize forces and stresses
163 self._update_accel()
164 if self.externalstress is not None:
165 self._update_force_eps()
167 def _set_externalstress_hydrostatic(self, externalstress, hydrostatic):
168 """Set self.externalstress of correct shape and self.hydrostatic
170 Parameters
171 ----------
172 externalstress: float or None
173 external stress
174 hydrostatic: bool
175 hydrostatic (fixed shape, variable volume) variations
176 """
177 self.hydrostatic = hydrostatic
179 if externalstress is None:
180 self.externalstress = externalstress
181 return None
183 # promote to ndarray to simplify code below
184 externalstress = np.asarray(externalstress)
185 if externalstress.shape == ():
186 externalstress = externalstress.reshape((-1))
188 # reshape to scalar (iff hydrostatic) or 3x3 matrix (general var cell)
189 if self.hydrostatic:
190 # external stress must be scalar
191 if externalstress.shape != (1,):
192 raise ValueError(
193 'externalstress must be scalar when hydrostatic, '
194 f"got '{externalstress}' with shape "
195 f'{externalstress.shape}'
196 )
197 externalstress = externalstress[0]
198 else:
199 # external stress must end up as 3x3 matrix
200 match externalstress.shape:
201 case (1,):
202 externalstress = externalstress * np.identity(3)
203 case (3,):
204 externalstress = np.diag(externalstress)
205 case (6,):
206 externalstress = voigt_6_to_full_3x3_stress(externalstress)
207 case (3, 3):
208 pass
209 case '_':
210 raise ValueError(
211 'externalstress must be scalar, 3-vector (diagonal), '
212 '6-vector (Voigt), or 3x3 matrix, '
213 f'got "{externalstress}" with shape '
214 f'{externalstress.shape}'
215 )
217 self.externalstress = externalstress
219 def _set_P_tau(self, P_tau):
220 """Set self.P_tau from value, or default based on self.T_tau or self.dt
222 Parameters
223 ----------
224 P_tau: float or None
225 time scale for pressure fluctuations
226 """
227 if P_tau is None:
228 if self.T_tau is not None:
229 P_tau = 20.0 * self.T_tau
230 warnings.warn(
231 'Got `externalstress` but missing `P_tau`, got '
232 f'`T_tau`, defaulting to 20 * `T_tau` = {P_tau}'
233 )
234 else:
235 P_tau = 1000.0 * self.dt
236 warnings.warn(
237 'Got `externalstress` but missing `P_tau` and '
238 f'`T_tau`, defaulting to 1000 * `timestep` = {P_tau}'
239 )
240 self.P_tau = P_tau
242 def _set_barostat_mass(self, P_mass, P_mass_factor):
243 """Set self.barostat_mass based on P_mass and P_mass_factor, as well as
244 variable cell timescale self.P_tau
246 Parameters
247 ----------
248 P_mass: float or None
249 barostat mass
250 P_mass_factor: float, default 1.0
251 factor to be applied relative to default value of heuristic
252 """
253 if P_mass is None:
254 # set P_mass with heuristic
255 #
256 # originally tried expression from Quigley Eq. 17
257 # W = 3 N k_B T / (2 pi / tau)^2
258 # empirically didn't work at all
259 #
260 # instead, using empirical value based on tests
261 # of various P_mass, supercell sizes
262 #
263 # supercell of 4 atom Al FCC cell
264 # VARY P_mass: looks like tau \prop sqrt(P_mass)
265 # sc 3 P_mass 1000 T 300 period 34.01360544217687
266 # sc 3 P_mass 10000 T 300 period 91.74311926605505
267 # VARY sc: looks like tau \prop 1/sqrt(N^1/3)
268 # sc 2 P_mass 10000.0 T 400 period 104.16666666666667 (32 atoms)
269 # sc 3 P_mass 10000.0 T 400 period 92.5925925925926 (108 atoms)
270 # sc 4 P_mass 10000.0 T 400 period 66.66666666666667 (256 atoms)
271 # tau = C * sqrt(P_mass) / N**(1/6)
272 # 66 fs for N = 4 * 4^3 = 256, P_mass = 10^4
273 # 66 fs = C * 10000**0.5 / 256.0**(1.0/6.0)
274 # C = 66 fs / (10000**0.5 / 256.0**(1.0/6.0))
275 # C = 1.6630957858612323 fs
276 # P_mass = ((tau / C) * N**(1/6)) ** 2
277 #
278 # note that constant here may be very system (bulk modulus?)
279 # dependent
280 if not self.P_tau > 0:
281 raise ValueError('Heuristic used for P_mass requires P_tau > 0')
282 C = 1.66 * units.fs
283 barostat_mass = (
284 ((self.P_tau / 4.0) / C) * (len(self.atoms) ** (1.0 / 6.0))
285 ) ** 2
286 warnings.warn(
287 f'Using heuristic P_mass {barostat_mass} '
288 f'from P_tau {self.P_tau}'
289 )
290 barostat_mass *= P_mass_factor
291 else:
292 barostat_mass = P_mass
294 self.barostat_mass = barostat_mass
296 def set_temperature(self, temperature_K, from_init=False):
297 """Set the internal parameters that depend on temperature
299 Parameters
300 ----------
301 temperature_K: float
302 temperature in K
303 """
304 self.temperature_K = temperature_K
306 # default to thermostats (for positions and cell DOFs) disabled
307 self.gamma = 0.0
308 self.barostat_gamma = 0.0
310 if self.temperature_K is None:
311 return
313 ############################################################
314 # position related quantities
315 if self.T_tau is None or self.T_tau <= 0:
316 raise RuntimeError(
317 f'Got temperature {self.temperature_K}, but Langevin '
318 f'time-scale T_tau {self.T_tau} is invalud'
319 )
320 self.gamma = 1.0 / self.T_tau
321 # sigma from before Eq. 4 of Leimkuhler
322 sigma = np.sqrt(2.0 * self.gamma * units.kB * self.temperature_K)
323 # prefactor from after Eq. 6
324 self.BAOAB_prefactor = (sigma / np.sqrt(2.0 * self.gamma)) * np.sqrt(
325 1.0 - np.exp(-2.0 * self.gamma * self.dt)
326 )
327 # does not include sqrt(mass), since that is different for
328 # each atom type
330 ############################################################
331 # cell related quantities
332 if self.externalstress is not None and not self.disable_cell_langevin:
333 self.barostat_gamma = 1.0 / self.P_tau
334 sigma = np.sqrt(
335 2.0 * self.barostat_gamma * units.kB * self.temperature_K
336 )
337 self._barostat_BAOAB_prefactor = (
338 (sigma / np.sqrt(2.0 * self.barostat_gamma))
339 * np.sqrt(1.0 - np.exp(-2.0 * self.barostat_gamma * self.dt))
340 * np.sqrt(self.barostat_mass)
341 )
342 # _does_ include sqrt(mass) factor
344 def _update_accel(self):
345 """Update position-acceleration from current positions via forces"""
346 self.accel = (self.atoms.get_forces().T / self.atoms.get_masses()).T
348 def _update_force_eps(self):
349 """Update cell force from current positions via stress"""
350 volume = self.atoms.get_volume()
351 if self.hydrostatic:
352 KE = self.atoms.get_kinetic_energy()
353 Tr_virial = -volume * np.trace(self.atoms.get_stress(voigt=False))
354 X = 1.0 / (3.0 * volume) * (2.0 * KE + Tr_virial)
356 # NB explicit dphi/dV term in Quigley Eq. 6 is old fashioned
357 # explicit volume dependence for long range tails. Stress/pressure
358 # comes in via sum r . f, which is
359 # Tr[virial] = - volume * Tr[stress]
360 self.force_eps = (
361 3.0 * volume * (X + self.externalstress)
362 + (3.0 / self.Ndof) * 2.0 * KE
363 )
364 else:
365 mom = self.atoms.get_momenta()
366 kinetic_stress_contrib = (mom.T / self.atoms.get_masses()) @ mom
367 virial = -volume * self.atoms.get_stress(voigt=False)
368 X = (1.0 / volume) * (kinetic_stress_contrib + virial)
370 self.force_eps = volume * (X + self.externalstress) + (
371 1.0 / self.Ndof
372 ) * np.trace(kinetic_stress_contrib) * np.identity(3)
374 def _BAOAB_B(self):
375 """Do a BAOAB B (velocity) half step"""
376 dvel = 0.5 * self.dt * self.accel
377 self.atoms.set_velocities(self.atoms.get_velocities() + dvel)
379 def _barostat_BAOAB_B(self):
380 """Do a barostat BAOAB B (cell momentum) half step"""
381 self.p_eps += 0.5 * self.dt * self.force_eps
383 def _BAOAB_A(self):
384 """Do a BAOAB A (position) half step"""
385 self.atoms.positions += 0.5 * self.dt * self.atoms.get_velocities()
387 def _barostat_BAOAB_A(self):
388 """Do a barostat BAOAB A (cell volume) half step"""
389 if self.hydrostatic:
390 volume = self.atoms.get_volume()
391 new_volume = volume * np.exp(
392 0.5 * self.dt * 3.0 * self.p_eps / self.barostat_mass
393 )
394 new_cell = self.atoms.cell * (new_volume / volume) ** (1.0 / 3.0)
395 else:
396 new_cell = (
397 self.atoms.cell
398 @ expm(0.5 * self.dt * self.p_eps / self.barostat_mass).T
399 )
401 self.atoms.set_cell(new_cell, True)
403 def _BAOAB_OU(self, drag_gamma_mod=0.0):
404 """Do a BAOAB Ornstein-Uhlenbeck position Langevin full step
406 Parameters
407 ----------
408 drag_gamma_mod: float, default 0
409 additional contribution to effective gamma used for drag on
410 velocities, e.g. from Langevin-Hoover
411 """
412 if self.gamma == 0 and np.all(drag_gamma_mod == 0):
413 return
415 vel = self.atoms.get_velocities()
417 if self.hydrostatic:
418 vel *= np.exp(-(self.gamma + drag_gamma_mod) * self.dt)
419 else:
420 vel = (
421 vel
422 @ expm(
423 -(self.gamma * np.identity(3) + drag_gamma_mod) * self.dt
424 ).T
425 )
427 if self.gamma != 0:
428 # here we divide by sqrt(m), since Leimkuhler definition includes
429 # sqrt(m) in numerator but that's for momentum, so for velocity we
430 # divide by m, i.e. net 1/sqrt(m)
431 vel_shape = self.atoms.positions.shape
432 masses = self.atoms.get_masses()
433 vel += (
434 self.BAOAB_prefactor
435 * (self.rng.normal(size=vel_shape).T / np.sqrt(masses)).T
436 )
438 self.atoms.set_velocities(vel)
440 def _barostat_BAOAB_OU(self):
441 """Do a barostat BAOAB Ornstein-Uhlenbeck cell volume Langevin full
442 step"""
443 if self.barostat_gamma == 0:
444 # i.e. temperature_K is None or disable_cell_langevin
445 return
447 self.p_eps *= np.exp(-self.barostat_gamma * self.dt)
449 if self.hydrostatic:
450 self.p_eps += self._barostat_BAOAB_prefactor * self.rng.normal()
451 else:
452 random_force = self.rng.normal(size=(3, 3))
453 # symmetrize to avoid rotation
454 random_force += random_force.T
455 random_force /= 2
456 self.p_eps += self._barostat_BAOAB_prefactor * random_force
458 def _barostat_BAOAB_OU_gamma_mod(self):
459 """Compute contribution to drag effective gamma applied to atom
460 velocities due to barostat velocity
461 """
462 if self.hydrostatic:
463 return (1.0 + 3.0 / self.Ndof) * self.p_eps / self.barostat_mass
464 else:
465 return (
466 self.p_eps + (1.0 / self.Ndof) * np.trace(self.p_eps)
467 ) / self.barostat_mass
469 def step(self):
470 """Do a time step"""
471 if self.externalstress is not None:
472 self._barostat_BAOAB_B()
473 self._BAOAB_B() # half step vel
475 if self.externalstress is not None:
476 self._barostat_BAOAB_A()
477 self._BAOAB_A() # half step pos
479 if self.externalstress is not None:
480 self.gamma_mod = self._barostat_BAOAB_OU_gamma_mod()
481 self._BAOAB_OU(self.gamma_mod) # OU vel
482 if self.externalstress is not None:
483 self._barostat_BAOAB_OU()
485 self._BAOAB_A() # half step pos
486 if self.externalstress is not None:
487 self._barostat_BAOAB_A()
489 self._update_accel() # update accel from final pos
490 if self.externalstress is not None:
491 self._update_force_eps() # update cell force
493 self._BAOAB_B() # half step vel
494 if self.externalstress is not None:
495 self._barostat_BAOAB_B()