Add ReaxFF potential

[cagrikymk]

This PR adds ReaxFF potential to JAX-MD.
The implementation is derived from JAX-ReaxFF

Since ReaxFF potential is quite complex, I added 4 new files which contain all the necessary logic to run ReaxFF simulations using JAX-MD.

These files are:

1. reaxff_helper.py This file contains helper functions such as tools to read and write force field parameters
2. reaxff_forcefield.py This file contains JAX compatible "ForceField" class that is designed to store ReaxFF parameters
3. reaxff_energy.py Includes ReaxFF energy functions and sparse iterative charge solver for QEQ charge equilibration model which can be potentially used as a standalone charge method.
4. reaxff_interactions.py Contains the functions responsible for interaction list generation (2-body, 3-body, 4-body, h-bond and long range interactions). Also it includes the necessary "allocate", "update", "reallocate" and "energy" functions compatible with other implementations in JAX-MD.

There are some tasks need to be completed before merging:

• Improving the interaction list update/allocate logic to minimize # recompilations
• Refactoring the code to make it more compatible with the rest of JAX-MD
• Including the necessary tests for the core functionality
• Finalizing the overall file/module structure

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[cagrikymk]

Hello @ekindogus,
I think there is a bug in "tpu_test.py" as the automated tests failed for my reaxff branch but the error seems to be caused by the tpu related testing.

Am I misinterpreting the failed tests?

Cagri

[MarcBerneman]

Hello @ekindogus, I think there is a bug in "tpu_test.py" as the automated tests failed for my reaxff branch but the error seems to be caused by the tpu related testing.

Am I misinterpreting the failed tests?

Cagri

Hi Cagri,

It might be unrelated, but your issue may have been resolved by a previous pull request of mine.

Marc

[cagrikymk]

Hello @MarcBerneman,
My branch is up-to-date with the main branch and it contains the changes from your PR. However, the tpu test is failing with the following error:

=================================== FAILURES ===================================
______________ ConvolutionalMDTest.test_nve_numdims=2_topology=() ______________
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

self = <tpu_test.ConvolutionalMDTest testMethod=test_nve_numdims=2_topology=()>
num_dims = 2, topology = ()

....

E       jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
E       The error occurred while tracing the function body_fn at /home/runner/work/jax-md/jax-md/jax_md/tpu.py:816 for while_loop. This value became a tracer due to JAX operations on these lines:
E       
E         operation a:bool[] = eq b c
E           from line /home/runner/work/jax-md/jax-md/jax_md/tpu.py:1505 (_update_grid_locations)
E       See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

[MarcBerneman]

Hello @MarcBerneman, My branch is up-to-date with the main branch and it contains the changes from your PR. However, the tpu test is failing with the following error:

=================================== FAILURES ===================================
______________ ConvolutionalMDTest.test_nve_numdims=2_topology=() ______________
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

self = <tpu_test.ConvolutionalMDTest testMethod=test_nve_numdims=2_topology=()>
num_dims = 2, topology = ()

....

E       jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
E       The error occurred while tracing the function body_fn at /home/runner/work/jax-md/jax-md/jax_md/tpu.py:816 for while_loop. This value became a tracer due to JAX operations on these lines:
E       
E         operation a:bool[] = eq b c
E           from line /home/runner/work/jax-md/jax-md/jax_md/tpu.py:1505 (_update_grid_locations)
E       See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

I shall investigate this once I find some free time.