Add periodic boundary conditions

[PabloAMC]

Currently Grad DFT only supports molecules. Add the capability to add periodic boundary conditions.

[jackbaker1001]

As a rough outline, this is what I think needs to be done.

1. Molecule objects will need to be able to store N_bands * N_k (number of k-points) Kohn-Sham orbitals instead of the currently supported N_bands. It is clear already that we may wish to create a new class all together for these periodic calculations (i.e, make a class Solid rather than Molecule) or we can rename Molecule to Atoms and have it handle both the periodic and isolated molecule cases.

2. Calculations for charge densities, 1RDMs and energies will now all require accumulation over the index k (k-points). This will also need to be taken into account when occupations are calculated.

We can update below when new things come up, but, interestingly, we may get away with not dealing with a lot of the more finicky parts of implementing PBC as (i) we never calculate KS Hamiltonian elements "directly" in a way which requires the explicit consideration of a periodic potential and interaction with image atoms: PySCF does this. Further KS Hamiltonians are calculated with automatic differentiation. So long as our energy and 1RDM expressions obey PBC (they should), then further autodiff calculated KS Hamiltonians should also respect PBC.

One additional consideration is to do with efficiency in BZ integration. Typically, we would like to work with weighted sums of k-points in the irreducible BZ. Hopefully this is already done in PySCF and we can extract it. If not, we may wish to work in spglib to do this.

[jackbaker1001]

@PabloAMC

Further to the above, I have figured out how to implement a periodic boundaries at the gamma point (BZ sampling at the Gamma point only; only converging the electronic structure with large supercells) using the existing Molecule object and some other small changes:

(1) When periodic PySCF mean field objects are read in by molecule_from_pyscf, an additional dimension is added to several arrays (molecular orbtial coefficients, density matrices etc) to denote the number of k-points. At the Gamma point, this is of course a redundant dimensions as there is only one k-point. Using np.squeeze this dimension can be removed making arrays once again compatible with the Molecule object.

(2) The definition of one electron and electron repulsion integrals (ERIs: which we use to compute the total energy to avoid having a potentially numerically unstable differentiable KS potential solver) change when in periodic boundaries. One electron integrals automatically are computed in this new way, but, in PBCs, (to my knowledge) PySCF does not implement exact electron repulsion integrals instead choosing to use the density fitted (df) approximation. These can be computed like:

from pyscf.pbc import df

repulsion_tensor = df.DF(cell).get_eri(compact=False).reshape(nao, nao, nao, nao)

where cell is the periodic system version of PySCF's Mole object and nao is the number of atomic orbitals. This expression can be used in place of the exact expression. This of course adds a layer of approximation to our solver, but, given that periodic df calculations can also be done in PySCF like:

from pyscf.pbc import gto, scf
cell = gto.M(
    atom = '''Na 0.0000 0.0000 0.0000
              H 1.4000 0.0000 0.0000''',
    a = '''3.8 0.0 0.0
           0.0 3.8 0.0
           0.0  0.0 3.8''',
    # pseudo = 'gth-pade',
    basis = 'sto-3g'
    )

kmf = scf.KUKS(cell, kpts=cell.make_kpts([1, 1, 1])).density_fit()
kmf = kmf.run()

we can learn from these types of calculations without approximation. I have verified that the total energies match in Grad DFT and PySCF when we compute the ERIs in this way. This works for all-electron and pseudopotential calculations. That being said, perhaps we should look into a method to compute the exact repulsion tensor efficiently.

I am going to make a PR for this gamma point version and move on to the full k-point sampling version next week. The full BZ sampling version will require further modifications including averaging over the BZ in various areas and a new container separate to Molecule (Solid) to handle the different shaped jax arrays.

[jackbaker1001]

Ok. The branch linked to this issue if for the full BZ sampling version I am working on. To summarize, the structure I am going for is to have a new module solid.py which mimics the structure of molcule.py. The main class Solid will mimic Molecule but hold arrays of different shapes to facilitate the extra k-points dimension. I will try to work with points in the irreducible BZ from the start as using the full grid of k-points is very inefficient. From some inspection, it looks like the PySCF cell object is able to compute the irreducible k points, but I am yet to figure out how to get their weights (without using spglib).

I will implement the mimic the key methods of Molecule in Solid, up to the level of being able to compute the total energy with LDA, GGA and possibly meta-GGA. If Exact exchange turns out to be easy, I will add this too.

More updates will follow here.

[jackbaker1001]

Just some progress updates here. I had previously looked to do all of this fully taking into advantage symmetries in the 1BZ. Unfortunately, for now, this seems a but out of scope for our initial implementation. There are a whole bunch of transforms (which map quantities between IBZ to full BZ) that need to be implemented in Jax which I did not account for (see https://github.com/pyscf/pyscf/blob/master/pyscf/pbc/symm/symmetry.py).

So, for now, although inefficient, we will work with the full 1BZ to avoud this stuff. I will implement functions with weights arguments however such that when this is sorted, we should be able to use the existing code for the version which takes full account of symmetry.