JAX MD Reference

JAX MD is a collection of tools for working with collections of particles in JAX. Different parts of the JAX MD tooling are designed to work independently and so you should feel free to pick and choose which ones you use for your particular problem.

Examples

• Examples
  • Microcanonical Ensemble (NVE)
  • Canonical Ensemble (NVT) - Nose-Hoover
  • Isothermal-Isobaric Ensemble (NPT) - Nose-Hoover
  • Microcanonical Ensemble (NVE) with Neighbor Lists
  • Neural Network Potentials
  • Equivariant Neural Network Potentials
  • Harmonic Minimization
  • ReaxFF NVE Simulation
  • Multi-Image Neighbor Lists for Small Periodic Boxes
  • Symmetric Molecular Dynamics (SyMD)
  • OPLSAA Torsion Scan
  • UMA (Universal Models for Atoms) — JAX-MD
• Units Examples
  • Metal Units (NVE Simulation)
  • Metal Units (NVT Simulation)
  • Metal Units (NPT Simulation)
  • Metal Units (NVK Simulation)
  • Metal Units (Berendsen Thermostat)
  • Metal Units (Velocity Rescaling)
  • Metal Units (CSVR Thermostat)
  • Metal Units (NPT Simulation - Melting)

Reference

• Defining Spaces
  • Spaces
  • Higher Order Functions
  • Helper Functions
• Spatial Partitioning
  • Neighbor Lists
  • Cell Lists
• Multi-Image Neighbor Lists
  • Neighbor List Construction
  • Data Structures
  • Capacity Estimation
  • Graph Neural Network Support
• Potential Energy Functions
  • Bond Potentials
  • Classical Potentials
  • Neural Network Potentials
  • Helper Functions
• Energy Minimization Routines
  • Minimization Routines
  • Data Types
• Simulation Routines
  • Deterministic Simulation Environments
  • Stochastic Simulation Environments
  • Helper Functions
  • Testing Functions
  • Data Types
• Neural Network Primitives
  • Behler-Parrinello Networks
  • Graph Neural Networks
• Useful Quantities
  • Forces
  • Structural Features
  • Physical Properties
  • Data Types
• Elasticity Routines
  • Primary Functions
  • Helper Functions
• Rigid Body Simulations
  • Quaternion Utilities
  • Rigid Body Simulation
  • Rigid Collections of Points
• Higher Order Functions
  • ParameterTree
  • ParameterTreeMapping
  • bond()
  • pair()
  • pair_neighbor_list()
  • triplet()
• Multi-Image Structure Maps
  • Pair Functions
• jax_md.interpolate module
  • canonicalize()
  • constant()
  • spline()
• jax_md.dataclasses module
  • asdict()
  • astuple()
  • dataclass()
  • field()
  • fields()
  • is_dataclass()
  • replace()
  • static_field()
  • unpack()
• jax_md.mm_forcefields.reaxff package
  • Submodules

Indices and tables

• Index

• Module Index

• Search Page