Preparing for JAX

[robertodr]

List of TODOs to make the code differentiable with JAX

• [x] Figure out the unit conversions: we should work in atomic units throughout. Conversions to/from AU should be done outside the d3 function.
• [x] Document input parameters to the d3 function. Related to #2.
• [x] Rearrange input parameters to the d3 function as suitable for usage with JAX (see the jax_tests repository)
• [x] Review and refactor unpacking of coordinates into separate arrays for X, Y, and Z. We want to work with a coordinate array shaped as 3*N_atoms
• [x] Use atomic numbers, rather than symbols to index atoms.

[robertodr]

Replying here to the Zulip chat:

• [ ] Document input parameters to the d3 function. Related to #2.

The d3 function takes some input parameters to decide which flavor of the D3 model to use and which atom pairs to include/exclude. The inclusion/exclusion is mentioned in issue #2. It would be good if we document these parameters in a documentation string. This is not high priority though.

• [ ] Rearrange input parameters to the d3 function as suitable for usage with JAX (see the jax_tests repository)

In jax_tests the function to differentiate only accepts atomic charges and coordinates. Here it takes those, plus other options. I believe that, in order to make d3 differentiable, we need to change the order in which parameters are passed.

[robertodr]

I think this is done, so we can move on to computing derivatives :confetti_ball: