Can we utilize sparsity in `grid_AO`?

[AlexanderMath]

In the image below we see ao and electron_repulsion share a similar range/distribution. We can remove a big chunk of elements from electron_repulsion using electron_repulsion[np.abs(electron_repulsion)<opts.threshold]. Running DFT for opts.threshold=1e-8 doesn't lead to large changes in the output values.

Question 1. Can we use a similar trick for ao? Q1 (a) How large are the errors if we similarly zero out ao[np.abs(ao)<1e-i]=0 for i=2,3,4,5,6,7,8,9? (b) How many floats do we save?

Question 2. With electron_repulsion we managed to perform a sparse multiplication using np.take/segment_sum. The ao is used in pmap_nanoDFT.py sharded between IPUs in def exchange_correlation(..., sharded_AO, ...). Most of the uses of sharded_AO can be written as matrix multiplications; is there any way we can utilize the additional sparsity.

[AlexanderMath]

import pyscf 
import matplotlib.pyplot as plt 
import pyscf.dft
import numpy as np

mol = pyscf.gto.Mole(atom=[["C", (0,0,i)] for i in range(3)])
mol.build()
grids            = pyscf.dft.gen_grid.Grids(mol)
grids.level      = 1 # {0,1,2,3}
grids.build()
grid_weights    = grids.weights                                 # (grid_size,) = (45624,) for C6H6
coord_str       = 'GTOval_cart_deriv1' if mol.cart else 'GTOval_sph_deriv1'
grid_AO         = mol.eval_gto(coord_str, grids.coords, 4) 
print(grid_AO)

for thresholds_exp in range(-10, -3):
  thresholds = 10**(thresholds_exp)
  grid_AO[np.abs(grid_AO) < thresholds] = 0
  print(thresholds_exp, np.sum(grid_AO==0) / grid_AO.size)

[blazejba]

The answer attempted here: https://github.com/graphcore-research/pyscf-ipu/pull/78