ase + jax-md = asax
asax exposes jax-md energy functions as ase calculators.
This is a very early-stage experiment. Please do not use it. Feel free to help out though!
Needs ase, jax-md, numpy. nose to run tests. (cd tests; nosetests).
Tests currently fail unless ase.calculators.lj is supports "smooth" cutoff. Recent master versions support this change, since this MR.
Codestyle black, roughly Google style otherwise.
poetry install in this folder will install an editable version.
jax-md is fundamentally built around functional programming: In the end, you get a function that maps positions -> energies. Implicitly in this, the positions are mapped to displacements (otherwise, the energy function wouldn't be translationally invariant).
This function is built up with the setup() method of the calculator sub-classes implemented here. Since it implicitly depends on the boundary conditions, this re-building needs to be repeated every time the unit cell changes. This logic is taken care of in asax.Calculator.
We only treat full periodic boundary conditions, i.e. 3D systems, or no boundary conditions.