Useful Quantities

Describes different physical quantities.

Forces

jax_md.quantity.force(energy_fn)

Computes the force as the negative gradient of an energy.

Return type:

Callable[..., Array]

jax_md.quantity.canonicalize_force(energy_or_force_fn)

Return type:

Callable[..., Array]

Structural Features

jax_md.quantity.cosine_angles(dR)

Returns cosine of angles for all atom triplets.

Parameters:

dR (Array) – Matrix of displacements; ndarray(shape=[num_atoms, num_neighbors, spatial_dim]).

Return type:

Array

Returns:

Tensor of cosine of angles; ndarray(shape=[num_atoms, num_neighbors, num_neighbors]).

jax_md.quantity.pair_correlation(displacement_or_metric, radii, sigma, species=None, eps=1e-07, compute_average=False)

Computes the pair correlation function at a mesh of distances.

The pair correlation function measures the number of particles at a given distance from a central particle. The pair correlation function is defined by

\[\begin{split}g(r) = <\sum_{i \\neq j}\delta(r - |r_i - r_j|)>\end{split}\]

We make the approximation,

\[\begin{split}\delta(r) \\approx {1 \over \sqrt{2\pi\sigma^2}e^{-r / (2\sigma^2)}}\end{split}\]

Parameters:

• displacement_or_metric (Union[DisplacementFn, MetricFn]) – A function that computes the displacement or distance between two points.

• radii (Array) – An array of radii at which we would like to compute \(g(r)\).

• sigima – A float specifying the width of the approximating Gaussian.

• species (Array | None) – An optional array specifying the species of each particle. If species is None then we compute a single \(g(r)\) for all particles, otherwise we compute one \(g(r)\) for each species.

• eps (float) – A small additive constant used to ensure stability if the radius is zero.

Returns:

A function g_fn that computes the pair correlation function for a collection of particles.

\(g(r)\) is calculated separately for each particle. For species=None, the output of g_fn is an array of shape (N, nr), where N is the number of particles passed to g_fn and nr is the size of radii (the number of points at which we calculate \(g(r)\). When species is specified, the output is a list of nspecies arrays, each of shape (N, nr), where nspecies is the number of unique species. If gofr is the output of g_fn, then gofr[si][i] gives the \(g(r)\) for particle i considering only pair particles of species si.

Note: when species is specified, the returned list is in the order of the sorted unique species indices, not the order in which they appear.

jax_md.quantity.pair_correlation_neighbor_list(displacement_or_metric, box_size, radii, sigma, species=None, dr_threshold=0.5, eps=1e-07, fractional_coordinates=False, format=NeighborListFormat.Dense, compute_average=False)

Computes the pair correlation function at a mesh of distances.

The pair correlation function measures the number of particles at a given distance from a central particle. The pair correlation function is defined by

\[\begin{split}g(r) = <\sum_{i \\neq j}\delta(r - |r_i - r_j|)>\end{split}\]

We make the approximation,

\[\begin{split}\delta(r) \\approx {1 \over \sqrt{2\pi\sigma^2}e^{-r / (2\sigma^2)}}\end{split}\]

This function uses neighbor lists to speed up the calculation.

Parameters:

• displacement_or_metric (Union[DisplacementFn, MetricFn]) – A function that computes the displacement or distance between two points.

• box_size (Union[Array, ndarray, bool, number, bool, int, float, complex]) – The size of the box containing the particles.

• radii (Array) – An array of radii at which we would like to compute \(g(r)\).

• sigima – A float specifying the width of the approximating Gaussian.

• species (Array | None) – An optional array specifying the species of each particle. If species is None then we compute a single \(g(r)\) for all particles, otherwise we compute one \(g(r)\) for each species.

• dr_threshold (float) – A float specifying the halo size of the neighbor list.

• eps (float) – A small additive constant used to ensure stability if the radius is zero.

• fractional_coordinates (bool) – Bool determining whether positions are stored in the unit cube or not.

• format (NeighborListFormat) – The format of the neighbor lists. Must be Dense or Sparse.

Returns:

neighbor_fn that constructs a neighbor list (see neighbor_list in partition.py for details). g_fn that computes the pair correlation function for a collection of particles given their position and a neighbor list.

Return type:

A pair of functions

\(g(r)\) is calculated separately for each particle. For species=None, the output of g_fn is an array of shape (N, nr), where N is the number of particles passed to g_fn and nr is the size of radii (the number of points at which we calculate \(g(r)\). When species is specified, the output is a list of nspecies arrays, each of shape (N, nr), where nspecies is the number of unique species. If gofr is the output of g_fn, then gofr[si][i] gives the \(g(r)\) for particle i considering only pair particles of species si.

Note: when species is specified, the returned list is in the order of the sorted unique species indices, not the order in which they appear.

jax_md.quantity.phop(displacement, window_size)

Computes the phop indicator of rearrangements.

phop is an indicator function that is effective at detecting when particles in a quiescent system have experienced a rearrangement. Qualitatively, phop measures when the average position of a particle has changed significantly.

Formally, given a window of size \(\Delta t\) we two averages before and after the current time,

\[E_A[f] = E_{t\in[t - \Delta t / 2, t]}[f(t)] E_B[f] = E_{t\in[t, t + \Delta t / 2]}[f(t)].\]

In terms of these expectations, phop is given by,

\[phop = \sqrt{E_A[(R_i(t) - E_B[R_i(t)])^2]E_B[(R_i(t) - E_A[R_i(t)])^2]}\]

phop was first introduced in Candelier et al. [1]

Parameters:

• displacement (DisplacementFn) – A function that computes displacements between pairs of particles. See spaces.py for details.

• window_size (int) – An integer specifying the number of positions that constitute the window.

Return type:

Tuple[Callable[[Array], PHopState], Callable[[PHopState, Array], PHopState]]

Returns:

A pair of functions, (init_fn, update_fn) that initialize the state of a phop measurement and update the state of a phop measurement to include new positions.

References

[1]

R. Candelier et al. “Spatiotemporal Hierarchy of Relaxation Events, Dynamical Heterogeneities, and Structural Reorganization in a Supercooled Liquid” Physical Review Letters 105, 135702 (2010).

Physical Properties

jax_md.quantity.kinetic_energy(*unused_args, momentum=None, velocity=None, mass=1.0)

Computes the kinetic energy of a system.

To avoid ambiguity, either momentum or velocity must be passed explicitly as a keyword argument.

Parameters:

• momentum (Array | None) – Array specifying the momentum of the system.

• velocity (Array | None) – Array specifying the velocity of the system.

• mass (Array | float) – Array specifying the mass of the constituents.

Return type:

Array | float

Returns:

The kinetic energy of the system.

jax_md.quantity.temperature(*unused_args, momentum=None, velocity=None, mass=1.0)

Computes the temperature of a system.

To avoid ambiguity, either momentum or velocity must be passed explicitly as a keyword argument.

Parameters:

• momentum (Array | None) – Array specifying the momentum of the system.

• velocity (Array | None) – Array specifying the velocity of the system.

• mass (Array | float) – Array specifying the mass of the constituents.

Return type:

Array | float

Returns:

The temperature of the system in units of the Boltzmann constant.

jax_md.quantity.pressure(energy_fn, position, box, kinetic_energy=0.0, **kwargs)

Computes the internal pressure of a system.

Parameters:

• energy_fn (Callable[..., Array]) – A function that computes the energy of the system. This function must take as an argument perturbation which perturbs the box shape. Any energy function constructed using smap or in energy.py with a standard space will satisfy this property.

• position (Array) – An array of particle positions.

• box (Union[Array, ndarray, bool, number, bool, int, float, complex]) – A box specifying the shape of the simulation volume. Used to infer the volume of the unit cell.

• kinetic_energy (float | Array) – A float specifying the kinetic energy of the system.

Return type:

float

Returns:

A float specifying the pressure of the system.

jax_md.quantity.stress(energy_fn, position, box, mass=1.0, velocity=None, **kwargs)

Computes the internal stress of a system.

Parameters:

• energy_fn (Callable[..., Array]) – A function that computes the energy of the system. This function must take as an argument perturbation which perturbs the box shape. Any energy function constructed using smap or in energy.py with a standard space will satisfy this property.

• position (Array) – An array of particle positions.

• box (Union[Array, ndarray, bool, number, bool, int, float, complex]) – A box specifying the shape of the simulation volume. Used to infer the volume of the unit cell.

• mass (Array | float) – The mass of the particles; only used to compute the kinetic contribution if velocity is not None.

• velocity (Array | None) – An array of atomic velocities.

Return type:

Array

Returns:

A float specifying the pressure of the system.

Data Types

class jax_md.quantity.PHopState(position_buffer, phop)