Coverage for ase / md / velocitydistribution.py: 78.33%
120 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +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"""
9import warnings
10from typing import Literal
12import numpy as np
14from ase import Atoms, units
15from ase.md.md import process_temperature
16from ase.parallel import DummyMPI, world
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
97def MaxwellBoltzmannDistribution(
98 atoms: Atoms,
99 temp: float | None = None,
100 *,
101 temperature_K: float | None = None,
102 comm=world,
103 communicator=None,
104 force_temp: bool = False,
105 rng=None,
106):
107 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
109 Parameters
110 ----------
111 atoms: Atoms object
112 The atoms. Their momenta will be modified.
114 temp: float (deprecated)
115 The temperature in eV. Deprecated, use ``temperature_K`` instead.
117 temperature_K: float
118 The temperature in Kelvin.
120 comm: MPI communicator, default: :data:`ase.parallel.world`
121 Communicator used to distribute an identical distribution to all tasks.
123 .. versionadded:: 3.26.0
125 communicator
127 .. deprecated:: 3.26.0
129 To be removed in ASE 3.27.0 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 if communicator is not None:
143 msg = (
144 '`communicator` has been deprecated since ASE 3.26.0 '
145 'and will be removed in ASE 3.27.0. Use `comm` instead.'
146 )
147 warnings.warn(msg, FutureWarning)
148 comm = communicator
149 if comm == 'serial':
150 msg = (
151 '`comm=="serial"` has been deprecated since ASE 3.26.0 '
152 'and will be removed in ASE 3.27.0. Use `comm=DummyMPI()` instead.'
153 )
154 warnings.warn(msg, FutureWarning)
155 comm = DummyMPI()
156 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
157 masses = atoms.get_masses()
158 momenta = _maxwellboltzmanndistribution(masses, temp, comm=comm, rng=rng)
159 atoms.set_momenta(momenta)
160 if force_temp:
161 force_temperature(atoms, temperature=temp, unit='eV')
164def Stationary(atoms: Atoms, preserve_temperature: bool = True):
165 "Sets the center-of-mass momentum to zero."
167 # Save initial temperature
168 temp0 = atoms.get_temperature()
170 p = atoms.get_momenta()
171 p0 = np.sum(p, 0)
172 # We should add a constant velocity, not momentum, to the atoms
173 m = atoms.get_masses()
174 mtot = np.sum(m)
175 v0 = p0 / mtot
176 p -= v0 * m[:, np.newaxis]
177 atoms.set_momenta(p)
179 if preserve_temperature:
180 force_temperature(atoms, temp0)
183def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
184 "Sets the total angular momentum to zero by counteracting rigid rotations."
186 # Save initial temperature
187 temp0 = atoms.get_temperature()
189 # Find the principal moments of inertia and principal axes basis vectors
190 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
191 # Calculate the total angular momentum and transform to principal basis
192 Lp = np.dot(basis, atoms.get_angular_momentum())
193 # Calculate the rotation velocity vector in the principal basis, avoiding
194 # zero division, and transform it back to the cartesian coordinate system
195 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
196 # We subtract a rigid rotation corresponding to this rotation vector
197 com = atoms.get_center_of_mass()
198 positions = atoms.get_positions()
199 positions -= com # translate center of mass to origin
200 velocities = atoms.get_velocities()
201 atoms.set_velocities(velocities - np.cross(omega, positions))
203 if preserve_temperature:
204 force_temperature(atoms, temp0)
207def n_BE(temp, omega):
208 """Bose-Einstein distribution function.
210 Args:
211 temp: temperature converted to eV (*units.kB)
212 omega: sequence of frequencies converted to eV
214 Returns
215 -------
216 Value of Bose-Einstein distribution function for each energy
218 """
220 omega = np.asarray(omega)
222 # 0K limit
223 if temp < eps_temp:
224 n = np.zeros_like(omega)
225 else:
226 n = 1 / (np.exp(omega / (temp)) - 1)
227 return n
230def phonon_harmonics(
231 force_constants,
232 masses,
233 temp=None,
234 *,
235 temperature_K=None,
236 rng=np.random.rand,
237 quantum=False,
238 plus_minus=False,
239 return_eigensolution=False,
240 failfast=True,
241):
242 r"""Return displacements and velocities that produce a given temperature.
244 Parameters
245 ----------
246 force_constants: array of size 3N x 3N
247 force constants (Hessian) of the system in eV/Ų
249 masses: array of length N
250 masses of the structure in amu
252 temp: float (deprecated)
253 Temperature converted to eV (T * units.kB). Deprecated, use
254 ``temperature_K``.
256 temperature_K: float
257 Temperature in Kelvin.
259 rng: function
260 Random number generator function, e.g., np.random.rand
262 quantum: bool
263 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
264 (classical limit)
266 plus_minus: bool
267 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
269 return_eigensolution: bool
270 return eigenvalues and eigenvectors of the dynamical matrix
272 failfast: bool
273 True for sanity checking the phonon spectrum for negative
274 frequencies at Gamma
276 Returns
277 -------
278 Displacements, velocities generated from the eigenmodes,
279 (optional: eigenvalues, eigenvectors of dynamical matrix)
281 Notes
282 -----
283 Excite phonon modes to specified temperature.
285 This excites all phonon modes randomly so that each contributes,
286 on average, equally to the given temperature. Both potential
287 energy and kinetic energy will be consistent with the phononic
288 vibrations characteristic of the specified temperature.
290 In other words the system will be equilibrated for an MD run at
291 that temperature.
293 force_constants should be the matrix as force constants, e.g.,
294 as computed by the ase.phonons module.
296 Let X_ai be the phonon modes indexed by atom and mode, w_i the
297 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
298 uniformly random numbers. Then
300 .. code-block:: none
303 1/2
304 _ / k T \ --- 1 _ 1/2
305 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
306 a \ m / --- w ai i i
307 a i i
310 1/2
311 _ / k T \ --- _ 1/2
312 v = | --- | > X (-2 ln Q ) sin (2 pi R )
313 a \ m / --- ai i i
314 a i
316 References
317 ----------
318 D. West and S. K. Estreicher, Phys. Rev. Lett. 96, 115504 (2006).
319 https://doi.org/10.1103/PhysRevLett.96.115504
321 """
323 # Handle the temperature units
324 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
326 # Build dynamical matrix
327 rminv = (masses**-0.5).repeat(3)
328 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
330 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
331 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
333 # Check for soft modes
334 if failfast:
335 zeros = w2_s[:3]
336 worst_zero = np.abs(zeros).max()
337 if worst_zero > 1e-3:
338 msg = 'Translational deviate from 0 significantly: '
339 raise ValueError(msg + f'{w2_s[:3]}')
341 w2min = w2_s[3:].min()
342 if w2min < 0:
343 msg = 'Dynamical matrix has negative eigenvalues such as '
344 raise ValueError(msg + f'{w2min}')
346 # First three modes are translational so ignore:
347 nw = len(w2_s) - 3
348 n_atoms = len(masses)
349 w_s = np.sqrt(w2_s[3:])
350 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
352 # Assign the amplitudes according to Bose-Einstein distribution
353 # or high temperature (== classical) limit
354 if quantum:
355 hbar = units._hbar * units.J * units.s
356 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
357 else:
358 A_s = np.sqrt(temp) / w_s
360 if plus_minus:
361 # create samples by multiplying the amplitude with +/-
362 # according to Eq. 5 in PRB 94, 075125
364 spread = (-1) ** np.arange(nw)
366 # gauge eigenvectors: largest value always positive
367 for ii in range(X_acs.shape[-1]):
368 vec = X_acs[:, :, ii]
369 max_arg = np.argmax(abs(vec))
370 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
372 # Create velocities und displacements from the amplitudes and
373 # eigenvectors
374 A_s *= spread
375 phi_s = 2.0 * np.pi * rng(nw)
377 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
378 # amplitude for displacement
379 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
380 v_ac /= np.sqrt(masses)[:, None]
382 # Assign displacements
383 d_ac = (A_s * X_acs).sum(axis=2)
384 d_ac /= np.sqrt(masses)[:, None]
386 else:
387 # compute the gaussian distribution for the amplitudes
388 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
389 # to avoid (highly improbable) NaN.
391 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
392 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
394 # assign amplitudes and phases
395 A_s *= spread
396 phi_s = 2.0 * np.pi * rng(nw)
398 # Assign velocities and displacements
399 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
400 v_ac /= np.sqrt(masses)[:, None]
402 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
403 d_ac /= np.sqrt(masses)[:, None]
405 if return_eigensolution:
406 return d_ac, v_ac, w2_s, X_is
407 # else
408 return d_ac, v_ac
411def PhononHarmonics(
412 atoms,
413 force_constants,
414 temp=None,
415 *,
416 temperature_K=None,
417 rng=np.random,
418 quantum=False,
419 plus_minus=False,
420 return_eigensolution=False,
421 failfast=True,
422):
423 r"""Excite phonon modes to specified temperature.
425 This will displace atomic positions and set the velocities so as
426 to produce a random, phononically correct state with the requested
427 temperature.
429 Parameters
430 ----------
431 atoms: :class:`~ase.Atoms`
432 Positions and momenta of this object are perturbed.
434 force_constants: ndarray of size 3N x 3N
435 Force constants for the the structure represented by atoms in eV/Ų
437 temp: float (deprecated).
438 Temperature in eV. Deprecated, use ``temperature_K`` instead.
440 temperature_K: float
441 Temperature in Kelvin.
443 rng: Random number generator
444 RandomState or other random number generator, e.g., np.random.rand
446 quantum: bool
447 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
448 (classical limit)
450 failfast: bool
451 True for sanity checking the phonon spectrum for negative frequencies
452 at Gamma.
454 """
456 # Receive displacements and velocities from phonon_harmonics()
457 d_ac, v_ac = phonon_harmonics(
458 force_constants=force_constants,
459 masses=atoms.get_masses(),
460 temp=temp,
461 temperature_K=temperature_K,
462 rng=rng.rand,
463 plus_minus=plus_minus,
464 quantum=quantum,
465 failfast=failfast,
466 return_eigensolution=False,
467 )
469 # Assign new positions (with displacements) and velocities
470 atoms.positions += d_ac
471 atoms.set_velocities(v_ac)