fix IdentityAddedMPO for singlet embedding

pyblock2/driver/core.py

@@ -244,6 +244,7 @@ def initialize_system(
         orb_sym=None,
         heis_twos=-1,
         heis_twosz=0,
+        singlet_embedding=True,
     ):
         bw = self.bw
         self.vacuum = bw.SX(0, 0, 0)
@@ -254,7 +255,17 @@ def initialize_system(
                 pg_irrep = self.pg_irrep
             else:
                 pg_irrep = 0
-        self.target = bw.SX(n_elec if heis_twosz == 0 else heis_twosz, spin, pg_irrep)
+        if not SymmetryTypes.SU2 in bw.symm_type:
+            singlet_embedding = False
+        if singlet_embedding:
+            assert heis_twosz == 0
+            self.target = bw.SX(n_elec + spin, 0, pg_irrep)
+            self.left_vacuum = bw.SX(spin, spin, 0)
+        else:
+            self.target = bw.SX(
+                n_elec if heis_twosz == 0 else heis_twosz, spin, pg_irrep
+            )
+            self.left_vacuum = self.vacuum
         self.n_sites = n_sites
         if orb_sym is None:
             self.orb_sym = bw.b.VectorUInt8([0] * self.n_sites)
@@ -751,6 +762,16 @@ def get_npdm(self, ket, pdm_type=1, bra=None, soc=False, site_type=1, iprint=0):
                 prule = bw.bs.ParallelRulePDM2QC(self.mpi)
             pmpo = bw.bs.ParallelMPO(pmpo, prule)
 
+        if soc:
+            mb_lv = mbra.info.left_dims_fci[0].quanta[0]
+            ml_lv = mket.info.left_dims_fci[0].quanta[0]
+            if mb_lv != ml_lv:
+                raise RuntimeError(
+                    "SOC 1PDM cannot be done with singlet_embedding for mpss"
+                    + " with different spins! Please consider setting "
+                    + "singlet_embedding=False in DMRGDriver.initialize_system(...)."
+                )
+
         pme = bw.bs.MovingEnvironment(pmpo, mbra, mket, "NPDM")
         pme.init_environments(iprint >= 2)
         pme.cached_contraction = True
@@ -897,11 +918,15 @@ def adjust_mps(self, ket, dot=1):
             elif ket.canonical_form[-1] == "C" and ket.canonical_form[-2] == "L":
                 ket.canonical_form = ket.canonical_form[:-1] + "S"
                 ket.center = ket.n_sites - 1
+            elif ket.center == ket.n_sites - 2 and ket.canonical_form[-2] == "L":
+                ket.center = ket.n_sites - 1
             if ket.canonical_form[0] == "M" and ket.canonical_form[1] == "R":
                 ket.canonical_form = "J" + ket.canonical_form[1:]
             elif ket.canonical_form[-1] == "M" and ket.canonical_form[-2] == "L":
                 ket.canonical_form = ket.canonical_form[:-1] + "T"
                 ket.center = ket.n_sites - 1
+            elif ket.center == ket.n_sites - 2 and ket.canonical_form[-2] == "L":
+                ket.center = ket.n_sites - 1
 
         ket.save_data()
         if self.mpi is not None:
@@ -1052,6 +1077,26 @@ def load_mps(self, tag, nroots=1):
         self.fix_restarting_mps(mps)
         return mps
 
+    def mps_change_singlet_embedding(self, mps, tag, forward):
+        cp_mps = mps.deep_copy(tag)
+        while cp_mps.center > 0:
+            cp_mps.move_left(self.ghamil.opf.cg, self.prule)
+        if forward:
+            cp_mps.to_singlet_embedding_wfn(self.ghamil.opf.cg, self.prule)
+        else:
+            cp_mps.from_singlet_embedding_wfn(self.ghamil.opf.cg, self.prule)
+        cp_mps.save_data()
+        cp_mps.info.save_data(self.scratch + "/%s-mps_info.bin" % tag)
+        if self.mpi is not None:
+            self.mpi.barrier()
+        return cp_mps
+
+    def mps_change_to_singlet_embedding(self, mps, tag):
+        return self.mps_change_singlet_embedding(mps, tag, forward=True)
+
+    def mps_change_from_singlet_embedding(self, mps, tag):
+        return self.mps_change_singlet_embedding(mps, tag, forward=False)
+
     def mps_change_precision(self, mps, tag):
         bw = self.bw
         assert tag != mps.info.tag
@@ -1089,7 +1134,9 @@ def get_random_mps(
             mps = bw.bs.MultiMPS(self.n_sites, center, dot, nroots)
         mps_info.tag = tag
         if full_fci:
-            mps_info.set_bond_dimension_full_fci(self.vacuum, self.vacuum)
+            mps_info.set_bond_dimension_full_fci(self.left_vacuum, self.vacuum)
+        else:
+            mps_info.set_bond_dimension_fci(self.left_vacuum, self.vacuum)
         if occs is not None:
             mps_info.set_bond_dimension_using_occ(bond_dim, bw.b.VectorDouble(occs))
         else:

pyblock2/unit_test/dmrg_excited.py

@@ -0,0 +1,92 @@
+import pytest
+import numpy as np
+from pyblock2._pyscf.ao2mo import integrals as itg
+from pyblock2.driver.core import DMRGDriver, SymmetryTypes
+
+pytestmark = pytest.mark.filterwarnings("ignore::DeprecationWarning")
+
+
+@pytest.fixture(scope="module", params=["N2"])
+def system_def(request):
+    from pyscf import gto
+
+    if request.param == "N2":
+        mol = gto.M(atom="N 0 0 0; N 0 0 1.1", basis="sto3g", symmetry="d2h", verbose=0)
+        return mol, 0, None, "N2"
+    elif request.param == "C2":
+        mol = gto.M(
+            atom="C 0 0 0; C 0 0 1.2425", basis="ccpvdz", symmetry="d2h", verbose=0
+        )
+        return mol, 2, 8, "C2"
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def singlet_embedding_type(request):
+    return request.param
+
+@pytest.fixture(scope="module", params=[0, 2, 4])
+def spin_type(request):
+    return request.param
+
+
+class TestExcitedDMRG:
+    def test_rhf(self, tmp_path, system_def, singlet_embedding_type, spin_type):
+        from pyscf import scf
+
+        mol, ncore, ncas, name = system_def
+        mf = scf.RHF(mol).run(conv_tol=1e-14)
+        if name == "N2":
+            assert abs(mf.e_tot - -107.49650051179789) < 1e-10
+        elif name == "C2":
+            assert abs(mf.e_tot - -75.386902377706) < 1e-10
+        ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(
+            mf, ncore, ncas
+        )
+
+        h1e[np.abs(h1e) < 1e-7] = 0
+        g2e[np.abs(g2e) < 1e-7] = 0
+
+        assert np.linalg.norm(h1e - h1e.transpose(1, 0).conj()) < 1e-7
+        assert np.linalg.norm(g2e - g2e.transpose(2, 3, 0, 1)) < 1e-7
+        assert np.linalg.norm(g2e - g2e.transpose(1, 0, 3, 2).conj()) < 1e-7
+        assert np.linalg.norm(g2e - g2e.transpose(3, 2, 1, 0).conj()) < 1e-7
+
+        driver = DMRGDriver(
+            scratch=str(tmp_path / "nodex"), symm_type=SymmetryTypes.SU2, n_threads=4
+        )
+
+        driver.initialize_system(
+            n_sites=ncas, n_elec=n_elec, spin=spin_type, orb_sym=orb_sym,
+            singlet_embedding=singlet_embedding_type
+        )
+        mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, iprint=1)
+        # fd = driver.write_fcidump(h1e, g2e, ecore=ecore)
+        # mpo = driver.get_conventional_qc_mpo(fd)
+        ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=5)
+        bond_dims = [250] * 8
+        noises = [1e-4] * 4 + [1e-5] * 4 + [0]
+        thrds = [1e-10] * 8
+        energies = driver.dmrg(
+            mpo,
+            ket,
+            n_sweeps=20,
+            bond_dims=bond_dims,
+            noises=noises,
+            thrds=thrds,
+            iprint=2,
+        )
+        if name == "N2":
+            if spin_type == 0:
+                assert abs(energies[0] - -107.65412244752470) < 1e-6
+                assert abs(energies[1] - -106.95962615467998) < 1e-6
+                assert abs(energies[2] - -106.94375693899154) < 1e-6
+            elif spin_type == 2:
+                assert abs(energies[0] - -106.93913285966788) < 1e-6
+                assert abs(energies[1] - -106.71173801496816) < 1e-6
+                assert abs(energies[2] - -106.70055113334190) < 1e-6
+            elif spin_type == 4:
+                assert abs(energies[0] - -107.03144947162717) < 1e-6
+                assert abs(energies[1] - -106.63379058932087) < 1e-6
+                assert abs(energies[2] - -106.62753894625438) < 1e-6
+
+        driver.finalize()

src/core/tensor_functions.hpp

@@ -894,6 +894,10 @@ template <typename S, typename FL> struct TensorFunctions {
             shared_ptr<OpProduct<S, FL>> op =
                 dynamic_pointer_cast<OpProduct<S, FL>>(expr);
             assert(op->b != nullptr);
+            // if (lop.count(op->a) == 0)
+            //     cout << "missing op->a : " << *op->a << endl;
+            // if (rop.count(op->b) == 0)
+            //     cout << "missing op->b : " << *op->b << endl;
             assert(lop.count(op->a) != 0 && rop.count(op->b) != 0);
             shared_ptr<SparseMatrix<S, FL>> lmat = lop.at(op->a);
             shared_ptr<SparseMatrix<S, FL>> rmat = rop.at(op->b);

src/dmrg/mpo.hpp

@@ -1521,7 +1521,7 @@ template <typename S, typename FL> struct IdentityAddedMPO : MPO<S, FL> {
                 }
             if (!found) {
                 x->data.push_back(i_op);
-                y->data.push_back(i_op * i_op);
+                y->data.push_back(m == 0 ? i_op : i_op * i_op);
                 assert(x->get_type() == SymTypes::RVec);
                 x->n = y->n = (int)x->data.size();
                 if (m == 0) {
@@ -1555,7 +1555,10 @@ template <typename S, typename FL> struct IdentityAddedMPO : MPO<S, FL> {
                 }
             if (!found) {
                 x->data.push_back(i_op);
-                y->data.push_back(i_op * i_op);
+                y->data.push_back(m == MPO<S, FL>::right_operator_names.size() -
+                                              1
+                                      ? i_op
+                                      : i_op * i_op);
                 assert(x->get_type() == SymTypes::CVec);
                 x->m = y->m = (int)x->data.size();
                 if (m == mpo->n_sites - 1) {

unit_test/debug_test_dmrg_auto.cpp

@@ -77,6 +77,9 @@ TYPED_TEST(TestDMRG, Test) {
     cout << "rescaled const = " << fcidump->e() << endl;
     cout << "INT end .. T = " << t.get_time() << endl;
 
+    const bool singlet_embedding = false;
+    int se_spin = 2;
+
     Random::rand_seed(1234);
 
     vector<uint8_t> orbsym = fcidump->template orb_sym<uint8_t>();
@@ -89,6 +92,9 @@ TYPED_TEST(TestDMRG, Test) {
     uint8_t isym = PointGroup::swap_pg(pg)(fcidump->isym());
     S target(fcidump->n_elec(), fcidump->twos(),
              S::pg_combine(isym, 0, fcidump->k_mod()));
+    if (singlet_embedding)
+        target = S(fcidump->n_elec() + se_spin, 0,
+                   S::pg_combine(isym, 0, fcidump->k_mod()));
     shared_ptr<HamiltonianQC<S, FL>> hamil =
         make_shared<HamiltonianQC<S, FL>>(vacuum, norb, pg_sym, fcidump);
 
@@ -166,6 +172,8 @@ TYPED_TEST(TestDMRG, Test) {
     // cout << mpo->get_blocking_formulas() << endl;
     // abort();
 
+    mpo = make_shared<IdentityAddedMPO<S, FL>>(mpo);
+
     ubond_t bond_dim = 200;
 
     // MPSInfo
@@ -175,6 +183,10 @@ TYPED_TEST(TestDMRG, Test) {
     // CCSD init
     shared_ptr<MPSInfo<S>> mps_info =
         make_shared<MPSInfo<S>>(norb, vacuum, target, mpo->basis);
+    if (singlet_embedding) {
+        S left_vacuum = S(se_spin, se_spin, 0);
+        mps_info->set_bond_dimension_full_fci(left_vacuum, vacuum);
+    }
     // mps_info->set_bond_dimension_full_fci();
     if (occs.size() == 0)
         mps_info->set_bond_dimension(bond_dim);

unit_test/mpi/test_dmrg_sa_1site_n2_sto3g.cpp

@@ -134,7 +134,7 @@ void TestOneSiteDMRGN2STO3GSA::test_dmrg(
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
-    dmrg->davidson_soft_max_iter = 2000;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     long double energy = dmrg->solve(10, mps->center == 0, 1E-8);

unit_test/mpi/test_dmrg_sa_n2_sto3g.cpp

@@ -136,6 +136,7 @@ void TestDMRGN2STO3GSA::test_dmrg(const vector<S> &targets,
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     long double energy = dmrg->solve(10, mps->center == 0, 1E-8);

unit_test/mpi/test_dmrg_sa_tto_n2_sto3g.cpp

@@ -136,7 +136,7 @@ void TestTTODMRGN2STO3GSA::test_dmrg(
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
-    dmrg->davidson_soft_max_iter = 2000;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     dmrg->solve(tto, mps->center == 0, 0);

unit_test/test_dmrg_sa_1site_n2_sto3g.cpp

@@ -90,7 +90,7 @@ void TestOneSiteDMRGN2STO3GSA<FL>::test_dmrg(
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
-    dmrg->davidson_soft_max_iter = 2000;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     FLL energy = dmrg->solve(10, mps->center == 0, 1E-8);

unit_test/test_dmrg_sa_n2_sto3g.cpp

@@ -90,6 +90,7 @@ void TestDMRGN2STO3GSA<FL>::test_dmrg(
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     dmrg->cutoff = 1E-20;

unit_test/test_dmrg_sa_tto_n2_sto3g.cpp

@@ -91,6 +91,7 @@ void TestTTODMRGN2STO3GSA<FL>::test_dmrg(
     shared_ptr<DMRG<S, FL, FL>> dmrg =
         make_shared<DMRG<S, FL, FL>>(me, bdims, noises);
     dmrg->iprint = 2;
+    dmrg->davidson_soft_max_iter = 500;
     dmrg->noise_type = NoiseTypes::ReducedPerturbativeCollected;
     dmrg->trunc_type = dmrg->trunc_type | TruncationTypes::RealDensityMatrix;
     dmrg->cutoff = 1E-20;