Coverage for /builds/hweiske/ase/ase/md/velocitydistribution.py: 82.46%
114 statements
« prev ^ index » next coverage.py v7.2.7, created at 2024-04-22 11:22 +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 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, temperature: float, unit: str = "K"):
27 """Force the (nuclear) temperature to a precise value.
29 Parameters:
30 atoms: ase.Atoms
31 the structure
32 temperature: float
33 nuclear temperature to set
34 unit: str
35 'K' or 'eV' as unit for the temperature
36 """
38 if unit == "K":
39 E_temp = temperature * units.kB
40 elif unit == "eV":
41 E_temp = temperature
42 else:
43 raise UnitError(f"'{unit}' is not supported, use 'K' or 'eV'.")
45 if temperature > eps_temp:
46 E_kin0 = atoms.get_kinetic_energy() / len(atoms) / 1.5
47 gamma = E_temp / E_kin0
48 else:
49 gamma = 0.0
50 atoms.set_momenta(atoms.get_momenta() * np.sqrt(gamma))
53def _maxwellboltzmanndistribution(masses, temp, communicator=None, rng=None):
54 """Return a Maxwell-Boltzmann distribution with a given temperature.
56 Paremeters:
58 masses: float
59 The atomic masses.
61 temp: float
62 The temperature in electron volt.
64 communicator: MPI communicator (optional)
65 Communicator used to distribute an identical distribution to
66 all tasks. Set to 'serial' to disable communication (setting to None
67 gives the default). Default: ase.parallel.world
69 rng: numpy RNG (optional)
70 The random number generator. Default: np.random
72 Returns:
74 A numpy array with Maxwell-Boltzmann distributed momenta.
75 """
76 if rng is None:
77 rng = np.random
78 if communicator is None:
79 communicator = world
80 xi = rng.standard_normal((len(masses), 3))
81 if communicator != 'serial':
82 communicator.broadcast(xi, 0)
83 momenta = xi * np.sqrt(masses * temp)[:, np.newaxis]
84 return momenta
87def MaxwellBoltzmannDistribution(
88 atoms: Atoms,
89 temp: Optional[float] = None,
90 *,
91 temperature_K: Optional[float] = None,
92 communicator=None,
93 force_temp: bool = False,
94 rng=None,
95):
96 """Set the atomic momenta to a Maxwell-Boltzmann distribution.
98 Parameters:
100 atoms: Atoms object
101 The atoms. Their momenta will be modified.
103 temp: float (deprecated)
104 The temperature in eV. Deprecated, use temperature_K instead.
106 temperature_K: float
107 The temperature in Kelvin.
109 communicator: MPI communicator (optional)
110 Communicator used to distribute an identical distribution to
111 all tasks. Set to 'serial' to disable communication. Leave as None to
112 get the default: ase.parallel.world
114 force_temp: bool (optional, default: False)
115 If True, the random momenta are rescaled so the kinetic energy is
116 exactly 3/2 N k T. This is a slight deviation from the correct
117 Maxwell-Boltzmann distribution.
119 rng: Numpy RNG (optional)
120 Random number generator. Default: numpy.random
121 If you would like to always get the identical distribution, you can
122 supply a random seed like `rng=numpy.random.RandomState(seed)`, where
123 seed is an integer.
124 """
125 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
126 masses = atoms.get_masses()
127 momenta = _maxwellboltzmanndistribution(masses, temp, communicator, rng)
128 atoms.set_momenta(momenta)
129 if force_temp:
130 force_temperature(atoms, temperature=temp, unit="eV")
133def Stationary(atoms: Atoms, preserve_temperature: bool = True):
134 "Sets the center-of-mass momentum to zero."
136 # Save initial temperature
137 temp0 = atoms.get_temperature()
139 p = atoms.get_momenta()
140 p0 = np.sum(p, 0)
141 # We should add a constant velocity, not momentum, to the atoms
142 m = atoms.get_masses()
143 mtot = np.sum(m)
144 v0 = p0 / mtot
145 p -= v0 * m[:, np.newaxis]
146 atoms.set_momenta(p)
148 if preserve_temperature:
149 force_temperature(atoms, temp0)
152def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
153 "Sets the total angular momentum to zero by counteracting rigid rotations."
155 # Save initial temperature
156 temp0 = atoms.get_temperature()
158 # Find the principal moments of inertia and principal axes basis vectors
159 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
160 # Calculate the total angular momentum and transform to principal basis
161 Lp = np.dot(basis, atoms.get_angular_momentum())
162 # Calculate the rotation velocity vector in the principal basis, avoiding
163 # zero division, and transform it back to the cartesian coordinate system
164 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
165 # We subtract a rigid rotation corresponding to this rotation vector
166 com = atoms.get_center_of_mass()
167 positions = atoms.get_positions()
168 positions -= com # translate center of mass to origin
169 velocities = atoms.get_velocities()
170 atoms.set_velocities(velocities - np.cross(omega, positions))
172 if preserve_temperature:
173 force_temperature(atoms, temp0)
176def n_BE(temp, omega):
177 """Bose-Einstein distribution function.
179 Args:
180 temp: temperature converted to eV (*units.kB)
181 omega: sequence of frequencies converted to eV
183 Returns:
184 Value of Bose-Einstein distribution function for each energy
186 """
188 omega = np.asarray(omega)
190 # 0K limit
191 if temp < eps_temp:
192 n = np.zeros_like(omega)
193 else:
194 n = 1 / (np.exp(omega / (temp)) - 1)
195 return n
198def phonon_harmonics(
199 force_constants,
200 masses,
201 temp=None,
202 *,
203 temperature_K=None,
204 rng=np.random.rand,
205 quantum=False,
206 plus_minus=False,
207 return_eigensolution=False,
208 failfast=True,
209):
210 r"""Return displacements and velocities that produce a given temperature.
212 Parameters:
214 force_constants: array of size 3N x 3N
215 force constants (Hessian) of the system in eV/Ų
217 masses: array of length N
218 masses of the structure in amu
220 temp: float (deprecated)
221 Temperature converted to eV (T * units.kB). Deprecated, use
222 ``temperature_K``.
224 temperature_K: float
225 Temperature in Kelvin.
227 rng: function
228 Random number generator function, e.g., np.random.rand
230 quantum: bool
231 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
232 (classical limit)
234 plus_minus: bool
235 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
237 return_eigensolution: bool
238 return eigenvalues and eigenvectors of the dynamical matrix
240 failfast: bool
241 True for sanity checking the phonon spectrum for negative
242 frequencies at Gamma
244 Returns:
246 Displacements, velocities generated from the eigenmodes,
247 (optional: eigenvalues, eigenvectors of dynamical matrix)
249 Purpose:
251 Excite phonon modes to specified temperature.
253 This excites all phonon modes randomly so that each contributes,
254 on average, equally to the given temperature. Both potential
255 energy and kinetic energy will be consistent with the phononic
256 vibrations characteristic of the specified temperature.
258 In other words the system will be equilibrated for an MD run at
259 that temperature.
261 force_constants should be the matrix as force constants, e.g.,
262 as computed by the ase.phonons module.
264 Let X_ai be the phonon modes indexed by atom and mode, w_i the
265 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
266 uniformly random numbers. Then
268 .. code-block:: none
271 1/2
272 _ / k T \ --- 1 _ 1/2
273 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
274 a \ m / --- w ai i i
275 a i i
278 1/2
279 _ / k T \ --- _ 1/2
280 v = | --- | > X (-2 ln Q ) sin (2 pi R )
281 a \ m / --- ai i i
282 a i
284 Reference: [West, Estreicher; PRL 96, 22 (2006)]
285 """
287 # Handle the temperature units
288 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
290 # Build dynamical matrix
291 rminv = (masses ** -0.5).repeat(3)
292 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
294 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
295 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
297 # Check for soft modes
298 if failfast:
299 zeros = w2_s[:3]
300 worst_zero = np.abs(zeros).max()
301 if worst_zero > 1e-3:
302 msg = "Translational deviate from 0 significantly: "
303 raise ValueError(msg + f"{w2_s[:3]}")
305 w2min = w2_s[3:].min()
306 if w2min < 0:
307 msg = "Dynamical matrix has negative eigenvalues such as "
308 raise ValueError(msg + f"{w2min}")
310 # First three modes are translational so ignore:
311 nw = len(w2_s) - 3
312 n_atoms = len(masses)
313 w_s = np.sqrt(w2_s[3:])
314 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
316 # Assign the amplitudes according to Bose-Einstein distribution
317 # or high temperature (== classical) limit
318 if quantum:
319 hbar = units._hbar * units.J * units.s
320 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
321 else:
322 A_s = np.sqrt(temp) / w_s
324 if plus_minus:
325 # create samples by multiplying the amplitude with +/-
326 # according to Eq. 5 in PRB 94, 075125
328 spread = (-1) ** np.arange(nw)
330 # gauge eigenvectors: largest value always positive
331 for ii in range(X_acs.shape[-1]):
332 vec = X_acs[:, :, ii]
333 max_arg = np.argmax(abs(vec))
334 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
336 # Create velocities und displacements from the amplitudes and
337 # eigenvectors
338 A_s *= spread
339 phi_s = 2.0 * np.pi * rng(nw)
341 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
342 # amplitude for displacement
343 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
344 v_ac /= np.sqrt(masses)[:, None]
346 # Assign displacements
347 d_ac = (A_s * X_acs).sum(axis=2)
348 d_ac /= np.sqrt(masses)[:, None]
350 else:
351 # compute the gaussian distribution for the amplitudes
352 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
353 # to avoid (highly improbable) NaN.
355 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
356 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
358 # assign amplitudes and phases
359 A_s *= spread
360 phi_s = 2.0 * np.pi * rng(nw)
362 # Assign velocities and displacements
363 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
364 v_ac /= np.sqrt(masses)[:, None]
366 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
367 d_ac /= np.sqrt(masses)[:, None]
369 if return_eigensolution:
370 return d_ac, v_ac, w2_s, X_is
371 # else
372 return d_ac, v_ac
375def PhononHarmonics(
376 atoms,
377 force_constants,
378 temp=None,
379 *,
380 temperature_K=None,
381 rng=np.random,
382 quantum=False,
383 plus_minus=False,
384 return_eigensolution=False,
385 failfast=True,
386):
387 r"""Excite phonon modes to specified temperature.
389 This will displace atomic positions and set the velocities so as
390 to produce a random, phononically correct state with the requested
391 temperature.
393 Parameters:
395 atoms: ase.atoms.Atoms() object
396 Positions and momenta of this object are perturbed.
398 force_constants: ndarray of size 3N x 3N
399 Force constants for the the structure represented by atoms in eV/Ų
401 temp: float (deprecated).
402 Temperature in eV. Deprecated, use ``temperature_K`` instead.
404 temperature_K: float
405 Temperature in Kelvin.
407 rng: Random number generator
408 RandomState or other random number generator, e.g., np.random.rand
410 quantum: bool
411 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
412 (classical limit)
414 failfast: bool
415 True for sanity checking the phonon spectrum for negative frequencies
416 at Gamma.
418 """
420 # Receive displacements and velocities from phonon_harmonics()
421 d_ac, v_ac = phonon_harmonics(
422 force_constants=force_constants,
423 masses=atoms.get_masses(),
424 temp=temp,
425 temperature_K=temperature_K,
426 rng=rng.rand,
427 plus_minus=plus_minus,
428 quantum=quantum,
429 failfast=failfast,
430 return_eigensolution=False,
431 )
433 # Assign new positions (with displacements) and velocities
434 atoms.positions += d_ac
435 atoms.set_velocities(v_ac)