Mihai Polceanu

The O(N^2) code from [the post](https://chemistry.stackexchange.com/questions/169429/integral-prescreening-in-calculation-of-electron-repulsion-integrals) looks like this: ``` import pyscf import numpy as np mol = pyscf.gto.Mole(atom=[["C", (0,0,0)], ["C", (10,2,3)]]) mol.build() ERI = mol.intor("int2e_sph", aosym="s1") N = mol.nao_nr()...

But it turns out we can do better, if we integrate the symmetries in the above screening strategy: ``` import pyscf import numpy as np def get_i_j(val, xnp=np, dtype=np.uint64): i...

The problem seems to be related to the use of `partial` ([docs here](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html#caching)). Caching the jitted nanoDFT function works faster (ugly code to prove it): ``` cached_jitted_nanoDFT = None def...

Based on what I could find, one way is to derive the class, for example: ``` class HashableMole(pyscf.gto.mole.Mole): def __init__(self): pyscf.gto.mole.Mole.__init__(self) print('Hello HashableMole') def _tree_flatten(self): children = () # arrays...

@AlexanderMath > Q1. What happens if we remove all the jax tree stuff? (I might just be misunderstanding, but don't see why it's needed in this case) In fact, nothing...

@AlexanderMath > Q2. Could we be a bit cheeky and do mol.__hash__ = self.nao_nr; mol.__eq__ = lambda self, other: ...? This seems to work: ``` pyscf.gto.mole.Mole.__hash__ = lambda self: hash(self.nao_nr())...