Coverage for ase / md / velocitydistribution.py: 82.61%
115 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-21 15:52 +0000
1"""Module for setting up velocity distributions such as Maxwell–Boltzmann.
3Currently, only a few functions are defined, such as
4MaxwellBoltzmannDistribution, which sets the momenta of a list of
5atoms according to a Maxwell-Boltzmann distribution at a given
6temperature.
7"""
9from typing import Literal
11import numpy as np
13from ase import Atoms, units
14from ase.md.md import process_temperature
15from ase.parallel import world
16from ase.utils import deprecated
18# define a ``zero'' temperature to avoid divisions by zero
19eps_temp = 1e-12
22class UnitError(Exception):
23 """Exception raised when wrong units are specified"""
26def force_temperature(
27 atoms: Atoms,
28 temperature: float,
29 unit: Literal['K', 'eV'] = 'K',
30):
31 """
32 Force the temperature of the atomic system to a precise value.
34 This function adjusts the momenta of the atoms to achieve the
35 exact specified temperature, overriding the Maxwell-Boltzmann
36 distribution. The temperature can be specified in either Kelvin
37 (K) or electron volts (eV).
39 Parameters
40 ----------
41 atoms
42 The atomic system represented as an ASE Atoms object.
43 temperature
44 The exact temperature to force for the atomic system.
45 unit
46 The unit of the specified temperature.
47 Can be either 'K' (Kelvin) or 'eV' (electron volts). Default is 'K'.
48 """
50 if unit == 'K':
51 target_temp = temperature * units.kB
52 elif unit == 'eV':
53 target_temp = temperature
54 else:
55 raise UnitError(f"'{unit}' is not supported, use 'K' or 'eV'.")
57 if temperature > eps_temp:
58 ndof = atoms.get_number_of_degrees_of_freedom()
59 current_temp = 2 * atoms.get_kinetic_energy() / ndof
60 scale = target_temp / current_temp
61 else:
62 scale = 0.0
64 atoms.set_momenta(atoms.get_momenta() * np.sqrt(scale))
67def _maxwellboltzmanndistribution(masses, temp, comm=world, rng=None):
68 """Return a Maxwell-Boltzmann distribution with a given temperature.
70 Parameters
71 ----------
72 masses: float
73 The atomic masses.
75 temp: float
76 The temperature in electron volt.
78 comm: MPI communicator (optional, default: ase.parallel.world)
79 Communicator used to distribute an identical distribution to all tasks.
81 rng: numpy RNG (optional)
82 The random number generator. Default: np.random
84 Returns
85 -------
86 np.ndarray
87 Maxwell-Boltzmann distributed momenta.
88 """
89 if rng is None:
90 rng = np.random
91 xi = rng.standard_normal((len(masses), 3))
92 comm.broadcast(xi, 0)
93 momenta = xi * np.sqrt(masses * temp)[:, np.newaxis]
94 return momenta
97@deprecated('Use thermalize_momenta', DeprecationWarning)
98def MaxwellBoltzmannDistribution(
99 atoms: Atoms,
100 *,
101 temperature_K: float,
102 comm=world,
103 force_temp: bool = False,
104 rng=None,
105):
106 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
108 .. versionremoved:: 3.29.0
109 The ``temp`` argument is removed. Use ``temperature_K`` instead.
111 .. deprecated:: 3.29.0
112 Use :func:`thermalize_momenta` instead.
114 Parameters
115 ----------
116 atoms: Atoms object
117 The atoms. Their momenta will be modified.
119 temperature_K: float
120 The temperature in Kelvin.
122 comm: MPI communicator, default: :data:`ase.parallel.world`
123 Communicator used to distribute an identical distribution to all tasks.
125 .. versionadded:: 3.26.0
127 .. versionremoved:: 3.29.0
129 ``communicator`` has been removed in favor of ``comm``.
131 force_temp: bool (optional, default: False)
132 If True, the random momenta are rescaled so the kinetic energy is
133 exactly 3/2 N k T. This is a slight deviation from the correct
134 Maxwell-Boltzmann distribution.
136 rng: Numpy RNG (optional)
137 Random number generator. Default: numpy.random
138 If you would like to always get the identical distribution, you can
139 supply a random seed like ``rng=numpy.random.RandomState(seed)``, where
140 seed is an integer.
141 """
142 thermalize_momenta(
143 atoms,
144 temperature_K,
145 comm=comm,
146 exact_temperature=force_temp,
147 rng=rng,
148 )
151def thermalize_momenta(
152 atoms: Atoms,
153 temperature_K: float,
154 *,
155 comm=world,
156 exact_temperature: bool = False,
157 rng=None,
158) -> None:
159 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
161 See :func:`MaxwellBoltzmannDistribution` for details of parameters.
163 .. versionadded:: 3.29.0
165 Parameters
166 ----------
167 atoms: Atoms object
168 The atoms. Their momenta will be modified.
170 temperature_K: float
171 The temperature in Kelvin.
173 comm: MPI communicator, default: :data:`ase.parallel.world`
174 Communicator used to distribute an identical distribution to all tasks.
176 exact_temperature: bool, default: False
177 If True, the random momenta are rescaled so the kinetic energy is
178 exactly 3/2 N k T. This is a slight deviation from the correct
179 Maxwell-Boltzmann distribution.
181 rng: Numpy RNG (optional)
182 Random number generator. Default: numpy.random
183 If you would like to always get the identical distribution, you can
184 supply a random seed like ``rng=numpy.random.RandomState(seed)``, where
185 seed is an integer.
186 """
187 temp = units.kB * temperature_K # K -> eV
188 masses = atoms.get_masses()
189 momenta = _maxwellboltzmanndistribution(masses, temp, comm=comm, rng=rng)
190 atoms.set_momenta(momenta, apply_constraint=True)
191 if exact_temperature:
192 force_temperature(atoms, temperature=temp, unit='eV')
195def Stationary(atoms: Atoms, preserve_temperature: bool = True):
196 "Sets the center-of-mass momentum to zero."
198 # Save initial temperature
199 temp0 = atoms.get_temperature()
201 p = atoms.get_momenta()
202 p0 = np.sum(p, 0)
203 # We should add a constant velocity, not momentum, to the atoms
204 m = atoms.get_masses()
205 mtot = np.sum(m)
206 v0 = p0 / mtot
207 p -= v0 * m[:, np.newaxis]
208 atoms.set_momenta(p)
210 if preserve_temperature:
211 force_temperature(atoms, temp0)
214def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
215 "Sets the total angular momentum to zero by counteracting rigid rotations."
217 # Save initial temperature
218 temp0 = atoms.get_temperature()
220 # Find the principal moments of inertia and principal axes basis vectors
221 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
222 # Calculate the total angular momentum and transform to principal basis
223 Lp = np.dot(basis, atoms.get_angular_momentum())
224 # Calculate the rotation velocity vector in the principal basis, avoiding
225 # zero division, and transform it back to the cartesian coordinate system
226 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
227 # We subtract a rigid rotation corresponding to this rotation vector
228 com = atoms.get_center_of_mass()
229 positions = atoms.get_positions()
230 positions -= com # translate center of mass to origin
231 velocities = atoms.get_velocities()
232 atoms.set_velocities(velocities - np.cross(omega, positions))
234 if preserve_temperature:
235 force_temperature(atoms, temp0)
238def n_BE(temp, omega):
239 """Bose-Einstein distribution function.
241 Args:
242 temp: temperature converted to eV (*units.kB)
243 omega: sequence of frequencies converted to eV
245 Returns
246 -------
247 Value of Bose-Einstein distribution function for each energy
249 """
251 omega = np.asarray(omega)
253 # 0K limit
254 if temp < eps_temp:
255 n = np.zeros_like(omega)
256 else:
257 n = 1 / (np.exp(omega / (temp)) - 1)
258 return n
261def phonon_harmonics(
262 force_constants,
263 masses,
264 temp=None,
265 *,
266 temperature_K=None,
267 rng=np.random.rand,
268 quantum=False,
269 plus_minus=False,
270 return_eigensolution=False,
271 failfast=True,
272):
273 r"""Return displacements and velocities that produce a given temperature.
275 Parameters
276 ----------
277 force_constants: array of size 3N x 3N
278 force constants (Hessian) of the system in eV/Ų
280 masses: array of length N
281 masses of the structure in amu
283 temp: float (deprecated)
284 Temperature converted to eV (T * units.kB). Deprecated, use
285 ``temperature_K``.
287 temperature_K: float
288 Temperature in Kelvin.
290 rng: function
291 Random number generator function, e.g., np.random.rand
293 quantum: bool
294 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
295 (classical limit)
297 plus_minus: bool
298 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
300 return_eigensolution: bool
301 return eigenvalues and eigenvectors of the dynamical matrix
303 failfast: bool
304 True for sanity checking the phonon spectrum for negative
305 frequencies at Gamma
307 Returns
308 -------
309 Displacements, velocities generated from the eigenmodes,
310 (optional: eigenvalues, eigenvectors of dynamical matrix)
312 Notes
313 -----
314 Excite phonon modes to specified temperature.
316 This excites all phonon modes randomly so that each contributes,
317 on average, equally to the given temperature. Both potential
318 energy and kinetic energy will be consistent with the phononic
319 vibrations characteristic of the specified temperature.
321 In other words the system will be equilibrated for an MD run at
322 that temperature.
324 force_constants should be the matrix as force constants, e.g.,
325 as computed by the ase.phonons module.
327 Let X_ai be the phonon modes indexed by atom and mode, w_i the
328 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
329 uniformly random numbers. Then
331 .. code-block:: none
334 1/2
335 _ / k T \ --- 1 _ 1/2
336 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
337 a \ m / --- w ai i i
338 a i i
341 1/2
342 _ / k T \ --- _ 1/2
343 v = | --- | > X (-2 ln Q ) sin (2 pi R )
344 a \ m / --- ai i i
345 a i
347 References
348 ----------
349 D. West and S. K. Estreicher, Phys. Rev. Lett. 96, 115504 (2006).
350 https://doi.org/10.1103/PhysRevLett.96.115504
352 """
354 # Handle the temperature units
355 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
357 # Build dynamical matrix
358 rminv = (masses**-0.5).repeat(3)
359 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
361 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
362 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
364 # Check for soft modes
365 if failfast:
366 zeros = w2_s[:3]
367 worst_zero = np.abs(zeros).max()
368 if worst_zero > 1e-3:
369 msg = 'Translational deviate from 0 significantly: '
370 raise ValueError(msg + f'{w2_s[:3]}')
372 w2min = w2_s[3:].min()
373 if w2min < 0:
374 msg = 'Dynamical matrix has negative eigenvalues such as '
375 raise ValueError(msg + f'{w2min}')
377 # First three modes are translational so ignore:
378 nw = len(w2_s) - 3
379 n_atoms = len(masses)
380 w_s = np.sqrt(w2_s[3:])
381 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
383 # Assign the amplitudes according to Bose-Einstein distribution
384 # or high temperature (== classical) limit
385 if quantum:
386 hbar = units._hbar * units.J * units.s
387 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
388 else:
389 A_s = np.sqrt(temp) / w_s
391 if plus_minus:
392 # create samples by multiplying the amplitude with +/-
393 # according to Eq. 5 in PRB 94, 075125
395 spread = (-1) ** np.arange(nw)
397 # gauge eigenvectors: largest value always positive
398 for ii in range(X_acs.shape[-1]):
399 vec = X_acs[:, :, ii]
400 max_arg = np.argmax(abs(vec))
401 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
403 # Create velocities und displacements from the amplitudes and
404 # eigenvectors
405 A_s *= spread
406 phi_s = 2.0 * np.pi * rng(nw)
408 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
409 # amplitude for displacement
410 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
411 v_ac /= np.sqrt(masses)[:, None]
413 # Assign displacements
414 d_ac = (A_s * X_acs).sum(axis=2)
415 d_ac /= np.sqrt(masses)[:, None]
417 else:
418 # compute the gaussian distribution for the amplitudes
419 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
420 # to avoid (highly improbable) NaN.
422 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
423 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
425 # assign amplitudes and phases
426 A_s *= spread
427 phi_s = 2.0 * np.pi * rng(nw)
429 # Assign velocities and displacements
430 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
431 v_ac /= np.sqrt(masses)[:, None]
433 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
434 d_ac /= np.sqrt(masses)[:, None]
436 if return_eigensolution:
437 return d_ac, v_ac, w2_s, X_is
438 # else
439 return d_ac, v_ac
442def PhononHarmonics(
443 atoms,
444 force_constants,
445 temp=None,
446 *,
447 temperature_K=None,
448 rng=np.random,
449 quantum=False,
450 plus_minus=False,
451 return_eigensolution=False,
452 failfast=True,
453):
454 r"""Excite phonon modes to specified temperature.
456 This will displace atomic positions and set the velocities so as
457 to produce a random, phononically correct state with the requested
458 temperature.
460 Parameters
461 ----------
462 atoms: :class:`~ase.Atoms`
463 Positions and momenta of this object are perturbed.
465 force_constants: ndarray of size 3N x 3N
466 Force constants for the the structure represented by atoms in eV/Ų
468 temp: float (deprecated).
469 Temperature in eV. Deprecated, use ``temperature_K`` instead.
471 temperature_K: float
472 Temperature in Kelvin.
474 rng: Random number generator
475 RandomState or other random number generator, e.g., np.random.rand
477 quantum: bool
478 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
479 (classical limit)
481 failfast: bool
482 True for sanity checking the phonon spectrum for negative frequencies
483 at Gamma.
485 """
487 # Receive displacements and velocities from phonon_harmonics()
488 d_ac, v_ac = phonon_harmonics(
489 force_constants=force_constants,
490 masses=atoms.get_masses(),
491 temp=temp,
492 temperature_K=temperature_K,
493 rng=rng.rand,
494 plus_minus=plus_minus,
495 quantum=quantum,
496 failfast=failfast,
497 return_eigensolution=False,
498 )
500 # Assign new positions (with displacements) and velocities
501 atoms.positions += d_ac
502 atoms.set_velocities(v_ac)