Coverage for /builds/debichem-team/python-ase/ase/md/velocitydistribution.py: 82.46%
114 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1# VelocityDistributions.py -- set up a velocity distribution
3"""Module for setting up velocity distributions such as Maxwell–Boltzmann.
5Currently, only a few functions are defined, such as
6MaxwellBoltzmannDistribution, which sets the momenta of a list of
7atoms according to a Maxwell-Boltzmann distribution at a given
8temperature.
9"""
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 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(atoms: Atoms,
27 temperature: float,
28 unit: Literal["K", "eV"] = "K"):
29 """
30 Force the temperature of the atomic system to a precise value.
32 This function adjusts the momenta of the atoms to achieve the
33 exact specified temperature, overriding the Maxwell-Boltzmann
34 distribution. The temperature can be specified in either Kelvin
35 (K) or electron volts (eV).
37 Parameters
38 ----------
39 atoms
40 The atomic system represented as an ASE Atoms object.
41 temperature
42 The exact temperature to force for the atomic system.
43 unit
44 The unit of the specified temperature.
45 Can be either 'K' (Kelvin) or 'eV' (electron volts). Default is 'K'.
46 """
48 if unit == "K":
49 target_temp = temperature * units.kB
50 elif unit == "eV":
51 target_temp = temperature
52 else:
53 raise UnitError(f"'{unit}' is not supported, use 'K' or 'eV'.")
55 if temperature > eps_temp:
56 ndof = atoms.get_number_of_degrees_of_freedom()
57 current_temp = 2 * atoms.get_kinetic_energy() / ndof
58 scale = target_temp / current_temp
59 else:
60 scale = 0.0
62 atoms.set_momenta(atoms.get_momenta() * np.sqrt(scale))
65def _maxwellboltzmanndistribution(masses, temp, communicator=None, rng=None):
66 """Return a Maxwell-Boltzmann distribution with a given temperature.
68 Paremeters:
70 masses: float
71 The atomic masses.
73 temp: float
74 The temperature in electron volt.
76 communicator: MPI communicator (optional)
77 Communicator used to distribute an identical distribution to
78 all tasks. Set to 'serial' to disable communication (setting to None
79 gives the default). Default: ase.parallel.world
81 rng: numpy RNG (optional)
82 The random number generator. Default: np.random
84 Returns:
86 A numpy array with Maxwell-Boltzmann distributed momenta.
87 """
88 if rng is None:
89 rng = np.random
90 if communicator is None:
91 communicator = world
92 xi = rng.standard_normal((len(masses), 3))
93 if communicator != 'serial':
94 communicator.broadcast(xi, 0)
95 momenta = xi * np.sqrt(masses * temp)[:, np.newaxis]
96 return momenta
99def MaxwellBoltzmannDistribution(
100 atoms: Atoms,
101 temp: Optional[float] = None,
102 *,
103 temperature_K: Optional[float] = None,
104 communicator=None,
105 force_temp: bool = False,
106 rng=None,
107):
108 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
110 Parameters:
112 atoms: Atoms object
113 The atoms. Their momenta will be modified.
115 temp: float (deprecated)
116 The temperature in eV. Deprecated, use temperature_K instead.
118 temperature_K: float
119 The temperature in Kelvin.
121 communicator: MPI communicator (optional)
122 Communicator used to distribute an identical distribution to
123 all tasks. Set to 'serial' to disable communication. Leave as None to
124 get the default: ase.parallel.world
126 force_temp: bool (optional, default: False)
127 If True, the random momenta are rescaled so the kinetic energy is
128 exactly 3/2 N k T. This is a slight deviation from the correct
129 Maxwell-Boltzmann distribution.
131 rng: Numpy RNG (optional)
132 Random number generator. Default: numpy.random
133 If you would like to always get the identical distribution, you can
134 supply a random seed like `rng=numpy.random.RandomState(seed)`, where
135 seed is an integer.
136 """
137 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
138 masses = atoms.get_masses()
139 momenta = _maxwellboltzmanndistribution(masses, temp, communicator, rng)
140 atoms.set_momenta(momenta)
141 if force_temp:
142 force_temperature(atoms, temperature=temp, unit="eV")
145def Stationary(atoms: Atoms, preserve_temperature: bool = True):
146 "Sets the center-of-mass momentum to zero."
148 # Save initial temperature
149 temp0 = atoms.get_temperature()
151 p = atoms.get_momenta()
152 p0 = np.sum(p, 0)
153 # We should add a constant velocity, not momentum, to the atoms
154 m = atoms.get_masses()
155 mtot = np.sum(m)
156 v0 = p0 / mtot
157 p -= v0 * m[:, np.newaxis]
158 atoms.set_momenta(p)
160 if preserve_temperature:
161 force_temperature(atoms, temp0)
164def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
165 "Sets the total angular momentum to zero by counteracting rigid rotations."
167 # Save initial temperature
168 temp0 = atoms.get_temperature()
170 # Find the principal moments of inertia and principal axes basis vectors
171 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
172 # Calculate the total angular momentum and transform to principal basis
173 Lp = np.dot(basis, atoms.get_angular_momentum())
174 # Calculate the rotation velocity vector in the principal basis, avoiding
175 # zero division, and transform it back to the cartesian coordinate system
176 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
177 # We subtract a rigid rotation corresponding to this rotation vector
178 com = atoms.get_center_of_mass()
179 positions = atoms.get_positions()
180 positions -= com # translate center of mass to origin
181 velocities = atoms.get_velocities()
182 atoms.set_velocities(velocities - np.cross(omega, positions))
184 if preserve_temperature:
185 force_temperature(atoms, temp0)
188def n_BE(temp, omega):
189 """Bose-Einstein distribution function.
191 Args:
192 temp: temperature converted to eV (*units.kB)
193 omega: sequence of frequencies converted to eV
195 Returns:
196 Value of Bose-Einstein distribution function for each energy
198 """
200 omega = np.asarray(omega)
202 # 0K limit
203 if temp < eps_temp:
204 n = np.zeros_like(omega)
205 else:
206 n = 1 / (np.exp(omega / (temp)) - 1)
207 return n
210def phonon_harmonics(
211 force_constants,
212 masses,
213 temp=None,
214 *,
215 temperature_K=None,
216 rng=np.random.rand,
217 quantum=False,
218 plus_minus=False,
219 return_eigensolution=False,
220 failfast=True,
221):
222 r"""Return displacements and velocities that produce a given temperature.
224 Parameters:
226 force_constants: array of size 3N x 3N
227 force constants (Hessian) of the system in eV/Ų
229 masses: array of length N
230 masses of the structure in amu
232 temp: float (deprecated)
233 Temperature converted to eV (T * units.kB). Deprecated, use
234 ``temperature_K``.
236 temperature_K: float
237 Temperature in Kelvin.
239 rng: function
240 Random number generator function, e.g., np.random.rand
242 quantum: bool
243 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
244 (classical limit)
246 plus_minus: bool
247 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
249 return_eigensolution: bool
250 return eigenvalues and eigenvectors of the dynamical matrix
252 failfast: bool
253 True for sanity checking the phonon spectrum for negative
254 frequencies at Gamma
256 Returns:
258 Displacements, velocities generated from the eigenmodes,
259 (optional: eigenvalues, eigenvectors of dynamical matrix)
261 Purpose:
263 Excite phonon modes to specified temperature.
265 This excites all phonon modes randomly so that each contributes,
266 on average, equally to the given temperature. Both potential
267 energy and kinetic energy will be consistent with the phononic
268 vibrations characteristic of the specified temperature.
270 In other words the system will be equilibrated for an MD run at
271 that temperature.
273 force_constants should be the matrix as force constants, e.g.,
274 as computed by the ase.phonons module.
276 Let X_ai be the phonon modes indexed by atom and mode, w_i the
277 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
278 uniformly random numbers. Then
280 .. code-block:: none
283 1/2
284 _ / k T \ --- 1 _ 1/2
285 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
286 a \ m / --- w ai i i
287 a i i
290 1/2
291 _ / k T \ --- _ 1/2
292 v = | --- | > X (-2 ln Q ) sin (2 pi R )
293 a \ m / --- ai i i
294 a i
296 Reference: [West, Estreicher; PRL 96, 22 (2006)]
297 """
299 # Handle the temperature units
300 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
302 # Build dynamical matrix
303 rminv = (masses ** -0.5).repeat(3)
304 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
306 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
307 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
309 # Check for soft modes
310 if failfast:
311 zeros = w2_s[:3]
312 worst_zero = np.abs(zeros).max()
313 if worst_zero > 1e-3:
314 msg = "Translational deviate from 0 significantly: "
315 raise ValueError(msg + f"{w2_s[:3]}")
317 w2min = w2_s[3:].min()
318 if w2min < 0:
319 msg = "Dynamical matrix has negative eigenvalues such as "
320 raise ValueError(msg + f"{w2min}")
322 # First three modes are translational so ignore:
323 nw = len(w2_s) - 3
324 n_atoms = len(masses)
325 w_s = np.sqrt(w2_s[3:])
326 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
328 # Assign the amplitudes according to Bose-Einstein distribution
329 # or high temperature (== classical) limit
330 if quantum:
331 hbar = units._hbar * units.J * units.s
332 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
333 else:
334 A_s = np.sqrt(temp) / w_s
336 if plus_minus:
337 # create samples by multiplying the amplitude with +/-
338 # according to Eq. 5 in PRB 94, 075125
340 spread = (-1) ** np.arange(nw)
342 # gauge eigenvectors: largest value always positive
343 for ii in range(X_acs.shape[-1]):
344 vec = X_acs[:, :, ii]
345 max_arg = np.argmax(abs(vec))
346 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
348 # Create velocities und displacements from the amplitudes and
349 # eigenvectors
350 A_s *= spread
351 phi_s = 2.0 * np.pi * rng(nw)
353 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
354 # amplitude for displacement
355 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
356 v_ac /= np.sqrt(masses)[:, None]
358 # Assign displacements
359 d_ac = (A_s * X_acs).sum(axis=2)
360 d_ac /= np.sqrt(masses)[:, None]
362 else:
363 # compute the gaussian distribution for the amplitudes
364 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
365 # to avoid (highly improbable) NaN.
367 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
368 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
370 # assign amplitudes and phases
371 A_s *= spread
372 phi_s = 2.0 * np.pi * rng(nw)
374 # Assign velocities and displacements
375 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
376 v_ac /= np.sqrt(masses)[:, None]
378 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
379 d_ac /= np.sqrt(masses)[:, None]
381 if return_eigensolution:
382 return d_ac, v_ac, w2_s, X_is
383 # else
384 return d_ac, v_ac
387def PhononHarmonics(
388 atoms,
389 force_constants,
390 temp=None,
391 *,
392 temperature_K=None,
393 rng=np.random,
394 quantum=False,
395 plus_minus=False,
396 return_eigensolution=False,
397 failfast=True,
398):
399 r"""Excite phonon modes to specified temperature.
401 This will displace atomic positions and set the velocities so as
402 to produce a random, phononically correct state with the requested
403 temperature.
405 Parameters:
407 atoms: ase.atoms.Atoms() object
408 Positions and momenta of this object are perturbed.
410 force_constants: ndarray of size 3N x 3N
411 Force constants for the the structure represented by atoms in eV/Ų
413 temp: float (deprecated).
414 Temperature in eV. Deprecated, use ``temperature_K`` instead.
416 temperature_K: float
417 Temperature in Kelvin.
419 rng: Random number generator
420 RandomState or other random number generator, e.g., np.random.rand
422 quantum: bool
423 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
424 (classical limit)
426 failfast: bool
427 True for sanity checking the phonon spectrum for negative frequencies
428 at Gamma.
430 """
432 # Receive displacements and velocities from phonon_harmonics()
433 d_ac, v_ac = phonon_harmonics(
434 force_constants=force_constants,
435 masses=atoms.get_masses(),
436 temp=temp,
437 temperature_K=temperature_K,
438 rng=rng.rand,
439 plus_minus=plus_minus,
440 quantum=quantum,
441 failfast=failfast,
442 return_eigensolution=False,
443 )
445 # Assign new positions (with displacements) and velocities
446 atoms.positions += d_ac
447 atoms.set_velocities(v_ac)