Coverage for /builds/ase/ase/ase/md/velocitydistribution.py: 78.33%
120 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +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, Optional
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: Optional[float] = None,
100 *,
101 temperature_K: Optional[float] = 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 Value of Bose-Einstein distribution function for each energy
217 """
219 omega = np.asarray(omega)
221 # 0K limit
222 if temp < eps_temp:
223 n = np.zeros_like(omega)
224 else:
225 n = 1 / (np.exp(omega / (temp)) - 1)
226 return n
229def phonon_harmonics(
230 force_constants,
231 masses,
232 temp=None,
233 *,
234 temperature_K=None,
235 rng=np.random.rand,
236 quantum=False,
237 plus_minus=False,
238 return_eigensolution=False,
239 failfast=True,
240):
241 r"""Return displacements and velocities that produce a given temperature.
243 Parameters:
245 force_constants: array of size 3N x 3N
246 force constants (Hessian) of the system in eV/Ų
248 masses: array of length N
249 masses of the structure in amu
251 temp: float (deprecated)
252 Temperature converted to eV (T * units.kB). Deprecated, use
253 ``temperature_K``.
255 temperature_K: float
256 Temperature in Kelvin.
258 rng: function
259 Random number generator function, e.g., np.random.rand
261 quantum: bool
262 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
263 (classical limit)
265 plus_minus: bool
266 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
268 return_eigensolution: bool
269 return eigenvalues and eigenvectors of the dynamical matrix
271 failfast: bool
272 True for sanity checking the phonon spectrum for negative
273 frequencies at Gamma
275 Returns:
277 Displacements, velocities generated from the eigenmodes,
278 (optional: eigenvalues, eigenvectors of dynamical matrix)
280 Purpose:
282 Excite phonon modes to specified temperature.
284 This excites all phonon modes randomly so that each contributes,
285 on average, equally to the given temperature. Both potential
286 energy and kinetic energy will be consistent with the phononic
287 vibrations characteristic of the specified temperature.
289 In other words the system will be equilibrated for an MD run at
290 that temperature.
292 force_constants should be the matrix as force constants, e.g.,
293 as computed by the ase.phonons module.
295 Let X_ai be the phonon modes indexed by atom and mode, w_i the
296 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
297 uniformly random numbers. Then
299 .. code-block:: none
302 1/2
303 _ / k T \ --- 1 _ 1/2
304 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
305 a \ m / --- w ai i i
306 a i i
309 1/2
310 _ / k T \ --- _ 1/2
311 v = | --- | > X (-2 ln Q ) sin (2 pi R )
312 a \ m / --- ai i i
313 a i
315 Reference: [West, Estreicher; PRL 96, 22 (2006)]
316 """
318 # Handle the temperature units
319 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
321 # Build dynamical matrix
322 rminv = (masses**-0.5).repeat(3)
323 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
325 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
326 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
328 # Check for soft modes
329 if failfast:
330 zeros = w2_s[:3]
331 worst_zero = np.abs(zeros).max()
332 if worst_zero > 1e-3:
333 msg = 'Translational deviate from 0 significantly: '
334 raise ValueError(msg + f'{w2_s[:3]}')
336 w2min = w2_s[3:].min()
337 if w2min < 0:
338 msg = 'Dynamical matrix has negative eigenvalues such as '
339 raise ValueError(msg + f'{w2min}')
341 # First three modes are translational so ignore:
342 nw = len(w2_s) - 3
343 n_atoms = len(masses)
344 w_s = np.sqrt(w2_s[3:])
345 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
347 # Assign the amplitudes according to Bose-Einstein distribution
348 # or high temperature (== classical) limit
349 if quantum:
350 hbar = units._hbar * units.J * units.s
351 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
352 else:
353 A_s = np.sqrt(temp) / w_s
355 if plus_minus:
356 # create samples by multiplying the amplitude with +/-
357 # according to Eq. 5 in PRB 94, 075125
359 spread = (-1) ** np.arange(nw)
361 # gauge eigenvectors: largest value always positive
362 for ii in range(X_acs.shape[-1]):
363 vec = X_acs[:, :, ii]
364 max_arg = np.argmax(abs(vec))
365 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
367 # Create velocities und displacements from the amplitudes and
368 # eigenvectors
369 A_s *= spread
370 phi_s = 2.0 * np.pi * rng(nw)
372 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
373 # amplitude for displacement
374 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
375 v_ac /= np.sqrt(masses)[:, None]
377 # Assign displacements
378 d_ac = (A_s * X_acs).sum(axis=2)
379 d_ac /= np.sqrt(masses)[:, None]
381 else:
382 # compute the gaussian distribution for the amplitudes
383 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
384 # to avoid (highly improbable) NaN.
386 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
387 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
389 # assign amplitudes and phases
390 A_s *= spread
391 phi_s = 2.0 * np.pi * rng(nw)
393 # Assign velocities and displacements
394 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
395 v_ac /= np.sqrt(masses)[:, None]
397 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
398 d_ac /= np.sqrt(masses)[:, None]
400 if return_eigensolution:
401 return d_ac, v_ac, w2_s, X_is
402 # else
403 return d_ac, v_ac
406def PhononHarmonics(
407 atoms,
408 force_constants,
409 temp=None,
410 *,
411 temperature_K=None,
412 rng=np.random,
413 quantum=False,
414 plus_minus=False,
415 return_eigensolution=False,
416 failfast=True,
417):
418 r"""Excite phonon modes to specified temperature.
420 This will displace atomic positions and set the velocities so as
421 to produce a random, phononically correct state with the requested
422 temperature.
424 Parameters:
426 atoms: ase.atoms.Atoms() object
427 Positions and momenta of this object are perturbed.
429 force_constants: ndarray of size 3N x 3N
430 Force constants for the the structure represented by atoms in eV/Ų
432 temp: float (deprecated).
433 Temperature in eV. Deprecated, use ``temperature_K`` instead.
435 temperature_K: float
436 Temperature in Kelvin.
438 rng: Random number generator
439 RandomState or other random number generator, e.g., np.random.rand
441 quantum: bool
442 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
443 (classical limit)
445 failfast: bool
446 True for sanity checking the phonon spectrum for negative frequencies
447 at Gamma.
449 """
451 # Receive displacements and velocities from phonon_harmonics()
452 d_ac, v_ac = phonon_harmonics(
453 force_constants=force_constants,
454 masses=atoms.get_masses(),
455 temp=temp,
456 temperature_K=temperature_K,
457 rng=rng.rand,
458 plus_minus=plus_minus,
459 quantum=quantum,
460 failfast=failfast,
461 return_eigensolution=False,
462 )
464 # Assign new positions (with displacements) and velocities
465 atoms.positions += d_ac
466 atoms.set_velocities(v_ac)