{"id":45181,"library":"jax-md","title":"JAX-MD","description":"JAX-MD is a differentiable, hardware-accelerated molecular dynamics library built on JAX. It supports both rigid and flexible molecules, NVE/NVT/NPT ensembles, and integrates with JAX's autograd for energy minimization and force computation. Current version 0.2.28, requires Python >=3.10. Active development with irregular releases.","status":"active","version":"0.2.28","language":"python","source_language":"en","source_url":"https://github.com/google/jax-md","tags":["molecular-dynamics","differentiable","jax","scientific-computing","gpu"],"install":[{"cmd":"pip install jax-md","lang":"bash","label":"Install from PyPI"},{"cmd":"pip install --upgrade jax jaxlib jax-md","lang":"bash","label":"Install with CPU/GPU support"}],"dependencies":[{"reason":"JAX-MD is built on JAX; must have compatible version (jax>=0.4).","package":"jax","optional":false},{"reason":"Provides XLA backend; required for GPU support.","package":"jaxlib","optional":true},{"reason":"Used for neural network potentials (optional).","package":"dm-haiku","optional":true},{"reason":"Used for optimization routines (optional).","package":"optax","optional":true}],"imports":[{"note":"jax_md is a package; submodules must be imported explicitly.","wrong":"import jax_md.simulate","symbol":"simulate","correct":"from jax_md import simulate"},{"note":"Same as above.","wrong":"import jax_md.energy","symbol":"energy","correct":"from jax_md import energy"},{"note":"Same pattern.","wrong":"import jax_md.space","symbol":"space","correct":"from jax_md import space"},{"note":"Neighbor list partitioning.","wrong":"import jax_md.partition","symbol":"partition","correct":"from jax_md import partition"}],"quickstart":{"code":"import jax\nimport jax.numpy as jnp\nfrom jax_md import simulate, energy, space, quantity\n\n# Set up a simple cubic lattice of particles\ndimension = 2\nbox_size = 5.0\ndisplacement_fn, shift_fn = space.periodic(box_size)\n\n# Create positions on a lattice\nN = 9\nlattice = space.initialize_canonical_lattice(N, box_size=box_size, dim=dimension, center=[0.0, 0.0])\npositions = lattice['position']\n\n# Define a soft sphere interaction\nenergy_fn = energy.soft_sphere(displacement_fn, sigma=1.0, epsilon=1.0)\n\n# Initialize neighbor list\nneighbor_fn = partition.neighbor_list(displacement_fn, box_size, r_cutoff=2.5, capacity_multiplier=1.2)\nneighbor_list = neighbor_fn.allocate(positions)\n\n# Wrap energy function with neighbor list\ndef total_energy(R, **kwargs):\n    return energy_fn(R, neighbor_list.idx)\n\n# Run simulation with NVE\ninit, apply = simulate.nve(energy_fn, shift_fn, dt=0.001, T0=1.0)\nstate = init(jax.random.PRNGKey(0), positions, neighbor_list=neighbor_list)\n\n# Simulate for 100 steps\nfor i in range(100):\n    state = apply(state, neighbor_list=neighbor_list)\n    if i % 10 == 0:\n        print(f'Step {i}, KE={quantity.kinetic_energy(state.velocity, state.mass):.3f}')\n\nprint('Quickstart complete.')","lang":"python","description":"A basic NVE (constant energy) simulation of soft spheres in 2D with a periodic box and neighbor lists."},"warnings":[{"fix":"Use `neighbor_fn.allocate(positions)` and pass the neighbor list via `state = init(..., neighbor_list=neighbor_list)` and each apply call: `state = apply(state, neighbor_list=neighbor_list)`.","message":"JAX-MD 0.2.x changed the neighbor list API: `neighbor_list` no longer returns a static object; must be passed through simulation state.","severity":"breaking","affected_versions":">=0.2.0"},{"fix":"Replace `quantity.kinetic_energy(state)` with `quantity.kinetic_energy(state.velocity, state.mass)`.","message":"The `quantity.kinetic_energy` function now requires `velocity` and `mass` as separate arguments; previously accepted a state object.","severity":"deprecated","affected_versions":">=0.2.20"},{"fix":"Always pass PRNGKey and neighbor lists through function arguments; never rely on mutation.","message":"JAX-MD uses JAX's functional programming; mutable state (e.g., random keys, neighbor lists) must be passed explicitly—global state will not update.","severity":"gotcha","affected_versions":"all"},{"fix":"Install appropriate jaxlib: `pip install jaxlib==0.4.28+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`","message":"GPU support requires installing jaxlib with CUDA; otherwise runs on CPU silently.","severity":"gotcha","affected_versions":"all"},{"fix":"Use `energy.lennard_jones_pair` with appropriate parameters.","message":"The `energy.lennard_jones` function is deprecated in favor of `energy.lennard_jones_pair` with corrected normalization.","severity":"deprecated","affected_versions":">=0.2.15"}],"env_vars":null,"search_vec":"'0.2.28':42 '3.10':45 'acceler':12 'activ':46 'autograd':33 'built':16 'comput':39,58 'current':40 'develop':47 'differenti':9,54 'dynam':14,53 'energi':35 'ensembl':27 'flexibl':24 'forc':38 'gpu':59 'hardwar':11 'hardware-acceler':10 'integr':29 'irregular':49 'jax':2,5,18,31,55 'jax-md':1,4 'librari':15 'md':3,6 'minim':36 'molecul':25 'molecular':13,52 'molecular-dynam':51 'nve/nvt/npt':26 'python':44 'releas':50 'requir':43 'rigid':22 'scientif':57 'scientific-comput':56 'support':20 'version':41","created_at":"2026-06-07T12:53:45.143190+00:00","updated_at":"2026-06-07T12:53:45.143190+00:00","problems":[{"fix":"Run `pip install jax-md`.","cause":"Package not installed.","error":"ModuleNotFoundError: No module named 'jax_md'"},{"fix":"Use `from jax_md import simulate`.","cause":"Incorrect import: using `import jax_md; jax_md.simulate` instead of explicit submodule import.","error":"AttributeError: module 'jax_md' has no attribute 'simulate'"},{"fix":"Pass the displacement function: `energy.smooth_barrier(displacement_fn, ...)`.","cause":"Energy functions require displacement function as first argument.","error":"TypeError: smooth_barrier() missing 1 required positional argument: 'displacement_fn'"},{"fix":"Construct neighbor lists outside the differentiated function, e.g., pre-allocate and pass as static argument.","cause":"Attempting to differentiate through a function that uses a non-differentiable operation (e.g., neighbor list construction inside a loop).","error":"jax.errors.UnexpectedTracerError: Cannot differentiate with respect to argument 0"}],"ecosystem":"pypi","meta_description":null,"install_score":null,"quickstart_score":null,"quickstart_tag":null,"pypi_latest":null,"cli_name":null,"cli_version":null,"type":"library","homepage":"https://github.com/google/jax-md","github":"https://github.com/google/jax-md","docs":null,"changelog":null,"pypi":null,"npm":null,"openapi_spec":null,"status_page":null,"smithery":null,"categories":["ai-ml"],"base_url":null,"auth_type":null,"provenance":{"verified_status":"passing","verified_at":"2026-06-29","last_verified":"2026-06-29","next_check":"2026-07-29","install_tag":null}}