Potential Energy Functions

Definitions of various standard energy functions.

Bond Potentials

jax_md.energy.simple_spring(dr, length=1, epsilon=1, alpha=2, **unused_kwargs)

Isotropic spring potential with a given rest length.

We define simple_spring to be a generalized Hookean spring with agreement when alpha = 2.

Return type:

Array

jax_md.energy.simple_spring_bond(displacement_or_metric, bond, bond_type=None, length=1, epsilon=1, alpha=2)

Convenience wrapper to compute energy of particles bonded by springs.

Return type:

Callable[[Array], Array]

Classical Potentials

jax_md.energy.soft_sphere(dr, sigma=1, epsilon=1, alpha=2, **unused_kwargs)

Finite ranged repulsive interaction between soft spheres.

This implements the overlap potential commonly used in jamming and soft matter studies:

\[U(r) = \frac{\epsilon}{\alpha} \left(1 - \frac{r}{\sigma}\right)^\alpha \quad \text{for } r < \sigma\]

and \(U(r) = 0\) for \(r \geq \sigma\).

Note: This is distinct from the inverse power law potential \(U(r) = \epsilon (\sigma/r)^\alpha\) used in some fluid simulations.

Reference:

O’Hern, C. S., Silbert, L. E., Liu, A. J., & Nagel, S. R. (2003). Jamming at zero temperature and zero applied stress: The epitome of disorder. Physical Review E, 68(1), 011306. https://doi.org/10.1103/PhysRevE.68.011306

Parameters:

• dr (Array) – An ndarray of shape [n, m] of pairwise distances between particles.

• sigma (Array | float) – Particle diameter. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• epsilon (Array | float) – Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• alpha (Array | float) – Exponent specifying interaction stiffness. Should either be a float point scalar or an ndarray whose shape is [n, m].

• unused_kwargs – Allows extra data (e.g. time) to be passed to the energy.

Return type:

Array

Returns:

Matrix of energies whose shape is [n, m].

jax_md.energy.soft_sphere_pair(displacement_or_metric, species=None, sigma=1.0, epsilon=1.0, alpha=2.0, per_particle=False)

Convenience wrapper to compute soft sphere energy over a system.

Return type:

Callable[[Array], Array]

jax_md.energy.soft_sphere_neighbor_list(displacement_or_metric, box_size, species=None, sigma=1.0, epsilon=1.0, alpha=2.0, dr_threshold=0.2, per_particle=False, fractional_coordinates=False, format=NeighborListFormat.OrderedSparse, neighbor_list_fn=<function neighbor_list>, pair_neighbor_list_fn=<function pair_neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute soft spheres using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.lennard_jones(dr, sigma=1, epsilon=1, **unused_kwargs)

Lennard-Jones interaction between particles with a minimum at sigma.

Parameters:

• dr (Array) – An ndarray of shape [n, m] of pairwise distances between particles.

• sigma (Array | float) – Distance between particles where the energy has a minimum. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• epsilon (Array | float) – Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• unused_kwargs – Allows extra data (e.g. time) to be passed to the energy.

Return type:

Array

Returns:

Matrix of energies of shape [n, m].

jax_md.energy.lennard_jones_pair(displacement_or_metric, species=None, sigma=1.0, epsilon=1.0, r_onset=2.0, r_cutoff=2.5, per_particle=False)

Convenience wrapper to compute Lennard-Jones energy over a system.

Return type:

Callable[[Array], Array]

jax_md.energy.lennard_jones_neighbor_list(displacement_or_metric, box_size, species=None, sigma=1.0, epsilon=1.0, r_onset=2.0, r_cutoff=2.5, dr_threshold=0.5, per_particle=False, fractional_coordinates=False, format=NeighborListFormat.OrderedSparse, neighbor_list_fn=<function neighbor_list>, pair_neighbor_list_fn=<function pair_neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute Lennard-Jones using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.morse(dr, sigma=1.0, epsilon=5.0, alpha=5.0, **unused_kwargs)

Morse interaction between particles with a minimum at sigma.

Parameters:

• dr (Array) – An ndarray of shape [n, m] of pairwise distances between particles.

• sigma (Array | float) – Distance between particles where the energy has a minimum. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• epsilon (Array | float) – Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• alpha (Array | float) – Range parameter. Should either be a floating point scalar or an ndarray whose shape is [n, m].

• unused_kwargs – Allows extra data (e.g. time) to be passed to the energy.

Return type:

Array

Returns:

Matrix of energies of shape [n, m].

jax_md.energy.morse_pair(displacement_or_metric, species=None, sigma=1.0, epsilon=5.0, alpha=5.0, r_onset=2.0, r_cutoff=2.5, per_particle=False)

Convenience wrapper to compute Morse energy over a system.

Return type:

Callable[[Array], Array]

jax_md.energy.morse_neighbor_list(displacement_or_metric, box_size, species=None, sigma=1.0, epsilon=5.0, alpha=5.0, r_onset=2.0, r_cutoff=2.5, dr_threshold=0.5, per_particle=False, fractional_coordinates=False, format=NeighborListFormat.OrderedSparse, neighbor_list_fn=<function neighbor_list>, pair_neighbor_list_fn=<function pair_neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute Morse using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.gupta_potential(displacement, p, q, r_0n, U_n, A, cutoff)

Gupta potential with default parameters for Au_55 cluster. Gupta potential was introduced by R. P. Gupta [1]. This potential uses parameters that were fit for bulk gold by Jellinek [2]. This particular implementation of the Gupta potential was introduced by Garzon and Posada-Amarillas [3].

Parameters:

• displacement – Function to compute displacement between two positions.

• p – Gupta potential parameter of the repulsive term that was fitted for bulk gold.

• q – Gupta potential parameter of the attractive term that was fitted for bulk gold.

• r_0n – Parameter that determines the length scale of the potential. This value was particularly fit for gold clusters of size 55 atoms.

• U_n – Parameter that determines the energy scale, fit particularly for gold clusters of size 55 atoms.

• A – Parameter that was obtained using the cohesive energy of the fcc gold metal.

• cutoff – Pairwise interactions that are farther than the cutoff distance will be ignored.

Returns:

A function that takes in positions of gold atoms (shape [n, 3] where n is the number of atoms) and returns the total energy of the system in units of eV.

References

[1]

R.P. Gupta, Phys. Rev. B 23, 6265 (1981)

[2]

J. Jellinek, in Metal-Ligand Interactions, edited by N. Russo and D. R. Salahub (Kluwer Academic, Dordrecht, 1996), p. 325.

[3]

I.L. Garzon, A. Posada-Amarillas, Phys. Rev. B 54, 16 (1996)

jax_md.energy.gupta_gold55(displacement, cutoff=8.0)

jax_md.energy.bks(dr, Q_sq, exp_coeff, exp_decay, attractive_coeff, repulsive_coeff, coulomb_alpha, cutoff, **unused_kwargs)

Beest-Kramer-van Santen (BKS) potential [4] which is commonly used to model silicas. This function computes the interaction between two given atoms within the Buckingham form [5] , following the implementation from Liu et al. [6] .

Parameters:

• dr (Array) – An ndarray of shape [n, m] of pairwise distances between particles.

• Q_sq (Array) – An ndarray of shape [n, m] of pairwise product of partial charges.

• exp_coeff (Array) – An ndarray of shape [n, m] that sets the scale of the exponential decay of the short-range interaction.

• attractive_coeff (Array) – An ndarray of shape [n, m] for the coefficient of the attractive 6th order term.

• repulsive_coeff (Array) – An ndarray of shape [n, m] for the coefficient of the repulsive 24th order term, to prevent the unphysical fusion of atoms.

• coulomb_alpha (Array) – Damping parameter for the approximation of the long-range coulombic interactions (a scalar).

• cutoff (float) – Cutoff distance for considering pairwise interactions.

• unused_kwargs – Allows extra data (e.g. time) to be passed to the energy.

Return type:

Array

Returns:

Matrix of energies of shape [n, m].

References

[4]

Van Beest, B. W. H., Gert Jan Kramer, and R. A. Van Santen. “Force fields for silicas and aluminophosphates based on ab initio calculations.” Physical Review Letters 64.16 (1990): 1955.

[5]

Carré, Antoine, et al. “Developing empirical potentials from ab initio simulations: The case of amorphous silica.” Computational Materials Science 124 (2016): 323-334.

[6]

Liu, Han, et al. “Machine learning Forcefield for silicate glasses.” arXiv preprint arXiv:1902.03486 (2019).

jax_md.energy.bks_pair(displacement_or_metric, species, Q_sq, exp_coeff, exp_decay, attractive_coeff, repulsive_coeff, coulomb_alpha, cutoff)

Convenience wrapper to compute BKS energy over a system.

Return type:

Callable[[Array], Array]

jax_md.energy.bks_neighbor_list(displacement_or_metric, box_size, species, Q_sq, exp_coeff, exp_decay, attractive_coeff, repulsive_coeff, coulomb_alpha, cutoff, dr_threshold=0.8, fractional_coordinates=False, format=NeighborListFormat.OrderedSparse, neighbor_list_fn=<function neighbor_list>, pair_neighbor_list_fn=<function pair_neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute BKS energy using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.bks_silica_pair(displacement_or_metric, species, cutoff=8.0)

Convenience wrapper to compute BKS energy for SiO2.

jax_md.energy.bks_silica_neighbor_list(displacement_or_metric, box_size, species, cutoff=8.0, dr_threshold=1.0, fractional_coordinates=False, format=NeighborListFormat.OrderedSparse, neighbor_list_fn=<function neighbor_list>, pair_neighbor_list_fn=<function pair_neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute BKS energy using neighbor lists.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.stillinger_weber(displacement, sigma=2.0951, A=7.049556277, B=0.6022245584, lam=21.0, gamma=1.2, epsilon=2.16826, three_body_strength=1.0, cutoff=3.77118)

Computes the Stillinger-Weber potential.

The Stillinger-Weber (SW) potential [7] which is commonly used to model silicon and similar systems. This function uses the default SW parameters from the original paper. The SW potential was originally proposed to model diamond in the diamond crystal phase and the liquid phase, and is known to give unphysical amorphous configurations [8] [9] . For this reason, we provide a three_body_strength parameter. Changing this number to 1.5 or 2.0 has been know to produce more physical amorphous phase, preventing most atoms from having more than four nearest neighbors. Note that this function currently assumes nearest-image-convention.

Parameters:

• displacement (DisplacementFn) – The displacement function for the space.

• sigma (float) – A scalar that sets the distance scale between neighbors.

• A (float) – A scalar that determines the scale of two-body term.

• B (float) – A scalar that determines the scale of the \(1 / r^p\) term.

• lam (float) – A scalar that determines the scale of the three-body term.

• epsilon (float) – A scalar that sets the total energy scale.

• gamma (float) – A scalar used to fit the angle interaction.

• three_body_strength (float) – A scalar that determines the relative strength of the angular interaction. Default value is 1.0, which works well for the diamond crystal and liquid phases. 1.5 and 2.0 have been used to model amorphous silicon.

Return type:

Callable[[Array], Array]

Returns:

A function that computes the total energy.

References

[7]

Stillinger, Frank H., and Thomas A. Weber. “Computer simulation of local order in condensed phases of silicon.” Physical review B 31.8 (1985): 5262.

[8]

Holender, J. M., and G. J. Morgan. “Generation of a large structure (105 atoms) of amorphous Si using molecular dynamics.” Journal of Physics: Condensed Matter 3.38 (1991): 7241.

[9]

Barkema, G. T., and Normand Mousseau. “Event-based relaxation of continuous disordered systems.” Physical review letters 77.21 (1996): 4358.

jax_md.energy.stillinger_weber_neighbor_list(displacement, box_size, sigma=2.0951, A=7.049556277, B=0.6022245584, lam=21.0, gamma=1.2, epsilon=2.16826, three_body_strength=1.0, cutoff=3.77118, dr_threshold=0.5, fractional_coordinates=False, format=NeighborListFormat.Dense, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute Stillinger-Weber using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.load_lammps_tersoff_parameters(file)

Reads Tersoff parameters from a LAMMPS file and returns parameter tables.

This function reads multi-element original Tersoff potential parameters from a file.

Parameters:

file (TextIO) – A parameter file that is written with lammps format.

Returns:

An array that contains Tersoff parameters.

Return type:

list[Dict[str, Any]]

jax_md.energy.tersoff(displacement, params, species=None)

Computes the Tersoff potential.

The Tersoff potential [1] which is commonly used to model semiconducting materials. The Tersoff potential was originally proposed to model various types of lattice with a simple functional form. For this reason, Tersoff model was introduced bond-order function to determine the strength of repulsive and attractive forces between atoms.

Parameters:

• displacement (DisplacementFn) – The displacement function for the space.

• params (list[Dict[str, Any]]) – A dictionary of parameters for the tersoff potential. Usually this should be loaded from lammps using the load_lammps_tersoff_parameters function.

• species (Array | None) – An array of species. Currently only None is supported.

Return type:

Callable[[Array], Array]

Returns:

A function that computes the total energy.

[1] J. Tersoff “New empirical approach for the structure and energy of covalent systems” Physical review B 37.12 (1988): 6991.

jax_md.energy.tersoff_neighbor_list(displacement, box_size, params, species=None, dr_threshold=0.5, disable_cell_list=False, fractional_coordinates=True, format=NeighborListFormat.Dense, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Computes the Tersoff potential.

The Tersoff potential [1] which is commonly used to model semiconducting materials. The Tersoff potential was originally proposed to model various types of lattice with a simple functional form. For this reason, Tersoff model was introduced bond-order function to determine the strength of repulsive and attractive forces between atoms.

Parameters:

• displacement (DisplacementFn) – The displacement function for the space.

• box_size (float) – A float or vector specifying the size of the simulation box.

• params (list[Dict[str, Any]]) – A dictionary of parameters for the tersoff potential. Usually this should be loaded from lammps using the load_lammps_tersoff_parameters function.

• species (Array | None) – An array of species. Currently only None is supported.

• dr_threshold (float) – A distance threshold that controls how often the neighor list is recomputed.

• fractional_coordinates (bool) – A boolean specifying whether coordinates are stored in the unit cube.

• format (NeighborListFormat) – Format of the neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

Returns:

A pair of functions. One that builds the neighbor list and one that computes the total energy.

[1] J. Tersoff “New empirical approach for the structure and energy of covalent systems” Physical review B 37.12 (1988): 6991.

jax_md.energy.tersoff_from_lammps_parameters_neighbor_list(displacement, box_size, f, dr_threshold=0.5, fractional_coordinates=True, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute Tersoff energy with LAMMPS parameters.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.load_lammps_eam_parameters(file)

Reads EAM parameters from a LAMMPS file and returns relevant spline fits.

This function reads single-element EAM potential fit parameters from a file in DYNAMO funcl format. In summary, the file contains:

• Line 1-3: Comments

• Line 4: Number of elements and the element type

• Line 5: The number of charge values that the embedding energy is evaluated on (num_drho), interval between the charge values (drho), the number of distances the pairwise energy and the charge density is evaluated on (num_dr), the interval between these distances (dr), and the cutoff distance (cutoff).

The lines that come after are the embedding function evaluated on num_drho charge values, charge function evaluated at num_dr distance values, and pairwise energy evaluated at num_dr distance values. Note that the pairwise energy is multiplied by distance (in units of eV x Angstroms).

For more details of the DYNAMO file format, see: https://sites.google.com/a/ncsu.edu/cjobrien/tutorials-and-guides/eam

Parameters:

f – File handle for the EAM parameters text file.

Return type:

Tuple[Callable[[Array], Array], Callable[[Array], Array], Callable[[Array], Array], float]

Returns:

A tuple containing three functions and a cutoff distance.

charge_fn:

A function that takes an ndarray of shape [n, m] of distances between particles and returns a matrix of charge contributions.

embedding_fn:

Function that takes an ndarray of shape [n] of charges and returns an ndarray of shape [n] of the energy cost of embedding an atom into the charge.

pairwise_fn:

A function that takes an ndarray of shape [n, m] of distances and returns an ndarray of shape [n, m] of pairwise energies.

cutoff:

Cutoff distance for the embedding_fn and pairwise_fn.

jax_md.energy.eam(displacement_or_metric, charge_fn, embedding_fn, pairwise_fn, axis=None)

Interatomic potential as approximated by embedded atom model (EAM).

This code implements the EAM approximation to interactions between metallic atoms. In EAM, the potential energy of an atom is given by two terms: a pairwise energy and an embedding energy due to the interaction between the atom and background charge density. The EAM potential for a single atomic species is often determined by three functions:

1. Charge density contribution of an atom as a function of distance.

2. Energy of embedding an atom in the background charge density.

3. Pairwise energy.

These three functions are usually provided as spline fits, and we follow the implementation and spline fits given by Mishin et al. [10] Note that in current implementation, the three functions listed above can also be expressed by a any function with the correct signature, including neural networks.

Parameters:

• displacement – A function that produces an ndarray of shape [n, m, spatial_dimension] of particle displacements from particle positions specified as an ndarray of shape [n, spatial_dimension] and [m, spatial_dimension] respectively.

• box_size – The size of the simulation box.

• charge_fn (Callable[[Array], Array]) – A function that takes an ndarray of shape [n, m] of distances between particles and returns a matrix of charge contributions.

• embedding_fn (Callable[[Array], Array]) – Function that takes an ndarray of shape [n] of charges and returns an ndarray of shape [n] of the energy cost of embedding an atom into the charge.

• pairwise_fn (Callable[[Array], Array]) – A function that takes an ndarray of shape [n, m] of distances and returns an ndarray of shape [n, m] of pairwise energies.

• cutoff – A float specifying the maximum interaction distance.

• dr_threshold – A float specifying the halo in the neighbor list.

• axis (Optional[Tuple[int, ...]]) – Specifies which axis the total energy should be summed over.

• fractional_coordinates – A boolean specifying whether or not the coordinates will be in the unit cube.

• format – The format of the neighbor list.

Return type:

Callable[[Array], Array]

Returns:

A tuple containing a function to build the neighbor list and function that computes the EAM energy of a set of atoms with positions given by an [n, spatial_dimension] ndarray.

References

[10]

Y. Mishin, D. Farkas, M.J. Mehl, DA Papaconstantopoulos, “Interatomic potentials for monoatomic metals from experimental data and ab initio calculations.” Physical Review B, 59 (1999)

jax_md.energy.eam_from_lammps_parameters(displacement, f)

Convenience wrapper to compute EAM energy with LAMMPS parameters.

Return type:

Callable[[Array], Array]

jax_md.energy.eam_neighbor_list(displacement_or_metric, box_size, charge_fn, embedding_fn, pairwise_fn, cutoff, dr_threshold=0.5, axis=None, fractional_coordinates=True, format=NeighborListFormat.Sparse, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute EAM using a neighbor list.

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

jax_md.energy.eam_from_lammps_parameters_neighbor_list(displacement, box_size, f, axis=None, dr_threshold=0.5, fractional_coordinates=True, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Convenience wrapper to compute EAM energy with parameters from LAMMPS using a neighbor list..

Return type:

Tuple[NeighborListFns, Callable[..., Array]]

Neural Network Potentials

jax_md.energy.behler_parrinello(displacement, *, key, species=None, mlp_sizes=(30, 30), sym_kwargs=None, per_particle=False, spatial_dimension=3, activation=<PjitFunction of <function tanh>>)

Build a Behler-Parrinello symmetry-function energy model.

Parameters:

• displacement (DisplacementFn) – Function to compute displacement between two positions.

• key (Array) – A JAX random key used to initialize model parameters.

• species (Array | None) – None or an integer array specifying the species of each particle.

• mlp_sizes (Tuple[int, ...]) – Layer widths for the per-particle MLP.

• sym_kwargs (Optional[Dict[str, Any]]) – Additional kwargs for symmetry function construction.

• per_particle (bool) – If True, return per-particle energies instead of the total.

• spatial_dimension (int) – The spatial dimension of the system (default 3).

• activation (Callable) – Activation function for the MLP (default jnp.tanh).

Returns:

An energy_fn(R, **kwargs) -> scalar. The underlying NNX module is accessible via energy_fn.model for training.

jax_md.energy.behler_parrinello_neighbor_list(displacement, box_size, *, key, species=None, mlp_sizes=(30, 30), sym_kwargs=None, dr_threshold=0.5, fractional_coordinates=False, format=NeighborListFormat.Sparse, neighbor_list_fn=<function neighbor_list>, activation=<PjitFunction of <function tanh>>, **neighbor_kwargs)

Build a Behler-Parrinello energy model with neighbor lists.

Parameters:

• displacement (DisplacementFn) – Function to compute displacement between two positions.

• box_size (float) – The size of the simulation volume.

• key (Array) – A JAX random key used to initialize model parameters.

• species (Array | None) – None or an integer array specifying the species of each particle.

• mlp_sizes (Tuple[int, ...]) – Layer widths for the per-particle MLP.

• sym_kwargs (Optional[Dict[str, Any]]) – Additional kwargs for symmetry function construction.

• dr_threshold (float) – Halo radius for neighbor list construction.

• fractional_coordinates (bool) – Whether coordinates are fractional.

• format (NeighborListFormat) – Neighbor list format (Dense or Sparse).

• activation (Callable) – Activation function for the MLP (default jnp.tanh).

Returns:

A tuple (neighbor_fn, energy_fn) where energy_fn(R, neighbor=nbrs) -> scalar. The underlying NNX module is accessible via energy_fn.model for training.

jax_md.energy.graph_network(displacement_fn, r_cutoff, *, key, nodes=None, spatial_dimension=3, n_recurrences=2, mlp_sizes=(64, 64), **model_kwargs)

Convenience wrapper around EnergyGraphNet model.

Parameters:

• displacement_fn (DisplacementFn) – Function to compute displacement between two positions.

• r_cutoff (float) – A floating point cutoff; Edges will be added to the graph for pairs of particles whose separation is smaller than the cutoff.

• key (Array) – A JAX random key used to initialize model parameters.

• nodes (Array | None) – None or an ndarray of shape [N, node_dim] specifying the state of the nodes. If None this is set to the zeros vector. Often, for a system with multiple species, this could be the species id.

• spatial_dimension (int) – The spatial dimension of the system (default 3).

• n_recurrences (int) – The number of steps of message passing in the graph network.

• mlp_sizes (Tuple[int, ...]) – A tuple specifying the layer-widths for the fully-connected networks used to update the states in the graph network.

• **model_kwargs – Additional kwargs forwarded to EnergyGraphNet (e.g. activation, kernel_init, bias_init).

Return type:

Callable[..., Array]

Returns:

An energy_fn(R, **kwargs) -> scalar. The underlying NNX module is accessible via energy_fn.model for training.

jax_md.energy.graph_network_neighbor_list(displacement_fn, box_size, r_cutoff, dr_threshold, *, key, nodes=None, spatial_dimension=3, n_recurrences=2, mlp_sizes=(64, 64), fractional_coordinates=False, format=NeighborListFormat.Sparse, neighbor_list_fn=<function neighbor_list>, **neighbor_kwargs)

Convenience wrapper around EnergyGraphNet model using neighbor lists.

Parameters:

• displacement_fn (DisplacementFn) – Function to compute displacement between two positions.

• box_size (Union[Array, ndarray, bool, number, bool, int, float, complex]) – The size of the simulation volume, used to construct neighbor list.

• r_cutoff (float) – A floating point cutoff; Edges will be added to the graph for pairs of particles whose separation is smaller than the cutoff.

• dr_threshold (float) – A floating point number specifying a “halo” radius that we use for neighbor list construction. See neighbor_list for details.

• key (Array) – A JAX random key used to initialize model parameters.

• nodes (Array | None) – None or an ndarray of shape [N, node_dim] specifying the state of the nodes. If None this is set to the zeros vector. Often, for a system with multiple species, this could be the species id.

• spatial_dimension (int) – The spatial dimension of the system (default 3).

• n_recurrences (int) – The number of steps of message passing in the graph network.

• mlp_sizes (Tuple[int, ...]) – A tuple specifying the layer-widths for the fully-connected networks used to update the states in the graph network.

• fractional_coordinates (bool) – A boolean specifying whether or not the coordinates will be in the unit cube.

• format (NeighborListFormat) – The format of the neighbor list. See partition.NeighborListFormat for details. Only Dense and Sparse formats are accepted. If the Dense format is used, then the graph network is constructed using the JAX MD backend, otherwise Jraph is used.

• **neighbor_kwargs – Additional kwargs forwarded to the neighbor list constructor and to EnergyGraphNet (e.g. activation, kernel_init, bias_init).

Returns:

A tuple (neighbor_fn, energy_fn) where energy_fn(R, neighbor=nbrs) -> scalar. The underlying NNX module is accessible via energy_fn.model for training.

Helper Functions

jax_md.energy.multiplicative_isotropic_cutoff(fn, r_onset, r_cutoff)

Takes an isotropic function and constructs a truncated function.

Given a function f:R -> R, we construct a new function f':R -> R such that f'(r) = f(r) for r < r_onset, f'(r) = 0 for r > r_cutoff, and f(r) is \(C^1\) everywhere. To do this, we follow the approach outlined in HOOMD Blue [11] (thanks to Carl Goodrich for the pointer). We construct a function S(r) such that S(r) = 1 for r < r_onset, S(r) = 0 for r > r_cutoff, and S(r) is \(C^1\). Then f'(r) = S(r)f(r).

Parameters:

• fn (Callable[..., Array]) – A function that takes an ndarray of distances of shape [n, m] as well as varargs.

• r_onset (float) – A float specifying the distance marking the onset of deformation.

• r_cutoff (float) – A float specifying the cutoff distance.

Return type:

Callable[..., Array]

Returns:

A new function with the same signature as fn, with the properties outlined above.

References

[11]

HOOMD Blue documentation. Accessed on 05/31/2019. https://hoomd-blue.readthedocs.io/en/stable/module-md-pair.html#hoomd.md.pair.pair

jax_md.energy.load_lammps_eam_parameters(file)

Reads EAM parameters from a LAMMPS file and returns relevant spline fits.

This function reads single-element EAM potential fit parameters from a file in DYNAMO funcl format. In summary, the file contains:

• Line 1-3: Comments

• Line 4: Number of elements and the element type

• Line 5: The number of charge values that the embedding energy is evaluated on (num_drho), interval between the charge values (drho), the number of distances the pairwise energy and the charge density is evaluated on (num_dr), the interval between these distances (dr), and the cutoff distance (cutoff).

The lines that come after are the embedding function evaluated on num_drho charge values, charge function evaluated at num_dr distance values, and pairwise energy evaluated at num_dr distance values. Note that the pairwise energy is multiplied by distance (in units of eV x Angstroms).

For more details of the DYNAMO file format, see: https://sites.google.com/a/ncsu.edu/cjobrien/tutorials-and-guides/eam

Parameters:

f – File handle for the EAM parameters text file.

Return type:

Tuple[Callable[[Array], Array], Callable[[Array], Array], Callable[[Array], Array], float]

Returns:

A tuple containing three functions and a cutoff distance.

charge_fn:

A function that takes an ndarray of shape [n, m] of distances between particles and returns a matrix of charge contributions.

embedding_fn:

Function that takes an ndarray of shape [n] of charges and returns an ndarray of shape [n] of the energy cost of embedding an atom into the charge.

pairwise_fn:

A function that takes an ndarray of shape [n, m] of distances and returns an ndarray of shape [n, m] of pairwise energies.

cutoff:

Cutoff distance for the embedding_fn and pairwise_fn.