Extra evaluations of aos for updateinternals
lkwagner opened this issue · 1 comments
lkwagner commented
Updateinternals is currently written as this in the Slater function
def updateinternals(self, e, epos, mask=None):
"""Update any internals given that electron e moved to epos. mask is a Boolean array
which allows us to update only certain walkers"""
s = int(e >= self._nelec[0])
if mask is None:
mask = [True] * epos.configs.shape[0]
eeff = e - s * self._nelec[0]
ao = self.orbitals.aos("GTOval_sph", epos, mask)
self._aovals[:, mask, e, :] = ao
mo = self.orbitals.mos(ao, s)
mo_vals = mo[:, self._det_occup[s]]
det_ratio, self._inverse[s][mask, :, :, :] = sherman_morrison_ms(
eeff, self._inverse[s][mask, :, :, :], mo_vals
)
self._updateval(det_ratio, s, mask)
Note that it recomputes the AOs, when almost always they were already computed in testvalue() beforehand:
def testvalue(self, e, epos, mask=None):
"""return the ratio between the current wave function and the wave function if
electron e's position is replaced by epos"""
s = int(e >= self._nelec[0])
ao = self.orbitals.aos("GTOval_sph", epos, mask)
mo = self.orbitals.mos(ao, s)
mo_vals = mo[..., self._det_occup[s]]
if len(epos.configs.shape) > 2:
mo_vals = mo_vals.reshape(
-1, epos.configs.shape[1], mo_vals.shape[1], mo_vals.shape[2]
)
return gpu.asnumpy(self._testrow(e, mo_vals, mask))
This is a pretty significant inefficiency since our profiling has found that computing the MOs is pretty expensive.
A simple resolution might be to have testvalue also return temporaries, which then could be optionally passed into updateinternals(). That way any optimization could happen at the algorithmic level, and the wave function object doesn't have to try to keep track of whether a given testvalue() needs to be saved. For example, in ECP evaluation this doesn't apply because we never call updateinternals().