repository·main·Indexed 23 days ago
https://github.com/jax-md/jax-mdA differentiable, hardware-accelerated molecular dynamics library built on JAX. It supports CPU, GPU, and TPU acceleration and provides tools for machine learning applications. Key features include the UMA (Universal Models for Atoms) interface for energy and force computation, and the jax_md.mm_forcefields package for building molecular mechanics force fields, including OPLS-AA implementation and various Coulomb handlers like CutoffCoulomb, EwaldCoulomb, and PMECoulomb.
JAX MD is a functional and data-driven library for accelerated, differentiable molecular dynamics. It leverages JAX to allow simulations to run on CPU, GPU, or TPU via XLA compilation. Because it is built on JAX, the entire simulation pipeline is end-to-end differentiable, enabling gradient-based optimization of simulation parameters.
Data in JAX MD is typically stored in arrays or tuples of arrays, and functions transform this data from one state to another.
The jax_md.partition module provides tools for spatial partitioning, which is essential for efficient neighbor searches in molecular dynamics.
Key components include:
neighbor_list: Generates neighbor lists for particles.neighbor_list_mask: Provides a mask for neighbor lists.to_jraph: Converts neighbor lists to a Jraph graph structure.to_dense: Converts neighbor lists to a dense representation.NeighborList: A class for managing neighbor list data.NeighborListFormat: Defines the format of the neighbor list.NeighborListFns: A collection of functions related to neighbor lists.CellList: A class for managing cell-based spatial partitioning.jax_md.mm_forcefields.reaxff package provides implementations of the ReaxFF (Reactive Force Field) method. It is organized into several submodules that handle different aspects of the ReaxFF energy and force calculations.The jax_md.mm_forcefields package provides utilities for building molecular mechanics (MM) force fields. The architecture is organized into several functional modules:
base.py: Defines core data structures like Topology (connectivity, exclusions, and 1–4 scaling masks), BondedParameters/NonbondedParameters (parameter bundles in kcal/mol and Å), and NonbondedOptions (runtime knobs for LJ and Coulomb terms).neighbor.py: Provides utilities for neighbor list management, including create_neighbor_list (a wrapper for jax_md.partition.neighbor_list), mask generation (make_exclusion_mask, make_14_table), and stable numerical primitives.io/: Contains parsers for input files. io/charmm.py handles CHARMM RTF/PRM and PDB files, while oplsaa/io.py provides higher-level orchestration to convert CHARMM files into (positions, Topology, Parameters) tuples.nonbonded/: Contains Coulomb handlers that all follow the same interface: CoulombHandler.energy(positions, charges, box, exclusion_mask, pair_14_mask, nlist, scale_14). Supported handlers include CutoffCoulomb, EwaldCoulomb, and PMECoulomb.oplsaa/: A composable implementation of the OPLS-AA force field, including topology creation, parameter validation, and energy assembly.JAX MD provides primitives for constructing Graph Neural Networks (GNNs) that are compatible with jax-md states, neighbor lists, end-to-end JIT compilation, and easy batching. These primitives are adapted from the jraph and graph_nets libraries.
Graphs in this framework are described by a GraphsTuple containing:
Key components include:
GraphIndependent: A layer that applies neural networks separately to node states, edge states, and globals (useful for encoding/decoding).GraphNetwork: A layer that transforms nodes, edges, and globals using neural networks via sum-message-aggregation (following the Battaglia et al. architecture).The jax_md.custom_smap module provides structure-mapped functions designed for systems with small periodic boxes where the cutoff radius $r_{\text{cut}}$ is greater than half the box length ($L/2$).
Standard jax_md.smap functions rely on the minimum image convention, which can fail in these scenarios. Instead, jax_md.custom_smap works with jax_md.custom_partition.NeighborListMultiImage. This approach uses explicit lattice shifts stored directly in the neighbor list to correctly handle periodic images, ensuring accuracy even when the box is smaller than the interaction range.
When working with small periodic boxes where the cutoff radius exceeds half the box size ($r_{\text{cut}} > L/2$), the standard minimum image convention (MIC) is insufficient because particles may interact with multiple periodic images of their neighbors. The jax_md.custom_partition module provides specialized neighbor list construction for these cases.
Key functions for construction include:
neighbor_list_multi_image: Constructs the neighbor list.neighbor_list_multi_image_mask: Provides a mask for the neighbor list.The jax_md.space module provides higher-order functions that take a space and return new functions. These are used to ensure that calculations like distances and displacements correctly account for the chosen boundary conditions (e.g., the shortest path in a periodic box).
Key higher-order functions include:
metric(space): Returns a function that calculates the distance between points according to the space's geometry.map_product(space, function): Maps a function over pairs of particles, accounting for periodic images.map_bond(space, function): Maps a function over bonded pairs.map_neighbor(space, function): Maps a function over neighboring particles.canonicalize_displacement_or_metric(space): Returns a function to canonicalize displacements or metrics within the given space.The quantity module provides tools to derive physical measurements from system configurations, velocities, energies, and forces. These are essential for thermodynamic analysis, simulation control, and mechanical response characterization.
Commonly used quantities:
quantity.force() (transforms energy to force), quantity.kinetic_energy(), and quantity.temperature().quantity.pressure(), quantity.stress(), quantity.volume(), quantity.volume_fraction(), quantity.bulk_modulus(), and elastic-tensor helpers.quantity.pair_correlation() and quantity.pair_correlation_neighbor_list() for radial distribution functions.quantity.gamma_from_stokes_law_3d() for estimating hydrodynamic drag.The a2c module contains specialized workflows for generating and managing crystal and amorphous structures.
Crystallization Utilities (a2c.crystallizer_utils):
get_subcells_to_crystallize() and get_subcells_to_crystallize_parallel(): Generate subcells for crystallization.subcells_to_structures() and valid_subcell(): Convert and filter candidate structures.Amorphous Generation (a2c.make_amorphous_utils):
random_packed_structure(): Generates random packed amorphous structures.Spaces define how distances and displacements are calculated between atoms, which is essential for handling boundary conditions like periodic boxes. A space is defined by a pair of functions: (displacement_fn, shift_fn).
displacement_fn(R_1, R_2): Computes the displacement vector between two points.shift_fn(R, dR): Displaces a point R by an amount dR.Supported spaces:
space.free(): Free boundary conditions.space.periodic(box_size): Periodic boundary conditions with a side length of box_size.space.periodic_general(box): Periodic parallelepiped defined by an affine transformation box applied to a unit cube.Mapping helpers to vectorize operations:
space.map_product(): Vectorizes a displacement or metric over all pairs.space.map_bond(): Maps over indexed bonds.space.map_neighbor(): Maps over neighbor-list entries.space.metric(): Converts a displacement function into a metric function.from jax_md import space
box_size = 25.0
displacement_fn, shift_fn = space.periodic(box_size)The ReaxFF implementation is split into the following functional modules:
jax_md.mm_forcefields.reaxff.reaxff_energy: Handles the calculation of ReaxFF energy components.jax_md.mm_forcefields.reaxff.reaxff_forcefield: Provides the main forcefield interface for ReaxFF.jax_md.mm_forcefields.reaxff.reaxff_helper: Contains helper functions and utilities for ReaxFF operations.jax_md.mm_forcefields.reaxff.reaxff_interactions: Manages the specific interaction terms used in the ReaxFF model.