Implementation of integrals over Gaussian orbitals

[hatemhelal]

This patch adds the following objects:

• Primitive: dataclass for a primitive Gaussian type orbital
• Orbital: dataclass for a linear combination of Primitive, these are sometimes called contractions in the literature.
• Basis: dataclass of Orbitals with an associated factory function basisset that can read from the basis-set-exchange database.
• Structure: dataclass for storing molecular structure (geometry and atom types)

Adds implementations as derived in Gaussian-expansion methods for molecular integrals for the following:

• overlap
• kinetic energy
• nuclear attraction
• electron-repulsion

All these integrals are implemented at the Primitive level and the results from these are used to build up the result of integration over the entire Basis.

Adds the following dependencies:

• basis-set-exchange
• periodictable
• sympy

Also adds a github action for running tests using the jax cpu backend by default and a few tests that run on the IPU model. Note that the ERI tests cannot be tested on the github action runners due to requiring more than the 7GB of memory available.

[awf]

in Gaussian-expansion methods for molecular integrals f
This link is paywalled? Is there an openly viewable equivalent?

[awf]

yay!

[awf]

Some comments, all non-blockers, as I've known this code from before.

I wonder if it might be better to move the folder e.g.

pyscf-ipu/pyscf_ipu/experimental -> pyscf-ipu/experimental/jax-integrals

[awf]

Is this correct, or to be changed to j?

Suggested change

	for kdx in range(idx + 1):
	for kdx in range(jdx + 1):

[hatemhelal]

So this needs to be left as idx as these loops generate the canonicalised i, j, k, l indices that handles the 8-fold permutation symmetries. To be honest, I was following the description in the four-index transformations book chapter - I think I understand that this is effectively pruning the complete 4-d mesh based on considering associativity between pairs. Put another way, the edit changing idx -> jdx generates the following 5 indices:

>>> list(gen_ijkl_edit(2))
[(0, 0, 0, 0), (1, 0, 0, 0), (1, 1, 0, 0), (1, 1, 1, 0), (1, 1, 1, 1)]

which is missing the (1, 0, 1, 0) index.

>>> list(gen_ijkl(2))
[(0, 0, 0, 0), (1, 0, 0, 0), (1, 0, 1, 0), (1, 1, 0, 0), (1, 1, 1, 0), (1, 1, 1, 1)]

[awf]

[Probably not for this PR but...] Duplicated information here, as atom positions are both in orbitals and structure - how far coudl we get if we explicitly convert from Structure -> Basis and don't maintain the handy back-pointer?

[hatemhelal]

I have thought about removing the structure field from the Basis class as you point out the atom positions are redundantly stored with in the orbital / primitives. The trouble is we need to know the structure for computing the occupied orbitals. Definitely needs some more thought!

[hatemhelal]

I wonder if it might be better to move the folder e.g.

pyscf-ipu/pyscf_ipu/experimental -> pyscf-ipu/experimental/jax-integrals

I'd like to leave this until it becomes clearer how much we are going to build into the experimental module. Either at some point functionality here stops becoming experimental and is moved into a top-level namespace or it stops being useful and we might as well delete it right?