Feature Request or Modifying Pair Potential

[moradza]

I want to add new 2- and 3-body potentials is there any documentation that I can start with. Particularly, I want to use a NN force field in the JAX-MD, trained using TensorFlow.

Thanks,
Alireza

[sschoenholz]

Thanks for the message! One of our big goals is to improve our documentation.

Simulations in JAX MD can use either 1) any "energy" function whose signature is energy_fn(R, **kwargs) mapping ndarrays of shape (N, dim) to real numbers or 2) any "force" function, force_fn(R, **kwargs) mapping from ndarrays of positions of shape(N, dim) to forces fo shape (N, dim). These functions either be written by hand or neural networks or whatever else you'd like as long as they share the above signature. We additionally include helper functions (in smap.py) that convert functions from pairwise distances / displacements to a function of the correct form above.

To start with adding a new potential I would look at either the energy functions defined in energy.py. In the JAX MD cookbook we also go through an example where we train a neural network that could be used as a drop in energy function by JAX MD.

Finally, note that JAX MD is written using JAX which, for the time being, can't obviously interoperate with tensorflow. To use a NN with JAX it should be written in JAX. If you have a neural network already trained in TF I would recommend either 1) retraining the NN with JAX or 2) loading the weights of the NN using numpy and using them to define a JAX neural network. This latter method might be the most direct but a little finicky.

In general, we'd love to help you get this use case working since it's one of the reasons we wanted to have JAX MD. Please don't hesitate to post followup questions. We will be working on a tutorial specifically targeting adding new potentials in the immediate future.

[moradza]

Thanks for your reply. I will go through the JAX-MD tutorial and cookbook. The problem that I am trying to solve is similar to Stillinger-Weber (SW) Potentials, in particular, it has two terms: 1. 2body similar to LJ pair potential 2. 3body potential as discussed in the lammps webpage https://lammps.sandia.gov/doc/pair_sw.html. I am not yet familiar enough with JAX-MD and I want to know your opinion if it can be used for a similar system as SW potential.

[sschoenholz]

Thanks for clarifying! I absolutely think that SW could be implemented in JAX MD pretty easily. Actually @ekindogus has been working on the underlying primitives that one would need to do this. We will followup here as we make progress.

[sschoenholz]

@ekindogus just merged PR #60 which includes the function quantity.cosine_angles(dR) which computes the cosine angle given a set of neighbors. Here dR is an [N, N_neighbors, dimension] array and cosine_angles(dR) returns an [N, N_neighbors, N_neighbors] array containing the cosine angle between pairs i, j, k. We think this is a useful primitive when approaching something like SW potentials.

Stay tuned for more in this direction soon.

[cpgoodri]

Hi @moradza, I have been working on a cookbook on how to set up custom potentials. If you are interested, see the PR: #76. If anything isn't clear or if there are any obvious gaps, please let me know.