import numpy as np
from pysisyphus.constants import BOHR2ANG, KBAU, VELO2E
from pysisyphus.hessian_proj import get_hessian_projector
from pysisyphus.xyzloader import make_trj_str
[docs]
def dump_coords(atoms, coords, trj_fn):
coords = np.array(coords)
coords = coords.reshape(-1, len(atoms), 3) * BOHR2ANG
trj_str = make_trj_str(atoms, coords)
with open(trj_fn, "w") as handle:
handle.write(trj_str)
[docs]
def kinetic_energy_from_velocities(masses, velocities):
"""Kinetic energy for given velocities and masses.
Parameters
----------
masses : 1d array, shape (number of atoms, )
Atomic masses in amu.
velocities : 2d array, (number of atoms, 3)
Atomic velocities in Bohr/fs.
Returns
-------
E_kin : float
Kinetic energy in Hartree.
"""
return np.sum(masses[:, None] / 2 * velocities**2) * VELO2E
[docs]
def kinetic_energy_for_temperature(atom_number, T, fixed_dof=0):
"""Kinetic energy for given temperature and number of atoms.
Each atom has three degrees of freedom (1/2 * 3 == 3/2).
Parameters
----------
atom_number : int
Number of atoms. Each atom has three degrees of freedom.
T : float
Temperature in Kelvin.
fixed_dof : int, optional, default=0
Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
velocity is removed.
Returns
-------
E_kin : float
Kinetic energy in Hartree.
"""
return (3 * atom_number - fixed_dof) / 2 * T * KBAU
[docs]
def temperature_for_kinetic_energy(atom_number, E_kin, fixed_dof=0):
"""Temperature for given kinetic energy and atom number.
Each atom has three degrees of freedom (1/2 * 3 == 3/2).
Parameters
----------
atom_number : int
Number of atoms. Each atom has three degrees of freedom.
E_kin : float
Kinetic energy in Hartree.
fixed_dof : int, optional, default=0
Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
velocity is removed.
Returns
-------
Temperature : float
Temperature in Kelvin.
"""
return 2 * E_kin / ((3 * atom_number - fixed_dof) * KBAU)
[docs]
def remove_com_velocity(v, masses, keep_norm=True):
"""Remove center-of-mass velocity.
Returned units vary with the given input units.
Parameters
----------
v : np.array, 2d, shape (number of atoms, 3)
Velocities.
masses : np.array, 1d, shape (number of atoms, )
Atomic masses.
keep_norm : bool, default=True
Whether to rescale v to its original norm, after removal of v_com.
Returns
-------
v : np.array
Velocities without center-of-mass velocity.
"""
org_norm = np.linalg.norm(v)
v_com = (v * masses[:, None] / np.sum(masses)).sum(axis=0)
v -= v_com
if keep_norm:
v /= np.linalg.norm(v) / org_norm
return v
[docs]
def scale_velocities_to_temperatue(masses, v, T_desired, fixed_dof=0):
"""Scale velocities to a given temperature.
Parameters
----------
masses : np.array, 1d, shape (number of atoms, )
Atomic masses in amu.
v : np.array, 2d, shape (number of atoms, 3)
(Unscaled) velocities in Bohr/fs.
T_desired : float
Desired temperature in Kelvin.
fixed_dof : int, optional, default=0
Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
velocity is removed.
Returns
-------
v : np.array, 2d, shape (number of atoms, 3)
Scaled velocities in Bohr/fs.
"""
E_ref = kinetic_energy_for_temperature(
len(masses), T_desired, fixed_dof=fixed_dof
) # in Hartree
E_cur = kinetic_energy_from_velocities(masses, v) # in Hartree
scale = (E_ref / E_cur) ** 0.5
v *= scale
# E_now = kinetic_energy_from_velocities(masses, v)
# np.testing.assert_allclose(E_now, E_ref)
return v
[docs]
def scale_velocities_to_energy(masses, v, E_desired):
"""Scale velocities to a given temperature.
Parameters
----------
masses : np.array, 1d, shape (number of atoms, )
Atomic masses in amu.
v : np.array, 2d, shape (number of atoms, 3)
(Unscaled) velocities in Bohr/fs.
E_desired : float
Desired kinetic energy in Hartree.
Returns
-------
v : np.array, 2d, shape (number of atoms, 3)
Scaled velocities in Bohr/fs.
"""
E_cur = kinetic_energy_from_velocities(masses, v) # in Hartree
scale = (E_desired / E_cur) ** 0.5
v *= scale
return v
[docs]
def unscaled_velocity_distribution(masses, T, seed=None):
"""
ρ ∝ exp(- v² * m / (2 * kT))
ln(ρ) ∝ - 1/2 * v² * m / kT
v² ∝ -2 * ln(ρ) * kT / m
v ∝ ±(-2 * ln(ρ) * kT / m)**0.5
v ∝ ±(-2 * k)**0.5 * (ln(ρ) * T / m)**0.5
The first term of the RHS is constant. As we later scale the
velocities we neglect it. Don't use these velocities unscaled!
v ∝ ±(ln(ρ) * T / m)**0.5
"""
if seed is not None:
np.random.seed(seed)
ps = np.random.standard_normal((len(masses), 3))
v = ps * np.sqrt(T / masses[:, None])
return v
[docs]
def get_mb_velocities(
masses, cart_coords, T, remove_com_v=True, remove_rot_v=True, seed=None
):
"""Initial velocities from Maxwell-Boltzmann distribution.
Parameters
----------
masses : np.array, 1d, shape (number of atoms, )
Atomic masses in amu.
cart_coords : iterable, 1d, shape (3 * number of atoms, )
Atomic cartesian coordinates. Needed for removal of rotation.
T : float
Temperature in Kelvin.
remove_com_v : bool, default=True, optional
Whether to remove center-of-mass velocity.
remove_rot_v : bool, default=True, optional
Whether to remove rotational velocity.
seed : int, default=None, optional
Seed for the random-number-generator.
Returns
-------
v : np.array, 2d, shape (number of atoms, 3)
Initial velocities in Bohr/fs.
"""
masses = np.array(masses)
# Initial velocities
v = unscaled_velocity_distribution(masses, T, seed=seed)
if (len(masses) == 1) and remove_com_v:
raise Exception("Removing COM velocity with only 1 atom is a bad idea!")
fixed_dof = 0
# Remove center-of-mass velocity
if remove_com_v:
v = remove_com_velocity(v, masses)
fixed_dof += 3
if remove_rot_v:
# Right now this also removes the translational components
P = get_hessian_projector(cart_coords, masses, full=True)
v = P.dot(v.flatten()).reshape(-1, 3)
fixed_dof = 6
# In Bohr/fs
v = scale_velocities_to_temperatue(masses, v, T, fixed_dof=fixed_dof)
return v
[docs]
def get_mb_velocities_for_geom(
geom, T, remove_com_v=True, remove_rot_v=True, seed=None
):
"""Initial velocities from Maxwell-Boltzmann distribution.
See 'get_mb_velocities' for explanation.
"""
return get_mb_velocities(
geom.masses,
geom.cart_coords,
T,
remove_com_v=remove_com_v,
remove_rot_v=remove_rot_v,
seed=seed,
)
[docs]
def energy_forces_getter_closure(geom):
def energy_forces_getter(coords):
results = geom.get_energy_and_forces_at(coords)
energy = results["energy"]
forces = results["forces"]
return energy, forces
return energy_forces_getter