diff --git a/vayesta/core/types/wf/fci.py b/vayesta/core/types/wf/fci.py index 4d900e50..a24083eb 100644 --- a/vayesta/core/types/wf/fci.py +++ b/vayesta/core/types/wf/fci.py @@ -536,11 +536,11 @@ def dummy_mo(self): return type(self.mo)(coeff_w_dummy, occ=self.mo.nocc) def _phys_ind_orbs(self): - return [np.array([i for i in range(y) if i not in x]) for x, y in zip(self.dummy_orbs, self.norb)] + return [np.array([i for i in range(y) if i not in x], dtype=np.int64) for x, y in zip(self.dummy_orbs, self.norb)] def _phys_ind_vir(self): return [ - np.array([i for i in range(y) if i + z not in x]) for x, y, z in zip(self.dummy_orbs, self.nvir, self.nocc) + np.array([i for i in range(y) if i + z not in x], dtype=np.int64) for x, y, z in zip(self.dummy_orbs, self.nvir, self.nocc) ] def make_rdm1(self, ao_basis=False, *args, **kwargs): @@ -560,7 +560,6 @@ def make_rdm2(self, ao_basis=False, *args, **kwargs): return (dm2[0][np.ix_(sa, sa, sa, sa)], dm2[1][np.ix_(sa, sa, sb, sb)], dm2[2][np.ix_(sb, sb, sb, sb)]) def as_cisd(self, *args, **kwargs): - self.check_norb() with replace_attr(self, mo=self.dummy_mo): wf_cisd = super().as_cisd(*args, **kwargs) va, vb = self._phys_ind_vir() diff --git a/vayesta/ewf/ufragment.py b/vayesta/ewf/ufragment.py index 8b57f61f..a015d9d6 100644 --- a/vayesta/ewf/ufragment.py +++ b/vayesta/ewf/ufragment.py @@ -43,35 +43,10 @@ def get_fragment_energy(self, c1, c2, hamil=None, fock=None, axis1="fragment", c e_corr: float Total fragment correlation energy contribution. """ - nocc = (c2[0].shape[1], c2[-1].shape[1]) - nvir = (c2[0].shape[2], c2[-1].shape[2]) - self.log.debugv("nocc= %d, %d nvir= %d, %d", *nocc, *nvir) - oa, ob = np.s_[: nocc[0]], np.s_[: nocc[1]] - va, vb = np.s_[nocc[0] :], np.s_[nocc[1] :] if axis1 == "fragment": pxa, pxb = self.get_overlap("proj|cluster-occ") + oaf, obf = np.s_[: pxa.shape[0]], np.s_[: pxb.shape[0]] - # --- Singles energy (zero for HF-reference) - if c1 is not None: - # if hasattr(eris, 'fock'): - # fa = eris.fock[0][oa,va] - # fb = eris.fock[1][ob,vb] - # else: - # fock = self.base.get_fock() - # fa = dot(self.c_active_occ[0].T, fock[0], self.c_active_vir[0]) - # fb = dot(self.c_active_occ[1].T, fock[1], self.c_active_vir[1]) - if fock is None: - fock = self.base.get_fock_for_energy() - fova = dot(self.cluster.c_active_occ[0].T, fock[0], self.cluster.c_active_vir[0]) - fovb = dot(self.cluster.c_active_occ[1].T, fock[1], self.cluster.c_active_vir[1]) - assert len(c1) == 2 - ca, cb = c1 - if axis1 == "fragment": - e_singles = einsum("ia,xi,xa->", fova, pxa, ca) + einsum("ia,xi,xa->", fovb, pxb, cb) - else: - e_singles = np.sum(fova * ca) + np.sum(fovb * cb) - else: - e_singles = 0 # Doubles energy # TODO: loop to reduce memory? if hamil is None: @@ -80,9 +55,22 @@ def get_fragment_energy(self, c1, c2, hamil=None, fock=None, axis1="fragment", c gab = hamil.get_eris_bare(block="ovOV") gbb = hamil.get_eris_bare(block="OVOV") + nocc = gab.shape[0], gab.shape[2] + nvir = gab.shape[1], gab.shape[3] + self.log.debugv("nocc= %d, %d nvir= %d, %d", *nocc, *nvir) + oa, ob = np.s_[:nocc[0]], np.s_[:nocc[1]] + va, vb = np.s_[:nvir[0]], np.s_[:nvir[1]] + if axis1 == "fragment": assert len(c2) == 4 caa, cab, cba, cbb = c2 + + # Remove padding + caa = caa[oaf, oa, va, va] + cab = cab[oaf, ob, va, vb] + cba = cba[obf, oa, vb, va] + cbb = cbb[obf, ob, vb, vb] + if c2ba_order == "ab": cba = cba.transpose(1, 0, 3, 2) e_doubles = ( @@ -93,9 +81,16 @@ def get_fragment_energy(self, c1, c2, hamil=None, fock=None, axis1="fragment", c + einsum("xi,xjab,iajb", pxa, cab, gab) / 2 + einsum("xi,xjab,jbia", pxb, cba, gab) / 2 ) + else: assert len(c2) == 3 caa, cab, cbb = c2 + + # Remove padding + caa = caa[oa, oa, va, va] + cab = cab[oa, ob, va, vb] + cbb = cbb[ob, ob, vb, vb] + e_doubles = ( einsum("ijab,iajb", caa, gaa) / 4 - einsum("ijab,ibja", caa, gaa) / 4 @@ -104,6 +99,43 @@ def get_fragment_energy(self, c1, c2, hamil=None, fock=None, axis1="fragment", c + einsum("ijab,iajb", cab, gab) ) + # --- Singles energy (zero for HF-reference) + if c1 is not None: + # if hasattr(eris, 'fock'): + # fa = eris.fock[0][oa,va] + # fb = eris.fock[1][ob,vb] + # else: + # fock = self.base.get_fock() + # fa = dot(self.c_active_occ[0].T, fock[0], self.c_active_vir[0]) + # fb = dot(self.c_active_occ[1].T, fock[1], self.c_active_vir[1]) + if fock is None: + fock = self.base.get_fock_for_energy() + fova = dot(self.cluster.c_active_occ[0].T, fock[0], self.cluster.c_active_vir[0]) + fovb = dot(self.cluster.c_active_occ[1].T, fock[1], self.cluster.c_active_vir[1]) + assert len(c1) == 2 + ca, cb = c1 + + # Remove padding + ca = ca[oaf, va] + cb = cb[obf, vb] + + e_singles = 0 + + if fova.shape[0] != 0 and fova.shape[1] != 0: + if axis1 == "fragment": + e_singles += einsum("ia,xi,xa->", fova, pxa, ca) + else: + e_singles += np.sum(fova * ca) + + if fovb.shape[0] != 0 and fovb.shape[1] != 0: + if axis1 == "fragment": + e_singles += einsum("ia,xi,xa->", fovb, pxb, cb) + else: + e_singles += np.sum(fovb * cb) + + else: + e_singles = 0 + e_singles = self.sym_factor * e_singles e_doubles = self.sym_factor * e_doubles e_corr = e_singles + e_doubles