Coverage for /builds/kinetik161/ase/ase/md/velocitydistribution.py: 82.46%
114 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +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 (nucl.) temperature to have 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 """Sets the 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, used 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 (optinal, default: False)
115 If True, random the 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 """
122 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
123 masses = atoms.get_masses()
124 momenta = _maxwellboltzmanndistribution(masses, temp, communicator, rng)
125 atoms.set_momenta(momenta)
126 if force_temp:
127 force_temperature(atoms, temperature=temp, unit="eV")
130def Stationary(atoms: Atoms, preserve_temperature: bool = True):
131 "Sets the center-of-mass momentum to zero."
133 # Save initial temperature
134 temp0 = atoms.get_temperature()
136 p = atoms.get_momenta()
137 p0 = np.sum(p, 0)
138 # We should add a constant velocity, not momentum, to the atoms
139 m = atoms.get_masses()
140 mtot = np.sum(m)
141 v0 = p0 / mtot
142 p -= v0 * m[:, np.newaxis]
143 atoms.set_momenta(p)
145 if preserve_temperature:
146 force_temperature(atoms, temp0)
149def ZeroRotation(atoms: Atoms, preserve_temperature: float = True):
150 "Sets the total angular momentum to zero by counteracting rigid rotations."
152 # Save initial temperature
153 temp0 = atoms.get_temperature()
155 # Find the principal moments of inertia and principal axes basis vectors
156 Ip, basis = atoms.get_moments_of_inertia(vectors=True)
157 # Calculate the total angular momentum and transform to principal basis
158 Lp = np.dot(basis, atoms.get_angular_momentum())
159 # Calculate the rotation velocity vector in the principal basis, avoiding
160 # zero division, and transform it back to the cartesian coordinate system
161 omega = np.dot(np.linalg.inv(basis), np.select([Ip > 0], [Lp / Ip]))
162 # We subtract a rigid rotation corresponding to this rotation vector
163 com = atoms.get_center_of_mass()
164 positions = atoms.get_positions()
165 positions -= com # translate center of mass to origin
166 velocities = atoms.get_velocities()
167 atoms.set_velocities(velocities - np.cross(omega, positions))
169 if preserve_temperature:
170 force_temperature(atoms, temp0)
173def n_BE(temp, omega):
174 """Bose-Einstein distribution function.
176 Args:
177 temp: temperature converted to eV (*units.kB)
178 omega: sequence of frequencies converted to eV
180 Returns:
181 Value of Bose-Einstein distribution function for each energy
183 """
185 omega = np.asarray(omega)
187 # 0K limit
188 if temp < eps_temp:
189 n = np.zeros_like(omega)
190 else:
191 n = 1 / (np.exp(omega / (temp)) - 1)
192 return n
195def phonon_harmonics(
196 force_constants,
197 masses,
198 temp=None,
199 *,
200 temperature_K=None,
201 rng=np.random.rand,
202 quantum=False,
203 plus_minus=False,
204 return_eigensolution=False,
205 failfast=True,
206):
207 r"""Return displacements and velocities that produce a given temperature.
209 Parameters:
211 force_constants: array of size 3N x 3N
212 force constants (Hessian) of the system in eV/Ų
214 masses: array of length N
215 masses of the structure in amu
217 temp: float (deprecated)
218 Temperature converted to eV (T * units.kB). Deprecated, use
219 ``temperature_K``.
221 temperature_K: float
222 Temperature in Kelvin.
224 rng: function
225 Random number generator function, e.g., np.random.rand
227 quantum: bool
228 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
229 (classical limit)
231 plus_minus: bool
232 Displace atoms with +/- the amplitude accoding to PRB 94, 075125
234 return_eigensolution: bool
235 return eigenvalues and eigenvectors of the dynamical matrix
237 failfast: bool
238 True for sanity checking the phonon spectrum for negative
239 frequencies at Gamma
241 Returns:
243 Displacements, velocities generated from the eigenmodes,
244 (optional: eigenvalues, eigenvectors of dynamical matrix)
246 Purpose:
248 Excite phonon modes to specified temperature.
250 This excites all phonon modes randomly so that each contributes,
251 on average, equally to the given temperature. Both potential
252 energy and kinetic energy will be consistent with the phononic
253 vibrations characteristic of the specified temperature.
255 In other words the system will be equilibrated for an MD run at
256 that temperature.
258 force_constants should be the matrix as force constants, e.g.,
259 as computed by the ase.phonons module.
261 Let X_ai be the phonon modes indexed by atom and mode, w_i the
262 phonon frequencies, and let 0 < Q_i <= 1 and 0 <= R_i < 1 be
263 uniformly random numbers. Then
265 .. code-block:: none
268 1/2
269 _ / k T \ --- 1 _ 1/2
270 R += | --- | > --- X (-2 ln Q ) cos (2 pi R )
271 a \ m / --- w ai i i
272 a i i
275 1/2
276 _ / k T \ --- _ 1/2
277 v = | --- | > X (-2 ln Q ) sin (2 pi R )
278 a \ m / --- ai i i
279 a i
281 Reference: [West, Estreicher; PRL 96, 22 (2006)]
282 """
284 # Handle the temperature units
285 temp = units.kB * process_temperature(temp, temperature_K, 'eV')
287 # Build dynamical matrix
288 rminv = (masses ** -0.5).repeat(3)
289 dynamical_matrix = force_constants * rminv[:, None] * rminv[None, :]
291 # Solve eigenvalue problem to compute phonon spectrum and eigenvectors
292 w2_s, X_is = np.linalg.eigh(dynamical_matrix)
294 # Check for soft modes
295 if failfast:
296 zeros = w2_s[:3]
297 worst_zero = np.abs(zeros).max()
298 if worst_zero > 1e-3:
299 msg = "Translational deviate from 0 significantly: "
300 raise ValueError(msg + f"{w2_s[:3]}")
302 w2min = w2_s[3:].min()
303 if w2min < 0:
304 msg = "Dynamical matrix has negative eigenvalues such as "
305 raise ValueError(msg + f"{w2min}")
307 # First three modes are translational so ignore:
308 nw = len(w2_s) - 3
309 n_atoms = len(masses)
310 w_s = np.sqrt(w2_s[3:])
311 X_acs = X_is[:, 3:].reshape(n_atoms, 3, nw)
313 # Assign the amplitudes according to Bose-Einstein distribution
314 # or high temperature (== classical) limit
315 if quantum:
316 hbar = units._hbar * units.J * units.s
317 A_s = np.sqrt(hbar * (2 * n_BE(temp, hbar * w_s) + 1) / (2 * w_s))
318 else:
319 A_s = np.sqrt(temp) / w_s
321 if plus_minus:
322 # create samples by multiplying the amplitude with +/-
323 # according to Eq. 5 in PRB 94, 075125
325 spread = (-1) ** np.arange(nw)
327 # gauge eigenvectors: largest value always positive
328 for ii in range(X_acs.shape[-1]):
329 vec = X_acs[:, :, ii]
330 max_arg = np.argmax(abs(vec))
331 X_acs[:, :, ii] *= np.sign(vec.flat[max_arg])
333 # Create velocities und displacements from the amplitudes and
334 # eigenvectors
335 A_s *= spread
336 phi_s = 2.0 * np.pi * rng(nw)
338 # Assign velocities, sqrt(2) compensates for missing sin(phi) in
339 # amplitude for displacement
340 v_ac = (w_s * A_s * np.sqrt(2) * np.cos(phi_s) * X_acs).sum(axis=2)
341 v_ac /= np.sqrt(masses)[:, None]
343 # Assign displacements
344 d_ac = (A_s * X_acs).sum(axis=2)
345 d_ac /= np.sqrt(masses)[:, None]
347 else:
348 # compute the gaussian distribution for the amplitudes
349 # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
350 # to avoid (highly improbable) NaN.
352 # Box Muller [en.wikipedia.org/wiki/Box–Muller_transform]:
353 spread = np.sqrt(-2.0 * np.log(1.0 - rng(nw)))
355 # assign amplitudes and phases
356 A_s *= spread
357 phi_s = 2.0 * np.pi * rng(nw)
359 # Assign velocities and displacements
360 v_ac = (w_s * A_s * np.cos(phi_s) * X_acs).sum(axis=2)
361 v_ac /= np.sqrt(masses)[:, None]
363 d_ac = (A_s * np.sin(phi_s) * X_acs).sum(axis=2)
364 d_ac /= np.sqrt(masses)[:, None]
366 if return_eigensolution:
367 return d_ac, v_ac, w2_s, X_is
368 # else
369 return d_ac, v_ac
372def PhononHarmonics(
373 atoms,
374 force_constants,
375 temp=None,
376 *,
377 temperature_K=None,
378 rng=np.random,
379 quantum=False,
380 plus_minus=False,
381 return_eigensolution=False,
382 failfast=True,
383):
384 r"""Excite phonon modes to specified temperature.
386 This will displace atomic positions and set the velocities so as
387 to produce a random, phononically correct state with the requested
388 temperature.
390 Parameters:
392 atoms: ase.atoms.Atoms() object
393 Positions and momenta of this object are perturbed.
395 force_constants: ndarray of size 3N x 3N
396 Force constants for the the structure represented by atoms in eV/Ų
398 temp: float (deprecated).
399 Temperature in eV. Deprecated, use ``temperature_K`` instead.
401 temperature_K: float
402 Temperature in Kelvin.
404 rng: Random number generator
405 RandomState or other random number generator, e.g., np.random.rand
407 quantum: bool
408 True for Bose-Einstein distribution, False for Maxwell-Boltzmann
409 (classical limit)
411 failfast: bool
412 True for sanity checking the phonon spectrum for negative frequencies
413 at Gamma.
415 """
417 # Receive displacements and velocities from phonon_harmonics()
418 d_ac, v_ac = phonon_harmonics(
419 force_constants=force_constants,
420 masses=atoms.get_masses(),
421 temp=temp,
422 temperature_K=temperature_K,
423 rng=rng.rand,
424 plus_minus=plus_minus,
425 quantum=quantum,
426 failfast=failfast,
427 return_eigensolution=False,
428 )
430 # Assign new positions (with displacements) and velocities
431 atoms.positions += d_ac
432 atoms.set_velocities(v_ac)