graphcore-research / pyscf-ipu

PySCF on IPU
https://github.com/graphcore-research/pyscf-ipu#readme
Apache License 2.0
41 stars 3 forks source link

Update basisset to use numpy over jax.numpy #126

Closed hatemhelal closed 9 months ago

hatemhelal commented 9 months ago

This updates the basisset factory function to prefer directly using numpy arrays over jax.numpy. Since a Basis consists of a number of constants we can directly manipulate those using numpy functions. For example, with this PR the Primitives are also normalised using numpy functions.

hatemhelal commented 9 months ago

Might we ever need to compute a gradient over the normalizing constant?

The only case I can think of where we might need that is if we are optimising the basis set exponents (alphas). For the use cases of structure optimisation and molecular dynamics we would expect the normalising term to be a constant since it doesn't depend on the basis centers.

That said, I think since the norm term is an optional we could support the basis set exponent optimisation use case by having a different factory function that creates a Primitive with a differentiable norm term.