WIP: nuclear gradients

[hatemhelal]

To perform tasks such as geometry optimisation or molecular dynamics we need a way to compute the gradient of the total energy with respect to the nuclei centers. This PR adds the analytic evaluation of components of the energy:

• overlap -> $\langle \nabla i | j \rangle$
• kinetic energy -> $-\frac{1}{2} \langle \nabla i | \nabla^2 |\phi_j \rangle$
• nuclear attraction -> $\langle \nabla i | V_N | j \rangle$
• two-electron integrals -> $\langle \nabla i j | k l \rangle$
• exchange-correlation -> tdb
• nuclear-nuclear energy -> $\nabla V_{NN}$

[awf]

Would be very nice to link to a derivation in sympy if possible.

[hatemhelal]

I had a half-hearted attempt of doing this in sympy but found it difficult to encourage it to collect terms together. A derivation would be easy enough to write out in the gto_integrals notebook.

[awf]

I realise this is sacrilege, but I wonder if a python loop over 0,1,2 might be easier to read, and ultimately be no slower, as it will be vectorisable by a higher level vmap.
E.g. something like (artist's impression)

grad = jnp.zero(3)
for axis in range(3):
  t1 = 2 * a.alpha * primitive_op(a.offset_lmn[axis], 1, b)
  t2 = a.offset_lmn[axis] * primitive_op(a.offset_lmn[axis], -1, b)
  grad = grad.at(axis).set(t1 - t2)

[awf]

Or if that turns out to be slower, it might serve as reference code in another function so testing can compare them.

[hatemhelal]

trying out a list comprehension like this:

    t1 = [primitive_op(a.offset_lmn(ax, 1), b) for ax in range(3)]
    t2 = [primitive_op(a.offset_lmn(ax, -1), b) for ax in range(3)]
    grad_out = 2 * a.alpha * jnp.stack(t1) - a.lmn * jnp.stack(t2)
    return grad_out

better or worse?

[awf]

Definitely easier to read - perf impact ?

[hatemhelal]

I did a very quick benchmark using the test case on the CPU backend:

	no-jit	with jit
before	1.06 s ± 2.43 ms	359 µs ± 1.93 µs
after	2.77 s ± 9.36 ms	368 µs ± 2.65 µs

so arguably no difference when we use JIT but quite a bit slower without. Happy to go with the more readable code over the "vectorise or die trying" version.

[awf]

What's the signature of this Callable? It's great to point to an example, even from this more generic code to less generic.

[hatemhelal]

I've added some doc-strings throughout this file and tried to document the signature required by callables being passed around.

[awf]

I like to define these local lambdas after the variables to which they refer have been defined. It puts the computation closer to the point of use.

(i.e. move it after line 37)

OTOH, batch_orbitals is already doing a lot of Python list comprehensions, so it might be clearer and the same complexity to treat b as a list of lists of primitives, and just inline the list comprehensions here?

[hatemhelal]

I'll have a think, this lambda is actually used in a few places now -> maybe it should be promoted from a function-local to a definition in a module?

[awf]

Does this jit change much? It feels a bit out of place, but I get that it might be needed.

[hatemhelal]

hard to say...I've removed it for now since I'd like to compare jit to no-jit when this is a bit more complete.

[awf]

This is a back-pointer into the list of which this is a member?
That's also a hint that the where(a.atom_index == atom_index) could be lifted.

[hatemhelal]

indeed, the problem I haven't solved is not redundantly storing the center on both the primitives and in the Structure