jax-md

This release contains a large refactor of the JAX MD simulation environments as well as new features built on this refactor including support for rigid body physics and preliminary support for simulations on TPU. This version also contains a preliminary version of Nequip as well as code to deserialize saved models.

• Refactor of simulation environments and rigid body physics with (@ellamichelleking, @xiyudu, @mc2engel, @rkruegs123)
• Preliminary TPU code with @Sohl-Dickstein.
• Neural network infrastructure + nequip with @simonbatzner, @doguscubuk, @amilmerchant.
• Improved handling of neighbor lists in triclinic systems.

• Adds an equivariant graph neural network example.
• Miscellaneous fixes to the renderer (added a non-verbose mode and fixed the render frame inference in the case of a static snapshot).
• Neighbor lists now default to dr_threshold = 0.
• Add ParameterTree construct to smap functions. These allow users to pass PyTrees of data into smap functions with control over how those parameters are mapped among particles.

• Small fixes to neighbor list / cell list capacity logic.
• Small fix to NPT simulations.
• Tersoff model added by @SangHyukYoo.
• Documentation improvements (with more contributions from @adrhill).

NPT Fixes, documentation, and neighbor list improvements.

Changes:

• Significant documentation improvements by @adrhill.
• Add custom mask function to neighbor lists by @oliverdutton.
• Small fixes to renderer.
• Fixes to NPT simulation environment using fixed equation for pressure in v0.1.23.
• Add SyMD (symmetric molecular dynamics) example notebook based on https://github.com/whitead/symd.

Improves and adds support for a number of mechanical properties.

Changes:

• Adds perturbation parameter to all spaces.
• Fixes pressure computation.
• Adds stress tensor computation to quantity.py.
• Carl Goodrich (@cpgoodri) added support for athermal linear elasticity computations in elasticity.py, along with a great colab notebook!
• Patrick Marks (@pmarks) improved neighbor lists by switching to int32 for indices.
• Maxi Lechner (@MaxiLechner) added an implicit differentiation notebook combining JAX MD with JaxOpt.

More neighbor list improvements.

Changes:

• Improves overflow checking in neighbor lists by also checking for buffer overflows in cell lists.
• Modernizes cell list API to mimic neighbor lists (adds allocate and update functions).
• Neighbor list no longer carries the cell list update function explicitly.
• Thread neighbor list keyword arguments through energy functions.
• Remove scaling of dr_threshold by radius in soft sphere and Lennard-Jones potentials.

Neighbor List Refactor

New Features:

• Adds notion of neighbor list format. Previous behavior is defined as NeighborListFormat.Dense.
• Adds two new neighbor list formats NeighborListFormat.Sparse and NeighborListFormat.OrderedSparse.
• Sparse neighbor lists can be converted to jraph.GraphsTuple using the partition.to_jraph function.
• Graph networks now use Jraph as a backend when neighbor list is specified in the NeighborListFormat.Sparse format.
• Improved the semantics of neighbor list creation / updating by splitting neighbor list function into a pair neighbor_fn.allocate and neighbor_fn.update. Neighbor lists include a convenience update function nbrs = nbrs.update(position).
• Adds neighbor list support for Stillinger-Weber and EAM potentials.