Coverage for /builds/ericyuan00000/ase/ase/md/velocitydistribution.py: 82.46%
114 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-06-18 01:20 +0000
1# fmt: off
3# VelocityDistributions.py -- set up a velocity distribution
5"""Module for setting up velocity distributions such as Maxwell–Boltzmann.
7Currently, only a few functions are defined, such as
8MaxwellBoltzmannDistribution, which sets the momenta of a list of
9atoms according to a Maxwell-Boltzmann distribution at a given
10temperature.
11"""
12from typing import Literal, Optional
14import numpy as np
16from ase import Atoms, units
17from ase.md.md import process_temperature
18from ase.parallel import world
20# define a ``zero'' temperature to avoid divisions by zero
21eps_temp = 1e-12
24class UnitError(Exception):
25 """Exception raised when wrong units are specified"""
28def force_temperature(atoms: Atoms,
29 temperature: float,
30 unit: Literal["K", "eV"] = "K"):
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, communicator=None, rng=None):
68 """Return a Maxwell-Boltzmann distribution with a given temperature.
70 Paremeters:
72 masses: float
73 The atomic masses.
75 temp: float
76 The temperature in electron volt.
78 communicator: MPI communicator (optional)
79 Communicator used to distribute an identical distribution to
80 all tasks. Set to 'serial' to disable communication (setting to None
81 gives the default). Default: ase.parallel.world
83 rng: numpy RNG (optional)
84 The random number generator. Default: np.random
86 Returns:
88 A numpy array with Maxwell-Boltzmann distributed momenta.
89 """
90 if rng is None:
91 rng = np.random
92 if communicator is None:
93 communicator = world
94 xi = rng.standard_normal((len(masses), 3))
95 if communicator != 'serial':
96 communicator.broadcast(xi, 0)
97 momenta = xi * np.sqrt(masses * temp)[:, np.newaxis]
98 return momenta
101def MaxwellBoltzmannDistribution(
102 atoms: Atoms,
103 temp: Optional[float] = None,
104 *,
105 temperature_K: Optional[float] = None,
106 communicator=None,
107 force_temp: bool = False,
108 rng=None,
109):
110 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
112 Parameters:
114 atoms: Atoms object
115 The atoms. Their momenta will be modified.
117 temp: float (deprecated)
118 The temperature in eV. Deprecated, use ``temperature_K`` instead.
120 temperature_K: float
121 The temperature in Kelvin.
123 communicator: MPI communicator (optional)
124 Communicator used to distribute an identical distribution to
125 all tasks. Set to 'serial' to disable communication. Leave as None to
126 get the default: ase.parallel.world
128 force_temp: bool (optional, default: False)
129 If True, the random momenta are rescaled so the kinetic energy is
130 exactly 3/2 N k T. This is a slight deviation from the correct
131 Maxwell-Boltzmann distribution.
133 rng: Numpy RNG (optional)
134 Random number generator. Default: numpy.random
135 If you would like to always get the identical distribution, you can
136 supply a random seed like ``rng=numpy.random.RandomState(seed)``, where
137 seed is an integer.
138 """
139 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
140 masses = atoms.get_masses()
141 momenta = _maxwellboltzmanndistribution(masses, temp, communicator, rng)
142 atoms.set_momenta(momenta)
143 if force_temp:
144 force_temperature(atoms, temperature=temp, unit="eV")
147def Stationary(atoms: Atoms, preserve_temperature: bool = True):
148 "Sets the center-of-mass momentum to zero."
150 # Save initial temperature
151 temp0 = atoms.get_temperature()
153 p = atoms.get_momenta()
154 p0 = np.sum(p, 0)
155 # We should add a constant velocity, not momentum, to the atoms
156 m = atoms.get_masses()
157 mtot = np.sum(m)
158 v0 = p0 / mtot
159 p -= v0 * m[:, np.newaxis]
160 atoms.set_momenta(p)
162 if preserve_temperature:
163 force_temperature(atoms, temp0)
166def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
167 "Sets the total angular momentum to zero by counteracting rigid rotations."
169 # Save initial temperature
170 temp0 = atoms.get_temperature()
172 # Find the principal moments of inertia and principal axes basis vectors
173 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
174 # Calculate the total angular momentum and transform to principal basis
175 Lp = np.dot(basis, atoms.get_angular_momentum())
176 # Calculate the rotation velocity vector in the principal basis, avoiding
177 # zero division, and transform it back to the cartesian coordinate system
178 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
179 # We subtract a rigid rotation corresponding to this rotation vector
180 com = atoms.get_center_of_mass()
181 positions = atoms.get_positions()
182 positions -= com # translate center of mass to origin
183 velocities = atoms.get_velocities()
184 atoms.set_velocities(velocities - np.cross(omega, positions))
186 if preserve_temperature:
187 force_temperature(atoms, temp0)
190def n_BE(temp, omega):
191 """Bose-Einstein distribution function.
193 Args:
194 temp: temperature converted to eV (*units.kB)
195 omega: sequence of frequencies converted to eV
197 Returns:
198 Value of Bose-Einstein distribution function for each energy
200 """
202 omega = np.asarray(omega)
204 # 0K limit
205 if temp < eps_temp:
206 n = np.zeros_like(omega)
207 else:
208 n = 1 / (np.exp(omega / (temp)) - 1)
209 return n
212def phonon_harmonics(
213 force_constants,
214 masses,
215 temp=None,
216 *,
217 temperature_K=None,
218 rng=np.random.rand,
219 quantum=False,
220 plus_minus=False,
221 return_eigensolution=False,
222 failfast=True,
223):
224 r"""Return displacements and velocities that produce a given temperature.
226 Parameters:
228 force_constants: array of size 3N x 3N
229 force constants (Hessian) of the system in eV/Ų
231 masses: array of length N
232 masses of the structure in amu
234 temp: float (deprecated)
235 Temperature converted to eV (T * units.kB). Deprecated, use
236 ``temperature_K``.
238 temperature_K: float
239 Temperature in Kelvin.
241 rng: function
242 Random number generator function, e.g., np.random.rand
244 quantum: bool
245 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
246 (classical limit)
248 plus_minus: bool
249 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
251 return_eigensolution: bool
252 return eigenvalues and eigenvectors of the dynamical matrix
254 failfast: bool
255 True for sanity checking the phonon spectrum for negative
256 frequencies at Gamma
258 Returns:
260 Displacements, velocities generated from the eigenmodes,
261 (optional: eigenvalues, eigenvectors of dynamical matrix)
263 Purpose:
265 Excite phonon modes to specified temperature.
267 This excites all phonon modes randomly so that each contributes,
268 on average, equally to the given temperature. Both potential
269 energy and kinetic energy will be consistent with the phononic
270 vibrations characteristic of the specified temperature.
272 In other words the system will be equilibrated for an MD run at
273 that temperature.
275 force_constants should be the matrix as force constants, e.g.,
276 as computed by the ase.phonons module.
278 Let X_ai be the phonon modes indexed by atom and mode, w_i the
279 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
280 uniformly random numbers. Then
282 .. code-block:: none
285 1/2
286 _ / k T \ --- 1 _ 1/2
287 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
288 a \ m / --- w ai i i
289 a i i
292 1/2
293 _ / k T \ --- _ 1/2
294 v = | --- | > X (-2 ln Q ) sin (2 pi R )
295 a \ m / --- ai i i
296 a i
298 Reference: [West, Estreicher; PRL 96, 22 (2006)]
299 """
301 # Handle the temperature units
302 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
304 # Build dynamical matrix
305 rminv = (masses ** -0.5).repeat(3)
306 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
308 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
309 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
311 # Check for soft modes
312 if failfast:
313 zeros = w2_s[:3]
314 worst_zero = np.abs(zeros).max()
315 if worst_zero > 1e-3:
316 msg = "Translational deviate from 0 significantly: "
317 raise ValueError(msg + f"{w2_s[:3]}")
319 w2min = w2_s[3:].min()
320 if w2min < 0:
321 msg = "Dynamical matrix has negative eigenvalues such as "
322 raise ValueError(msg + f"{w2min}")
324 # First three modes are translational so ignore:
325 nw = len(w2_s) - 3
326 n_atoms = len(masses)
327 w_s = np.sqrt(w2_s[3:])
328 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
330 # Assign the amplitudes according to Bose-Einstein distribution
331 # or high temperature (== classical) limit
332 if quantum:
333 hbar = units._hbar * units.J * units.s
334 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
335 else:
336 A_s = np.sqrt(temp) / w_s
338 if plus_minus:
339 # create samples by multiplying the amplitude with +/-
340 # according to Eq. 5 in PRB 94, 075125
342 spread = (-1) ** np.arange(nw)
344 # gauge eigenvectors: largest value always positive
345 for ii in range(X_acs.shape[-1]):
346 vec = X_acs[:, :, ii]
347 max_arg = np.argmax(abs(vec))
348 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
350 # Create velocities und displacements from the amplitudes and
351 # eigenvectors
352 A_s *= spread
353 phi_s = 2.0 * np.pi * rng(nw)
355 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
356 # amplitude for displacement
357 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
358 v_ac /= np.sqrt(masses)[:, None]
360 # Assign displacements
361 d_ac = (A_s * X_acs).sum(axis=2)
362 d_ac /= np.sqrt(masses)[:, None]
364 else:
365 # compute the gaussian distribution for the amplitudes
366 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
367 # to avoid (highly improbable) NaN.
369 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
370 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
372 # assign amplitudes and phases
373 A_s *= spread
374 phi_s = 2.0 * np.pi * rng(nw)
376 # Assign velocities and displacements
377 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
378 v_ac /= np.sqrt(masses)[:, None]
380 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
381 d_ac /= np.sqrt(masses)[:, None]
383 if return_eigensolution:
384 return d_ac, v_ac, w2_s, X_is
385 # else
386 return d_ac, v_ac
389def PhononHarmonics(
390 atoms,
391 force_constants,
392 temp=None,
393 *,
394 temperature_K=None,
395 rng=np.random,
396 quantum=False,
397 plus_minus=False,
398 return_eigensolution=False,
399 failfast=True,
400):
401 r"""Excite phonon modes to specified temperature.
403 This will displace atomic positions and set the velocities so as
404 to produce a random, phononically correct state with the requested
405 temperature.
407 Parameters:
409 atoms: ase.atoms.Atoms() object
410 Positions and momenta of this object are perturbed.
412 force_constants: ndarray of size 3N x 3N
413 Force constants for the the structure represented by atoms in eV/Ų
415 temp: float (deprecated).
416 Temperature in eV. Deprecated, use ``temperature_K`` instead.
418 temperature_K: float
419 Temperature in Kelvin.
421 rng: Random number generator
422 RandomState or other random number generator, e.g., np.random.rand
424 quantum: bool
425 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
426 (classical limit)
428 failfast: bool
429 True for sanity checking the phonon spectrum for negative frequencies
430 at Gamma.
432 """
434 # Receive displacements and velocities from phonon_harmonics()
435 d_ac, v_ac = phonon_harmonics(
436 force_constants=force_constants,
437 masses=atoms.get_masses(),
438 temp=temp,
439 temperature_K=temperature_K,
440 rng=rng.rand,
441 plus_minus=plus_minus,
442 quantum=quantum,
443 failfast=failfast,
444 return_eigensolution=False,
445 )
447 # Assign new positions (with displacements) and velocities
448 atoms.positions += d_ac
449 atoms.set_velocities(v_ac)