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.

Reference

• Defining Spaces
  • Spaces
  • Higher Order Functions
  • Helper Functions
• Spatial Partitioning
  • Neighbor Lists
  • Cell Lists
• 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()

Indices and tables

• Index

• Module Index

• Search Page