WagnerGroup/pyqmc

Extra evaluations of aos for updateinternals

lkwagner opened this issue · 1 comments

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().

This was addressed by #346