def test_2_2(self): refstate = cache.refstate["cn_sto3g"] a = nosym_like(refstate.fock("o1v1")).set_random() b = nosym_like(refstate.fock("v1o1")).set_random() no, nv = a.shape res = direct_sum("ia+bj", a, b) assert res.shape == (no, nv, nv, no) ref = a.to_ndarray()[:, :, None, None] + b.to_ndarray()[None, None, :, :] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("-ia+bj->jabi", a, b) assert res.shape == (no, nv, nv, no) ref = -a.to_ndarray()[:, :, None, None] + b.to_ndarray()[None, None, :, :] ref = ref.transpose((3, 1, 2, 0)) assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("-ia-bj->abij", a, b) assert res.shape == (nv, nv, no, no) ref = -a.to_ndarray()[:, :, None, None] - b.to_ndarray()[None, None, :, :] ref = ref.transpose((1, 2, 0, 3)) assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_nontrivial_diagonal(self): refstate = cache.refstate["cn_sto3g"] mtcs = [ nosym_like(refstate.eri("o1o1v1v1")).set_random(), nosym_like(refstate.eri("o1v1o1v1")).set_random(), nosym_like(refstate.eri("o1o1v1v1")).set_random() ] mnps = [m.to_ndarray() for m in mtcs] res = einsum("ijab,ibkc,kjad->cd", mtcs[0], mtcs[1], mtcs[2]).diagonal() ref = np.einsum("ijab,ibkc,kjad->cd", mnps[0], mnps[1], mnps[2]).diagonal() assert res.needs_evaluation assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_nontrivial_direct_sum(self): refstate = cache.refstate["cn_sto3g"] oeo = nosym_like(refstate.orbital_energies("o1")).set_random() oovv = empty_like(refstate.eri("o1o1v1v1")).set_random() oev = nosym_like(refstate.orbital_energies("v1")).set_random() interm = np.einsum("ijac,ijcb->ab", oovv.to_ndarray(), oovv.to_ndarray()) ref = (interm[:, :, None, None] - oeo.to_ndarray()[None, None, :, None] - oev.to_ndarray()[None, None, None, :]) ref = ref.transpose((2, 0, 3, 1)) res = direct_sum("ab-i-c->iacb", einsum("ijac,ijcb->ab", oovv, oovv), oeo, oev) assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_1_2_1(self): refstate = cache.refstate["cn_sto3g"] oeo = nosym_like(refstate.orbital_energies("o1")).set_random() fvv = nosym_like(refstate.fock("v1v1")).set_random() oev = nosym_like(refstate.orbital_energies("v1")).set_random() res = direct_sum("i+ab-c", oeo, fvv, oev) ref = (oeo.to_ndarray()[:, None, None, None] + fvv.to_ndarray()[None, :, :, None] - oev.to_ndarray()[None, None, None, :]) assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("ab-i-c->iacb", fvv, oeo, oev) ref = (fvv.to_ndarray()[:, :, None, None] - oeo.to_ndarray()[None, None, :, None] - oev.to_ndarray()[None, None, None, :]) ref = ref.transpose((2, 0, 3, 1)) assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_nontrivial_symmetrisation(self): refstate = cache.refstate["cn_sto3g"] mtcs = [ nosym_like(refstate.eri("o1o1v1v1")).set_random(), nosym_like(refstate.eri("o1v1o1v1")).set_random(), nosym_like(refstate.eri("o1o1v1v1")).set_random() ] mnps = [m.to_ndarray() for m in mtcs] res = (mtcs[0] / mtcs[1].transpose((0, 2, 3, 1))) res -= 2 * mtcs[2].antisymmetrise((0, 1)) res = res.symmetrise((0, 1), (2, 3)) ref = (mnps[0] / mnps[1].transpose((0, 2, 3, 1))) ref -= (mnps[2] - mnps[2].transpose((1, 0, 2, 3))) ref = 0.5 * (ref + ref.transpose((1, 0, 3, 2))) assert res.needs_evaluation assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_1_1(self): refstate = cache.refstate["cn_sto3g"] a = nosym_like(refstate.orbital_energies("o1")).set_random() b = nosym_like(refstate.orbital_energies("o1")).set_random() res = direct_sum("i+j", a, b) ref = a.to_ndarray()[:, None] + b.to_ndarray()[None, :] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("i+j->ji", a, b) ref = a.to_ndarray()[None, :] + b.to_ndarray()[:, None] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("i-j", a, b) ref = a.to_ndarray()[:, None] - b.to_ndarray()[None, :] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("-i+j->ji", a, b) ref = -a.to_ndarray()[None, :] + b.to_ndarray()[:, None] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14) res = direct_sum("-i-j", a, b) ref = -a.to_ndarray()[:, None] - b.to_ndarray()[None, :] assert_allclose(res.to_ndarray(), ref, rtol=1e-10, atol=1e-14)
def test_einsum_2_2_2(self): # (1, 2, 2) in C++ refstate = cache.refstate["h2o_sto3g"] a = nosym_like(refstate.fock("o1v1")) b = nosym_like(refstate.fock("v1v1")) self.base_test("ij,jk->ik", a, b)
def test_einsum_1_2_1_ji(self): # (1,1,2) in C++ refstate = cache.refstate["h2o_sto3g"] a = nosym_like(refstate.orbital_energies("o1")) b = nosym_like(refstate.fock("v1o1")) self.base_test("i,ji->j", a, b)
def test_einsum_1_1_0(self): refstate = cache.refstate["h2o_sto3g"] a = nosym_like(refstate.orbital_energies("o1")) b = nosym_like(refstate.orbital_energies("v1")) self.base_test("i,j->ji", a, b)
def test_thc(self): refstate = cache.refstate["h2o_sto3g"] a = nosym_like(refstate.foo) b = nosym_like(refstate.fov) c = nosym_like(refstate.foo) self.base_test("ij,ia,ik->jka", a, b, c)
def test_contract_4_2_4_perm(self): # (1, 4, 2) in C++ refstate = cache.refstate["cn_sto3g"] a = empty_like(refstate.eri("o1o1v1v1")) b = empty_like(refstate.fock("v1v1")) out = nosym_like(refstate.eri("o1v1o1v1")) self.base_test("ijkl,km->imjl", a, b, out)
def test_contract_2_4_2(self): # (2, 2, 4) in C++ refstate = cache.refstate["h2o_sto3g"] a = empty_like(refstate.fock("o1o1")) b = empty_like(refstate.eri("o1v1o1v1")) out = nosym_like(refstate.fock("v1v1")) self.base_test("ij,ikjl->kl", a, b, out)
def test_contract_2_1_1(self): # (1,2,1) in C++ refstate = cache.refstate["h2o_sto3g"] a = empty_like(refstate.fock("o1v1")) b = empty_like(refstate.orbital_energies("v1")) out = nosym_like(refstate.orbital_energies("o1")) self.base_test("ij,j->i", a, b, out)
def test_contract_1_2_1_ij(self): # (1,1,2) in C++ refstate = cache.refstate["h2o_sto3g"] a = nosym_like(refstate.orbital_energies("o1")) b = nosym_like(refstate.fock("o1v1")) out = nosym_like(refstate.orbital_energies("v1")) self.base_test("i,ij->j", a, b, out)
def test_contract_4_4_2(self): # (3, 4, 4) in C++ refstate = cache.refstate["cn_sto3g"] a = empty_like(refstate.eri("o1v1o1v1")) b = empty_like(refstate.eri("o1o1o1v1")) out = nosym_like(refstate.fock("o1v1")) self.base_test("iajb,jikb->ka", a, b, out)
def test_contract_4_4_4_pjoi(self): # (2, 4, 4) in C++ refstate = cache.refstate["cn_sto3g"] a = empty_like(refstate.eri("o1o1v1v1")) b = empty_like(refstate.eri("v1v1v1v1")) out = nosym_like(refstate.eri("o1v1o1v1")) self.base_test("opvw,vijw->pjoi", a, b, out)