utils

Auxiliary routines for initial velocities

Functions

get_random_vel(→ jax.Array)	Generate random atomic velocities using a Maxwell-Boltzmann distribution
remove_com_moment(→ jax.Array)	Zero the linear center-of-mass momentum.
remove_angular_moment(→ tuple[jax.Array, jax.Array])	Zero the global angular momentum.
clean_momenta(→ tuple[jax.Array, jax.Array])	Remove any relevant external momenta
inertia_tensor(→ jax.Array)	Compute the inertia tensor for a given set of point particles
angular_moment(→ jax.Array)	Compute the angular moment of a set of point particles
angular_velocity(→ jax.Array)	Derive the angular velocity from the angular moment and the inertia tensor
rigid_body_angular_velocities(→ jax.Array)	Generate the velocities of a set of atoms that move as a rigid body.
get_ndof_internal_md(→ int)	Return the effective number of internal degrees of freedom for MD simulations
cell_symmetrize(→ tuple[jax.Array, jax.Array, ...)	Symmetrizes the unit cell tensor, and updates the position vectors
cell_lower(→ tuple[jax.Array, jax.Array])	Transform the cell tensor to its lower diagonal form. The transformation
get_random_vel_press(→ jax.Array)	Generates symmetric tensor of barostat velocities
get_ndof_baro(→ int)	Calculates the number of degrees of freedom associated with the cell fluctuation
stabilized_cholesky_decomp(→ jax.Array)	Do LDL^T and transform to MM^T with negative diagonal entries of D put equal to zero

Module Contents

utils.get_random_vel(temp0: float | jax.Array, scalevel0: bool, masses: jax.Array, key: jax.Array, select: jax.Array | None = None) → jax.Array

Generate random atomic velocities using a Maxwell-Boltzmann distribution

Arguments:

temp0

The temperature for the Maxwell-Boltzmann distribution.

scalevel0

When set to True, the velocities are rescaled such that the instantaneous temperature coincides with temp0.

masses

An (N,) array with atomic masses.

Optional arguments:

select

When given, this must be an array of integer indexes. Only for these atoms (masses) initial velocities will be generated.

Returns: An (N, 3) array with random velocities. When the select option is used, the shape of the results is (M, 3), where M is the length of the select array.

utils.remove_com_moment(vel: jax.Array, masses: jax.Array) → jax.Array

Zero the linear center-of-mass momentum.

Arguments:

vel

An (N, 3) array with atomic velocities. This array is modified in-place.

masses

An (N,) array with atomic masses

The zero linear COM momentum is achieved by subtracting translational rigid body motion from the atomic velocities.

utils.remove_angular_moment(pos: jax.Array, vel: jax.Array, masses: jax.Array) → tuple[jax.Array, jax.Array]

Zero the global angular momentum.

Arguments:

pos

An (N, 3) array with atomic positions. This array is not modified.

vel

An (N, 3) array with atomic velocities. This array is modified in-place.

masses

An (N,) array with atomic masses

The zero angular momentum is achieved by subtracting angular rigid body motion from the atomic velocities. (The angular momentum is measured with respect to the center of mass to avoid that this routine reintroduces a linear COM velocity. This is also beneficial for the numerical stability.)

utils.clean_momenta(pos: jax.Array, vel: jax.Array, masses: jax.Array, cell: IMLCV.new_yaff.system.YaffCell) → tuple[jax.Array, jax.Array]

Remove any relevant external momenta

Arguments:

pos

An (N, 3) array with atomic positions. This array is not modified.

vel

An (N, 3) array with atomic velocities. This array is modified in-place.

masses

An (N,) array with atomic masses

cell

A Cell instance describing the periodic boundary conditions.

utils.inertia_tensor(pos: jax.Array, masses: jax.Array) → jax.Array

Compute the inertia tensor for a given set of point particles

Arguments:

pos

An (N, 3) array with atomic positions.

masses

An (N,) array with atomic masses.

Returns: a (3, 3) array containing the inertia tensor.

utils.angular_moment(pos: jax.Array, vel: jax.Array, masses: jax.Array) → jax.Array

Compute the angular moment of a set of point particles

Arguments:

pos

An (N, 3) array with atomic positions.

vel

An (N, 3) array with atomic velocities.

masses

An (N,) array with atomic masses.

Returns: a (3,) array with the angular momentum vector.

utils.angular_velocity(ang_mom: jax.Array, iner_tens: jax.Array, epsilon: float = 1e-10) → jax.Array

Derive the angular velocity from the angular moment and the inertia tensor

Arguments:

ang_mom

An (3,) array with angular momenta.

iner_tens

A (3, 3) array with the inertia tensor.

Optional arguments:

epsilon

A threshold for the low eigenvalues of the inertia tensor. When an eigenvalue is below this threshold, it is assumed to be zero plus some (irrelevant) numerical noise.

Returns: An (3,) array with the angular velocity vector.

In principle these routine should merely return:

jnp.linalg.solve(inter_tens, ang_mom).

However, when the inertia tensor has zero eigenvalues (linear molecules, single atoms), this routine will use a proper pseudo inverse of the inertia tensor.

utils.rigid_body_angular_velocities(pos: jax.Array, ang_vel: jax.Array) → jax.Array

Generate the velocities of a set of atoms that move as a rigid body.

Arguments:

pos

An (N, 3) array with atomic positions.

ang_vel

An (3,) array with the angular velocity vector of the rigid body.

Returns: An (N, 3) array with atomic velocities in the rigid body. Note that the linear momentum is zero.

utils.get_ndof_internal_md(natom: int, nper: int) → int

Return the effective number of internal degrees of freedom for MD simulations

Arguments:

natom

The number of atoms

nper

The number of periodic boundary conditions (0 for isolated systems)

utils.cell_symmetrize(ff: IMLCV.new_yaff.ff.YaffFF, vector_list=None, tensor_list=None) → tuple[jax.Array, jax.Array, list[jax.Array], list[jax.Array]]

Symmetrizes the unit cell tensor, and updates the position vectors

Arguments:

ff

A ForceField instance

Optional arguments:

vector_list

A list of numpy vectors which should be transformed under the symmetrization. Note that the positions are already transformed automatically

tensor_list

A list of numpy tensors of rank 2 which should be transformed under the symmetrization.

utils.cell_lower(rvecs: jax.Array) → tuple[jax.Array, jax.Array]

Transform the cell tensor to its lower diagonal form. The transformation is described here https://lammps.sandia.gov/doc/Howto_triclinic.html, bearing in mind that YAFF stores cell vectors as rows, not columns.

Arguments:

rvecs

A [3x3] NumPy array representing a cell tensor

Returns:

newrvecs

A [3x3] NumPy array representing a lower-diagonal form of rvecs

rot

A [3x3] matrix representing the rotation matrix to go from rvecs to newrvecs

utils.get_random_vel_press(mass: jax.Array, temp: jax.Array, key: jax.Array) → jax.Array

Generates symmetric tensor of barostat velocities

Arguments:

mass

The Barostat mass

temp

The temperature at which the velocities are selected

utils.get_ndof_baro(dim: int, anisotropic: bool, vol_constraint: bool) → int

Calculates the number of degrees of freedom associated with the cell fluctuation

Arguments:

dim

The dimension of the system

anisotropic

Boolean value determining whether anisotropic cell fluctuations are allowed

vol_constraint

Boolean value determining whether the cell volume can change

utils.stabilized_cholesky_decomp(mat: jax.Array) → jax.Array

Do LDL^T and transform to MM^T with negative diagonal entries of D put equal to zero Assume mat is square and symmetric (but not necessarily positive definite).