sany: su2 fix npdm & to sz

pyblock2/driver/core.py

@@ -882,12 +882,30 @@ def initialize_system(
                 self.left_vacuum = self.vacuum if left_vacuum is None else left_vacuum
             elif bw.qargs == ("U1Fermi", "U1", "AbelianPG"):
                 self.vacuum = bw.SX(0, 0, 0)
-                self.target = bw.SX(n_elec, spin, pg_irrep)
-                self.left_vacuum = self.vacuum if left_vacuum is None else left_vacuum
+                if left_vacuum is None:
+                    self.target = bw.SX(n_elec, spin, pg_irrep)
+                    self.left_vacuum = (
+                        self.vacuum if left_vacuum is None else left_vacuum
+                    )
+                else:
+                    self.target = bw.SX(
+                        n_elec + left_vacuum.n, spin - left_vacuum.twos, pg_irrep
+                    )
+                    self.left_vacuum = left_vacuum
             elif bw.qargs == ("U1Fermi", "SU2", "SU2", "AbelianPG"):
                 self.vacuum = bw.SX(0, 0, 0, 0)
-                self.target = bw.SX(n_elec, spin, spin, pg_irrep)
-                self.left_vacuum = self.vacuum if left_vacuum is None else left_vacuum
+                if singlet_embedding and left_vacuum is None:
+                    self.target = bw.SX(n_elec + spin % 2, 0, 0, pg_irrep)
+                    self.left_vacuum = bw.SX(spin % 2, spin, spin, 0)
+                elif singlet_embedding and left_vacuum is not None:
+                    assert spin == left_vacuum.twos
+                    self.target = bw.SX(n_elec + left_vacuum.n, 0, 0, pg_irrep)
+                    self.left_vacuum = left_vacuum
+                else:
+                    self.target = bw.SX(n_elec, spin, spin, pg_irrep)
+                    self.left_vacuum = (
+                        self.vacuum if left_vacuum is None else left_vacuum
+                    )
             else:
                 raise RuntimeError("target argument required for custom symmetry.")
         elif target is None:
@@ -984,8 +1002,8 @@ def initialize_system(
             }
             std_ops_su2 = {
                 "": np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),  # identity
-                "C": np.array([[0, 0, 0], [1, 0, 0], [0, -2 ** 0.5, 0]]), # +
-                "D": np.array([[0, 2 ** 0.5, 0], [0, 0, 1], [0, 0, 0]]), # -
+                "C": np.array([[0, 0, 0], [1, 0, 0], [0, -(2**0.5), 0]]),  # +
+                "D": np.array([[0, 2**0.5, 0], [0, 0, 1], [0, 0, 0]]),  # -
             }
             if bw.qargs == ("U1Fermi", "AbelianPG") and "SGF" in bw.hints:
                 site_basis, site_ops = [], []
@@ -2301,7 +2319,12 @@ def get_string_quantum(self, expr, idxs):
                         "C": super_self.bw.SX(1, self.orb_sym[ix]),
                         "D": super_self.bw.SX(-1, -self.orb_sym[ix]),
                     }
-                return sum([qs(0)[""]] + [qs(ix)[ex] for ex, ix in zip(expr, idxs)])
+                from functools import reduce
+
+                return reduce(
+                    lambda a, b: a + b,
+                    [qs(0)[""]] + [qs(ix)[ex] for ex, ix in zip(expr, idxs)],
+                )
 
             def deallocate(self):
                 """Release memory."""
@@ -2433,17 +2456,11 @@ def init_site_ops(self):
                     else:
                         for m in range(self.n_sites):
                             qs = {self.vacuum}
-                            for iter in range(20):
-                                new_qs = set()
-                                for q in qs:
-                                    for k, _ in site_basis[m]:
-                                        new_q = q + k
-                                        for iq in range(new_q.count):
-                                            new_qs.add(new_q[iq])
-                                        new_q = q - k
-                                        for iq in range(new_q.count):
-                                            new_qs.add(new_q[iq])
-                                qs = new_qs
+                            for q, _ in site_basis[m]:
+                                for k, _ in site_basis[m]:
+                                    new_q = q - k
+                                    for iq in range(new_q.count):
+                                        qs.add(new_q[iq])
                             for q in sorted(qs):
                                 mat = super_self.bw.brs.SparseMatrixInfo(i_alloc)
                                 mat.initialize(
@@ -2482,9 +2499,16 @@ def init_site_ops(self):
                                         if dqs is None:
                                             dqs = [xdqs[ix] for ix in range(xdqs.count)]
                                         else:
-                                            dqs = [dq for dq in dqs if any(dq == xdqs[ix] for ix in range(xdqs.count))]
+                                            dqs = [
+                                                dq
+                                                for dq in dqs
+                                                if any(
+                                                    dq == xdqs[ix]
+                                                    for ix in range(xdqs.count)
+                                                )
+                                            ]
                                         blocks.append((q_map[i][1], q_map[j][1], mat))
-                            
+
                             assert dqs is not None
                             assert len(dqs) >= 1
                             dq = dqs[0]
@@ -2506,8 +2530,12 @@ def get_site_string_op(self, m, expr):
                     if SymmetryTypes.SU2 in super_self.symm_type:
                         l = super_self.bw.b.SpinRecoupling.get_level(expr, 0)
                         a = self.get_site_string_op(m, expr[l.left_idx : l.mid_idx - 1])
-                        b = self.get_site_string_op(m, expr[l.mid_idx : l.right_idx - 1])
-                        dq = self.get_su2_string_quantum(expr, [m] * (l.left_cnt + l.right_cnt))
+                        b = self.get_site_string_op(
+                            m, expr[l.mid_idx : l.right_idx - 1]
+                        )
+                        dq = self.get_su2_string_quantum(
+                            expr, [m] * (l.left_cnt + l.right_cnt)
+                        )
                         r = super_self.bw.bs.SparseMatrix(d_alloc)
                         r.allocate(self.find_site_op_info(m, dq))
                         self.opf.product(0, a, b, r)
@@ -2534,13 +2562,15 @@ def init_string_quanta(self, exprs, term_l, left_vacuum):
                             r = super_self.bw.VectorSX([self.vacuum] * (term_l[ix] + 1))
                             r[-1] = self.get_su2_string_quantum(expr, [])
                             lacc = 0
-                            while True: # (.+(.+(.+.)0)0)0
+                            while True:  # (.+(.+(.+.)0)0)0
                                 l = super_self.bw.b.SpinRecoupling.get_level(expr, 0)
                                 if l.right_idx == -1:
                                     break
                                 exprr = expr[l.mid_idx : l.right_idx - 1]
                                 lacc += l.left_cnt
-                                r[lacc] = (r[-1] - self.get_su2_string_quantum(exprr, []))[0]
+                                r[lacc] = (
+                                    r[-1] - self.get_su2_string_quantum(exprr, [])
+                                )[0]
                                 expr = exprr
                             rr.append(r)
                         return rr
@@ -2575,7 +2605,9 @@ def get_string_quanta(self, ref, expr, idxs, k):
                         return ref[k], ref[-1] - ref[k]
                     else:
                         l, r = ref[k], ref[-1] - ref[k]
-                        pexpr = [x for x in expr if not x.isdigit() and x not in '()[]+-,;']
+                        pexpr = [
+                            x for x in expr if not x.isdigit() and x not in "()[]+-,;"
+                        ]
                         for j, (ex, ix) in enumerate(zip(pexpr, idxs)):
                             if ex in orb_dependent_ops:
                                 ipg = self.orb_sym[ix]
@@ -2601,20 +2633,29 @@ def get_su2_string_quantum(self, expr, idxs):
                         else:
                             for m in range(self.n_sites):
                                 if expr in self.site_norm_ops[m]:
-                                    xq = self.site_norm_ops[m][expr].info.delta_quantum[0]
+                                    xq = self.site_norm_ops[m][expr].info.delta_quantum[
+                                        0
+                                    ]
                                     if expr in orb_dependent_ops:
                                         xq.pg = 0
                                     return xq
                     else:
                         l = super_self.bw.b.SpinRecoupling.get_level(expr, 0)
                         qs = super_self.bw.b.SpinRecoupling.get_quanta(expr, l, False)
-                        a = self.get_su2_string_quantum(expr[l.left_idx : l.mid_idx - 1], idxs[:l.left_cnt])
-                        b = self.get_su2_string_quantum(expr[l.mid_idx : l.right_idx - 1], idxs[l.left_cnt:])
+                        a = self.get_su2_string_quantum(
+                            expr[l.left_idx : l.mid_idx - 1], idxs[: l.left_cnt]
+                        )
+                        b = self.get_su2_string_quantum(
+                            expr[l.mid_idx : l.right_idx - 1],
+                            idxs[l.left_cnt : l.left_cnt + l.right_cnt],
+                        )
                         c = a + b
                         nab_idx = c.non_abelian_indices()
                         assert len(qs) == len(nab_idx)
                         for ic in range(c.count):
-                            if all(c[ic].values[ix] == iq for iq, ix in zip(qs, nab_idx)):
+                            if all(
+                                c[ic].values[ix] == iq for iq, ix in zip(qs, nab_idx)
+                            ):
                                 return c[ic]
                         return c[0]
 
@@ -2627,7 +2668,12 @@ def get_string_quantum(self, expr, idxs):
                             k: v.info.delta_quantum
                             for k, v in self.site_norm_ops[ix].items()
                         }
-                        return sum([qs(0)[""]] + [qs(ix)[ex] for ex, ix in zip(expr, idxs)])
+                        from functools import reduce
+
+                        return reduce(
+                            lambda a, b: a + b,
+                            [qs(0)[""]] + [qs(ix)[ex] for ex, ix in zip(expr, idxs)],
+                        )
 
                 def deallocate(self):
                     """Release memory."""
@@ -4421,7 +4467,7 @@ def dmrg(
                 dmrg.davidson_type = getattr(bw.b.DavidsonTypes, dav_type)
         dmrg.davidson_shift = davidson_shift
         if noise_type is None:
-            noise_type = 'ReducedPerturbativeCollected'
+            noise_type = "ReducedPerturbativeCollected"
         dmrg.noise_type = getattr(bw.b.NoiseTypes, noise_type)
         if lowmem_noise:
             dmrg.noise_type = dmrg.noise_type | bw.b.NoiseTypes.LowMem
@@ -6862,11 +6908,11 @@ def mps_change_complex(self, mps, tag):
     def mps_flip_twos(self, mps):
         """
         Flip the sign of projection spin in the MPS.
-        
+
         Args:
             mps : MPS
                 The input MPS.
-        
+
         Returns:
             mps : MPS
                 The output MPS.
@@ -6937,11 +6983,21 @@ def mps_change_symm(self, mps, tag, target):
             self.mpi.barrier()
         mps.info.load_mutable()
         mps.load_mutable()
+        su2_to_sz = (
+            len(mps.info.vacuum.su2_indices()) != 0 and len(target.u1_indices()) != 0
+        )
+        xtarget = target[0]
+        if su2_to_sz:
+            lv = mps.info.left_dims_fci[0].quanta[0]
+            xtarget.n = xtarget.n + lv.n
+            xtarget.twos = 0
         umps = bw.bs.trans_unfused_mps_to_sany(
-            bw.bs.UnfusedMPS(mps), tag, self.ghamil.opf.cg, target
+            bw.bs.UnfusedMPS(mps), tag, self.ghamil.opf.cg, xtarget
         )
         if self.mpi is not None:
             self.mpi.barrier()
+        if su2_to_sz:
+            umps.resolve_singlet_embedding(target.twos)
         zmps = umps.finalize(self.prule)
 
         zmps.info.save_data(self.scratch + "/%s-mps_info.bin" % tag)
@@ -7325,9 +7381,13 @@ def get_spin_projection_mpo(
         it = np.ones((1, 1, 1, 1))
         pympo = None
         for ixw, (xt, wt) in enumerate(zip(xts, wts)):
-            if self.mpi is None or not mpi_split or (
-                mpi_split 
-                and self.mpi.rank == min(ixw, len(wts) - 1 - ixw) % self.mpi.size
+            if (
+                self.mpi is None
+                or not mpi_split
+                or (
+                    mpi_split
+                    and self.mpi.rank == min(ixw, len(wts) - 1 - ixw) % self.mpi.size
+                )
             ):
                 ct = np.cos(xt / 2) * it
                 st, mt = np.sin(xt / 2) * it, -np.sin(xt / 2) * it

src/core/hamiltonian.hpp

@@ -255,9 +255,16 @@ template <typename S, typename FL> struct Hamiltonian {
     shared_ptr<SparseMatrixInfo<S>> find_site_op_info(uint16_t i, S q) const {
         auto p = lower_bound(site_op_infos[i].begin(), site_op_infos[i].end(),
                              q, SparseMatrixInfo<S>::cmp_op_info);
-        if (p == site_op_infos[i].end() || p->first != q)
-            return nullptr;
-        else
+        if (p == site_op_infos[i].end() || p->first != q) {
+            // in general Hamiltonian, allow skipping empty infos
+            shared_ptr<SparseMatrixInfo<S>> info =
+                make_shared<SparseMatrixInfo<S>>(nullptr);
+            info->n = 0;
+            info->delta_quantum = q;
+            info->is_fermion = q.is_fermion();
+            info->is_wavefunction = false;
+            return info;
+        } else
             return p->second;
     }
     // get the delta quantum of an operator string

src/core/operator_functions.hpp

@@ -187,6 +187,7 @@ template <typename S, typename FL> struct OperatorFunctions {
             return;
         S adq = a->info->delta_quantum, cdq = c->info->delta_quantum;
         // possible assert failure here due to the same bra/ket tag
+        // or due to left/right bond dims larger than fci
         assert(adq == cdq && a->info->n >= c->info->n);
         for (int ic = 0, ia = 0; ic < c->info->n; ia++, ic++) {
             while (a->info->quanta[ia] != c->info->quanta[ic])

src/core/sparse_matrix.hpp

@@ -823,6 +823,7 @@ struct SparseMatrixInfo<
     }
     void deallocate() {
         assert(n != -1);
+        assert(alloc != nullptr || n == 0);
         alloc->deallocate((uint32_t *)quanta,
                           n * (sizeof(S) >> 2) + n + _DBL_MEM_SIZE(n));
         alloc = nullptr;

src/core/state_info.hpp

@@ -461,13 +461,11 @@ template <typename S1, typename S2, typename = void, typename = void>
 struct TransStateInfo {
     static shared_ptr<StateInfo<S2>>
     forward(const shared_ptr<StateInfo<S1>> &si, S2 ref) {
-        throw runtime_error("TransStateInfo::forward: not implemented");
-        return nullptr;
+        return TransStateInfo<S2, S1>::backward(si, ref);
     }
     static shared_ptr<StateInfo<S1>>
     backward(const shared_ptr<StateInfo<S2>> &si, S1 ref) {
-        throw runtime_error("TransStateInfo::backward: not implemented");
-        return nullptr;
+        return TransStateInfo<S2, S1>::forward(si, ref);
     }
     static shared_ptr<StateInfo<S2>>
     backward_connection(const shared_ptr<StateInfo<S2>> &si,
@@ -637,9 +635,30 @@ struct TransStateInfo<S, S, typename S::is_sany_t, typename S::is_sany_t> {
         for (int i = 0; i < si->n; i++) {
             S q = si->quanta[i];
             S z = ref;
-            if (q.symm_len() == 3 && q.types[0] == SAnySymmTypes::U1Fermi &&
-                q.types[1] == SAnySymmTypes::U1 &&
-                q.types[2] == SAnySymmTypes::AbelianPG) {
+            if (q.symm_len() == 4 && q.types[0] == SAnySymmTypes::U1Fermi &&
+                q.types[1] == SAnySymmTypes::SU2 &&
+                q.types[2] == SAnySymmTypes::SU2 &&
+                q.types[3] == SAnySymmTypes::AbelianPG) {
+                if (ref.symm_len() == 3 &&
+                    ref.types[0] == SAnySymmTypes::U1Fermi &&
+                    ref.types[1] == SAnySymmTypes::U1 &&
+                    ref.types[2] == SAnySymmTypes::AbelianPG) {
+                    assert(q.values[1] == q.values[2]);
+                    z.values[0] = q.values[0];
+                    z.values[2] = q.values[3];
+                    for (int j = -q.values[2]; j <= q.values[2]; j += 2) {
+                        z.values[1] = j;
+                        mp[z] += si->n_states[i];
+                    }
+                } else
+                    throw runtime_error("TransStateInfo::forward: Unsupported "
+                                        "target symm type: " +
+                                        Parsing::to_string(q) + " -> " +
+                                        Parsing::to_string(ref));
+            } else if (q.symm_len() == 3 &&
+                       q.types[0] == SAnySymmTypes::U1Fermi &&
+                       q.types[1] == SAnySymmTypes::U1 &&
+                       q.types[2] == SAnySymmTypes::AbelianPG) {
                 if (ref.symm_len() == 2 &&
                     ref.types[0] == SAnySymmTypes::U1Fermi &&
                     ref.types[1] == SAnySymmTypes::AbelianPG) {

src/dmrg/mps_unfused.hpp

@@ -782,6 +782,7 @@ template <typename S, typename FL> struct UnfusedMPS {
                     info->left_dims[i]->quanta[j] - lq;
             info->left_dims[i]->sort_states();
         }
+        info->check_bond_dimensions();
         info->save_mutable();
         info->deallocate_mutable();
         shared_ptr<SparseTensor<S, FL>> rst =