Molecule similarity aware jax.vmap(nanoDFT)

[AlexanderMath]

Our DFT function considers only a single molecule. We could compute DFT for a batch of molecules using jax.vmap.

Problem: This consumes more on-chip memory.

Idea: Take a molecule a. Construct molecule b by modifying the position a single atom. Most floats we need to store for a and b overlap: only O(N^3) floats will differ. If the molecule has n_atom<64 atoms we could then make batch_size=64<n_atom such "one-atom" perturbations while increasing memory consumption by only batch_size*n_atom^3.

Note: From deep learning perspective this feels similar to data augmentation. From DFT we are just re-using stuff when doing batch_size=64 similar DFTs.

[AlexanderMath]

Note: This also reduces FLOPs

Note: Consider GD wrt density_matrix as in D4FT. If we instead do GD wrt NN(mol)=dm and then DFT_energy(NN(mol)) a single DFT iteration is sufficient to evaluate how well NN did.

Note: This also works if we change within {C,N,O,F}.