"""Scales or otherwise modifies the particle velocities during simulation to
simulate coupling to an external heat bath with temperature T₀.
"""
import numpy as np
import logging
from mpi4py import MPI
from typing import Callable
from .logger import Logger
from .input_parser import Config
def cancel_com_momentum(velocities, config, comm=MPI.COMM_WORLD):
com_velocity = comm.allreduce(np.sum(velocities[...], axis=0), MPI.SUM)
velocities[...] = velocities[...] - com_velocity / config.n_particles
return velocities
def generate_initial_velocities(velocities, config, prng, comm=MPI.COMM_WORLD):
kT_start = config.gas_constant * config.start_temperature
n_particles_ = velocities.shape[0]
velocities[...] = prng.normal(
loc=0, scale=kT_start / config.mass, size=(n_particles_, 3)
)
velocities = cancel_com_momentum(velocities, config, comm=comm)
kinetic_energy = comm.allreduce(
0.5 * config.mass * np.sum(velocities**2), MPI.SUM
)
start_kinetic_energy_target = (
1.5 * config.gas_constant * config.n_particles * config.start_temperature
)
factor = np.sqrt(1.5 * config.n_particles * kT_start / kinetic_energy)
velocities[...] = velocities[...] * factor
kinetic_energy = comm.allreduce(
0.5 * config.mass * np.sum(velocities**2), MPI.SUM
)
Logger.rank0.log(
logging.INFO,
(
f"Initialized {config.n_particles} velocities, target kinetic "
f"energy: {start_kinetic_energy_target}, actual kinetic energy "
f"generated: {kinetic_energy}"
),
)
return velocities
[docs]def _random_gaussian(prng: np.random.Generator) -> float:
"""Draw a single random number from the standard normal distribution
Generate a number from the Gaussian distribution centered at zero with unit
standard deviation, :math:`N(0, 1)`.
Parameters
----------
prng : np.random.Generator
Numpy object that provides a stream of random bits
Returns
-------
float
A random number drawn from :math:`N(0, 1)`.
"""
return prng.normal()
[docs]def _random_chi_squared(prng: np.random.Generator, M: int) -> float:
"""Draw the sum of :code:`M` squared normally distributed values
The value is generated by the Gamma distribution, in lieu of generating
:code:`M` Gaussian distributed numbers and summing their squares.
Parameters
----------
prng : np.random.Generator
Numpy object that provides a stream of random bits
M : int
Number of standard normally distributed numbers in the sum.
Returns
-------
float
The sum of :code:`M` squared normally distributed values centered at
zero with unit standard deviation.
Notes
-----
The sum of the squares of k independent standard normal random variables is
distributed according to a :math:`\\chi^2` distribution.
.. math::
\\chi^2_1 + \\chi_2^2 + \\chi_3^2 + \\dots + \\chi_M^2
\\sim
\\sigma^2\\chi^2(k)
This is a special case of the :math:`\\Gamma` distribution,
:math:`\\Gamma(k/2, 2)`, and may be generated by
.. math::
\\chi^2(k) \\sim 2 \\Gamma(k/2, 2)
References
----------
Knuth, D.E. 1981, Seminumerical Algorithms, 2nd ed., vol. 2 of The Art of
Computer Programming (Reading, MA: Addison-Wesley), pp. 120ff.
J. H. Ahrens and U. Dieter, Computing 12 (1974), 223-246.
"""
return prng.chisquare(M)
[docs]def csvr_thermostat(
velocity: np.ndarray,
names: np.ndarray,
config: Config,
prng: np.random.Generator,
comm: MPI.Intracomm = MPI.COMM_WORLD,
random_gaussian: Callable[[np.random.Generator], float] = _random_gaussian,
random_chi_squared: Callable[
[np.random.Generator, int, float], float
] = _random_chi_squared,
remove_center_of_mass_momentum: bool = True,
) -> np.ndarray:
"""Canonical sampling through velocity rescaling thermostat
Implements the CSVR thermostat. Rescales the system kinetic energy by a
stochastically chosen factor to keep the temperature constant. Requires
communcation of the kinetic energies calculated locally for each MPI rank.
The random numbers sampled through :code:`random_gaussian` and
:code:`random_chi_squared` are broadcast from the root rank to the other
ranks to ensure the scaling is performed with the same stochasticity for
all particles in the full system.
The velocities are cleaned of center of mass momentum before the thermostat
is applied, and the center of mass momentum is subsequently reapplied. This
is performed for each thermostat coupling group, i.e. the center of mass
momenta of each *group* is separately removed and reapplied after
thermostatting.
The implementation here is based on the derivation presented in the 2008
Comput. Phys. Commun paper, not in the original 2007 J. Chem. Phys. paper.
Parameters
----------
velocity : (N, D) numpy.ndarray
Array of velocities of N particles in D dimensions.
names : (N,) numpy.ndarray
Array of particle names.
config : Config
Configuration dataclass containing simulation metadata and parameters.
prng : np.random.Generator
Numpy object that provides a stream of random bits
comm : MPI.Intracomm, optional
MPI communicator to use for rank commuication. Defaults to
MPI.COMM_WORLD.
random_gaussian : callable, optional
Function for generating standard normally distributed numbers.
random_chi_squared : callable, optional
Function for generating :math:`\\chi^2`-distributed numbers
remove_center_of_mass_momentum : bool, optional
If True, the center of mass of each coupling group is removed before
the thermostat is applied. The center of mass momenta are added back
after the kinetic energy rescaling is complete.
See Also
--------
_random_gaussian :
Used to sample Gaussian-distributed numbers.
_random_chi_squared :
Used to sample :math:`\\chi^2`-distributed numbers.
hymd.input_parser.Config :
Configuration dataclass handler.
References
----------
G. Bussi, D. Donadio, and M. Parrinello, J. Chem. Phys. 126, 014101 (2007).
G. Bussi and M. Parrinello, Comput. Phys. Commun. 179, 26-29, (2008).
"""
if not any(config.thermostat_coupling_groups):
config.thermostat_coupling_groups = [config.unique_names.copy()]
for i, group in enumerate(config.thermostat_coupling_groups):
ind = np.where(
np.logical_or.reduce(list(names == np.string_(t) for t in group))
)
group_n_particles = comm.allreduce(len(ind[0]), MPI.SUM)
# Clean velocities of center of mass momentum
if remove_center_of_mass_momentum and group_n_particles > 1:
com_velocity = comm.allreduce(np.sum(velocity[ind], axis=0))
velocity_clean = velocity[ind] - com_velocity / group_n_particles
K = comm.allreduce(0.5 * config.mass * np.sum(velocity_clean[...] ** 2))
else:
K = comm.allreduce(0.5 * config.mass * np.sum(velocity[...] ** 2))
K_target = (
1.5 * config.gas_constant * group_n_particles * config.target_temperature
)
N_f = 3 * group_n_particles
c = np.exp(-(config.time_step * config.respa_inner) / config.tau)
# Random numbers are identical across ranks because all ranks have
# the same prng after the initialization of the velocities
R = random_gaussian(prng)
SNf = random_chi_squared(prng, N_f - 1)
alpha2 = (
c
+ (1 - c) * (SNf + R**2) * K_target / (N_f * K)
+ 2 * R * np.sqrt(c * (1 - c) * K_target / (N_f * K))
)
dK = K * (alpha2 - 1)
alpha = np.sqrt(alpha2)
if remove_center_of_mass_momentum and group_n_particles > 1:
# Assign velocities and reapply the previously removed center
# of mass momentum removed
velocity_clean *= alpha
velocity[ind] = velocity_clean + com_velocity / group_n_particles
else:
velocity *= alpha
config.thermostat_work += dK