npdm index masks

Author: hczhai

docs/source/theory/spin-orbital.rst

@@ -169,7 +169,7 @@ First consider the second and third terms
       + \frac{1}{2} \sum_{ij\in L, kl\in R} v_{klij} a_i^\dagger a_j a_k^\dagger a_l \\
     =&\ \sum_{ij\in L, kl\in R} v_{ijkl} a_i^\dagger a_j a_k^\dagger a_l
      = \sum_{ij\in L} a_i^\dagger a_j \sum_{kl\in R} v_{ijkl} a_k^\dagger a_l
-     = \sum_{ij\in L} \hat{B}_{ij} \hat{Q}_{ij\prime}^{R}
+     = \sum_{ij\in L} \hat{B}_{ij} \hat{Q}_{ij}^{R\prime}
 
 For the other two terms,
 
@@ -182,12 +182,12 @@ For the other two terms,
         -\frac{1}{2} \sum_{il\in L, jk\in R} v_{klij} a_i^\dagger a_l a_k^\dagger a_j \\
     =&\ -\sum_{il\in L, jk\in R} v_{ijkl} a_i^\dagger a_l a_k^\dagger a_j \\
     =&\ -\sum_{il\in L} a_i^\dagger a_l \sum_{jk\in R} v_{ijkl} a_k^\dagger a_j
-    =  \sum_{il\in L} \hat{B}_{il} \hat{Q}_{il\prime\prime}^{R}
+    =  \sum_{il\in L} \hat{B}_{il} \hat{Q}_{il}^{R\prime\prime}
 
 Then
 
 .. math::
-     \hat{Q}_{ij}^{R} =  \hat{Q}_{ij\prime}^{R} + \hat{Q}_{ij\prime\prime}^{R}
+     \hat{Q}_{ij}^{R} =  \hat{Q}_{ij}^{R\prime} + \hat{Q}_{ij}^{R\prime\prime}
     = \sum_{kl\in R} \big( v_{ijkl} - v_{ilkj} \big) a_k^\dagger a_l
 
 Normal/Complementary Partitioning

docs/source/theory/su2.rst

@@ -315,7 +315,7 @@ we have
     + a_{p\beta}^\dagger a_{q\beta}^\dagger a_{r\beta} a_{s\beta}
     + a_{p\alpha}^\dagger a_{q\beta}^\dagger a_{r\alpha} a_{s\beta}
     + a_{p\beta}^\dagger a_{q\alpha}^\dagger a_{r\beta} a_{s\alpha} \\
-    =&\ \sum_{\sigma\sigma'} a_{p\sigma}^\dagger a_{q\sigma'}^\dagger a_{r\sigma} s_{s\sigma'}
+    =&\ \sum_{\sigma\sigma'} a_{p\sigma}^\dagger a_{q\sigma'}^\dagger a_{r\sigma} a_{s\sigma'}
 
 Another case
 

docs/source/tutorial/hubbard.ipynb

@@ -897,7 +897,7 @@
     "$$\n",
     "\n",
     "$$\n",
-    "    -t \\sum_{\\langle i,j \\rangle, \\sigma} \\left( c_{i,\\sigma}^\\dagger c_{j,\\sigma} + \\mathrm{H.c.} \\right)\n",
+    "    -t \\sum_{\\langle i,j \\rangle, \\sigma} \\left( a_{i,\\sigma}^\\dagger a_{j,\\sigma} + \\mathrm{H.c.} \\right)\n",
     "    = -\\sqrt{2} t \\sum_{\\langle i,j \\rangle} \\left[\n",
     "        \\big( E_i \\big)^{[1/2]} \\otimes_{[0]} \\big( F_{j} \\big)^{[1/2]} + \\big( E_{j} \\big)^{[1/2]} \\otimes_{[0]} \\big( F_i \\big)^{[1/2]}\n",
     "        \\right]\n",

pyblock2/driver/core.py

@@ -5498,6 +5498,7 @@ def get_npdm(
         algo_type=None,
         npdm_expr=None,
         mask=None,
+        index_masks=None,
         simulated_parallel=0,
         fused_contraction_rotation=True,
         cutoff=1e-24,
@@ -5537,6 +5538,9 @@ def get_npdm(
             mask : None or list[int] or list[list[int]]
                 The mask for setting repeated indices for the operator expression.
                 Default is None, meaning that all indices can be different.
+            index_masks : None or list[list[int]] or list[list[list[int]]]
+                The list of allowed site indices for each operator in the operator expression.
+                Default is None, meaning that all site indices are allowed.
             simulated_parallel : int
                 Number of processors for simulating parallel algorithm serially.
                 Default is zero, meaning that the serial algorithm is used if
@@ -5629,7 +5633,7 @@ def get_npdm(
                 for _ in range(pdm_type - 1):
                     op_str = su2_coupling % op_str
             perm = bw.b.SpinPermScheme.initialize_su2(
-                pdm_type * 2, op_str, True,
+                int(pdm_type * 2), op_str, True,
                 mask=bw.b.VectorUInt16() if mask is None else bw.b.VectorUInt16(mask),
                 max_n_sites=ket.n_sites,
             )
@@ -5646,41 +5650,33 @@ def get_npdm(
             if mask is None:
                 perms = bw.b.VectorSpinPermScheme(
                     [
-                        bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True,
+                        bw.b.SpinPermScheme.initialize_sz(len(cd), cd, True,
                             mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) if fermionic_ops is None else
-                        bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops,
+                        bw.b.SpinPermScheme.initialize_sany(len(cd), cd, fermionic_ops,
                             mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) for cd in op_str
                     ]
                 )
             elif len(mask) != 0 and not isinstance(mask[0], int):
                 assert len(mask) == len(op_str)
-                pts = (
-                    [pdm_type] * len(op_str) if isinstance(pdm_type, int) else pdm_type
-                )
                 perms = bw.b.VectorSpinPermScheme(
                     [
                         bw.b.SpinPermScheme.initialize_sz(
-                            pt * 2, cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
+                            len(cd), cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
                         ) if fermionic_ops is None else
                         bw.b.SpinPermScheme.initialize_sany(
-                            pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
-                        )
-                        for cd, xm, pt in zip(op_str, mask, pts)
+                            len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
+                        ) for cd, xm in zip(op_str, mask)
                     ]
                 )
             else:
-                pts = (
-                    [pdm_type] * len(op_str) if isinstance(pdm_type, int) else pdm_type
-                )
                 perms = bw.b.VectorSpinPermScheme(
                     [
                         bw.b.SpinPermScheme.initialize_sz(
-                            pt * 2, cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
+                            len(cd), cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
                         ) if fermionic_ops is None else
                         bw.b.SpinPermScheme.initialize_sany(
-                            pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
-                        )
-                        for cd, pt in zip(op_str, pts)
+                            len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
+                        ) for cd in op_str
                     ]
                 )
         elif SymmetryTypes.SGF in bw.symm_type:
@@ -5693,41 +5689,34 @@ def get_npdm(
             if mask is None:
                 perms = bw.b.VectorSpinPermScheme(
                     [
-                        bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True,
+                        bw.b.SpinPermScheme.initialize_sz(len(cd), cd, True,
                             mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) if fermionic_ops is None else
-                        bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops,
+                        bw.b.SpinPermScheme.initialize_sany(len(cd), cd, fermionic_ops,
                             mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) for cd in op_str
                     ]
                 )
             elif len(mask) != 0 and not isinstance(mask[0], int):
                 assert len(mask) == len(op_str)
-                pts = (
-                    [pdm_type] * len(op_str) if isinstance(pdm_type, int) else pdm_type
-                )
                 perms = bw.b.VectorSpinPermScheme(
                     [
                         bw.b.SpinPermScheme.initialize_sz(
-                            pt * 2, cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
+                            len(cd), cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
                         ) if fermionic_ops is None else
                         bw.b.SpinPermScheme.initialize_sany(
-                            pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
+                            len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
                         )
-                        for cd, xm, pt in zip(op_str, mask, pts)
+                        for cd, xm in zip(op_str, mask)
                     ]
                 )
             else:
-                pts = (
-                    [pdm_type] * len(op_str) if isinstance(pdm_type, int) else pdm_type
-                )
                 perms = bw.b.VectorSpinPermScheme(
                     [
                         bw.b.SpinPermScheme.initialize_sz(
-                            pt * 2, cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
+                            len(cd), cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
                         ) if fermionic_ops is None else
                         bw.b.SpinPermScheme.initialize_sany(
-                            pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
-                        )
-                        for cd, pt in zip(op_str, pts)
+                            len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
+                        ) for cd in op_str
                     ]
                 )
         elif SymmetryTypes.SAny in bw.symm_type:
@@ -5761,6 +5750,19 @@ def get_npdm(
                     ]
                 )
 
+        if index_masks is not None:
+            if any(any(isinstance(im, (int, np.int64)) for im in ims) for ims in index_masks if len(ims) != 0):
+                index_masks = [index_masks]
+            if len(index_masks) == 1:
+                index_masks = index_masks * len(perms)
+            assert len(perms) == len(index_masks)
+            for perm, ims in zip(perms, index_masks):
+                f = lambda x: [int(px) for px in x]
+                if self.reorder_idx is not None:
+                    rev_idx = np.argsort(self.reorder_idx)
+                    f = lambda x: [rev_idx[int(px)] for px in x]
+                perm.index_mask = bw.b.VectorVectorUInt16([bw.b.VectorUInt16(f(x)) for x in ims])
+
         if iprint >= 1:
             print("npdm string =", op_str)
 
@@ -5949,7 +5951,7 @@ def get_npdm(
             for ip in range(len(npdms)):
                 npdms[ip] = np.asarray(npdms[ip])
 
-        if self.reorder_idx is not None:
+        if self.reorder_idx is not None and index_masks is None:
             rev_idx = np.argsort(self.reorder_idx)
             for ip in range(len(npdms)):
                 for i in range(npdms[ip].ndim):