Merge pull request #333 from sjdu10/clean_quimb

add codes of 2-site MPS-MPO fitting for fermionic TN

Author: jcmgray

ci/requirements/py-openblas.yml

@@ -16,3 +16,6 @@ dependencies:
   - pytest-cov
   - scipy<1.15.0
   - tqdm
+  - pip
+  - pip:
+    - git+https://github.com/jcmgray/symmray.git@main

quimb/tensor/fitting.py

@@ -4,6 +4,7 @@
 
 from ..utils import check_opt
 from .contraction import contract_strategy
+from .array_ops import isfermionic
 
 
 def tensor_network_distance(
@@ -94,6 +95,16 @@ def tensor_network_distance(
     if method == "dense":
         tnA = tnA.contract(output_inds=oix, preserve_tensor=True)
         tnB = tnB.contract(output_inds=oix, preserve_tensor=True)
+        if isfermionic(tnA.data):
+            # if fermion tensor, flip dual outer indices in A
+            data = tnA.data
+            dual_outer_axs = tuple(
+                ax
+                for ax, ix in enumerate(tnA.inds)
+                if (ix in oix) and not data.indices[ax].dual
+            )
+            if dual_outer_axs:
+                tnA.modify(data=data.phase_flip(*dual_outer_axs))
 
     # overlap method
     if xAA is None:

quimb/tensor/tensor_1d_compress.py

@@ -23,6 +23,7 @@
 from .tensor_arbgeom import tensor_network_apply_op_vec
 from .tensor_arbgeom_compress import tensor_network_ag_compress
 from .tensor_builder import TN_matching, rand_tensor
+from .array_ops import isfermionic
 from .tensor_core import (
     Tensor,
     TensorNetwork,
@@ -1280,6 +1281,10 @@ def _tn1d_fit_sum_sweep_1site(
     N = len(site_tags)
     K = len(tn_overlaps)
 
+    fermion = isfermionic(tn_fit.tensors[0].data)
+    if fermion:
+        raise NotImplementedError("Fermionic 1-site fitting not implemented, use 2-site (bsz=2).")
+
     if max_bond is not None:
         current_bond_dim = tn_fit.max_bond()
         if current_bond_dim < max_bond:
@@ -1461,7 +1466,9 @@ def _tn1d_fit_sum_sweep_2site(
         left_inds = tuple(ix for ix in tfi0.inds if ix != bond)
         right_inds = tuple(ix for ix in tfi1.inds if ix != bond)
         tfinew = None
-
+        
+        fermion = isfermionic(tfi0.data)
+        
         for k, tn_overlap in enumerate(tn_overlaps):
             # form local overlap
             tnik = (
@@ -1470,6 +1477,14 @@ def _tn1d_fit_sum_sweep_2site(
                 | envs["R", i + 1, k]
             )
 
+            if fermion:
+                # keep track of the environment legs for possible phase flips later
+                left_env_ind, right_env_ind = None, None
+                if type(envs["L", i, k]) is Tensor:
+                    (left_env_ind,) = tfi0.bonds(envs["L", i, k])
+                if type(envs["R", i + 1, k]) is Tensor:
+                    (right_env_ind,) = tfi1.bonds(envs["R", i + 1, k])
+
             # remove old tensors
             del tnik["__FIT__", site0]
             del tnik["__FIT__", site1]
@@ -1479,13 +1494,43 @@ def _tn1d_fit_sum_sweep_2site(
                 all, optimize=optimize, output_inds=left_inds + right_inds
             )
 
+            if fermion:
+                # flip the dual indices of the environment legs if needed
+                lind_id, rind_id = None, None
+                lind_id = tfiknew.inds.index(left_env_ind) if left_env_ind is not None else None
+                rind_id = tfiknew.inds.index(right_env_ind) if right_env_ind is not None else None
+                if lind_id is not None and rind_id is not None:
+                    if tfiknew.data.duals[lind_id]:
+                        tfiknew.data.phase_flip(lind_id, inplace=True)
+                    else:
+                        tfiknew.data.phase_flip(rind_id, inplace=True)
+                elif lind_id is not None and rind_id is None:
+                    if tfiknew.data.duals[lind_id]:
+                        tfiknew.data.phase_flip(lind_id, inplace=True)
+                elif lind_id is None and rind_id is not None:
+                    if tfiknew.data.duals[rind_id]:
+                        tfiknew.data.phase_flip(rind_id, inplace=True)
+
             # sum into fitted tensor
             if tfinew is None:
                 tfinew = tfiknew
             else:
                 tfinew += tfiknew
 
-        tfinew.conj_()
+        if fermion: 
+            # when conjugating tn_fit (a ftn) the dual outer indices will be flipped (phase_dual=True)
+            # but ftensor.conj() assumes phase_dual=False by default
+            # we need to manually flip the dual 'physical' indices in tfinew before conjugating
+            data = tfinew.data
+            dual_outer_axs = tuple(
+                ax
+                for ax, ix in enumerate(tfinew.inds)
+                if ix not in (left_env_ind, right_env_ind) and data.indices[ax].dual # or tfixnew.data.duals[ax]
+            )
+            if dual_outer_axs:
+                tfinew.modify(data=data.phase_flip(*dual_outer_axs))
+        
+        tfinew.conj_() 
 
         tfinew0, tfinew1 = tfinew.split(
             max_bond=max_bond,
@@ -1509,6 +1554,13 @@ def _tn1d_fit_sum_sweep_2site(
         tfi0.modify(data=tfinew0.data, left_inds=tfinew0.left_inds)
         tfi1.modify(data=tfinew1.data, left_inds=tfinew1.left_inds)
 
+        if fermion:
+            # deal with the global signs generated during conjugation
+            for ts in (tfi0|tfi1).tensors:
+                if len(ts.data._oddpos) % 2 == 1:
+                    ts.data.phase_global(inplace=True)
+
+
     return max_tdiff
 
 

tests/test_tensor/test_fmps_fitting.py

@@ -0,0 +1,240 @@
+import importlib
+
+import numpy as np
+import pytest
+
+requires_symmray = pytest.mark.skipif(
+    importlib.util.find_spec("symmray") is None,
+    reason="symmray not installed",
+)
+
+@requires_symmray
+@pytest.fixture
+def get_fpeps_and_norm():
+    import symmray as sr
+    # fPEPS parameters
+    Lx = int(4)
+    Ly = int(4)
+    symmetry = "Z2"
+    D = 4
+    # Load PEPS
+    fpeps = sr.PEPS_fermionic_rand(
+        symmetry, Lx, Ly, bond_dim=D, phys_dim=4, seed=42
+    )
+    fpeps.equalize_norms_()
+    for ts in fpeps.tensors:
+        ts.data.phase_sync(inplace=True)
+    fpeps_norm = fpeps.make_norm()
+    # benchmark_norm = fpeps_norm.contract_boundary_from_xmax(xrange=(0, Lx-1), max_bond=256, cutoff=0.0, mode='direct').contract()
+    benchmark_norm = np.float64(9.347604511732736e18)
+    return fpeps_norm, benchmark_norm
+
+@requires_symmray
+@pytest.mark.parametrize("from_which", ["xmin", "xmax", "ymin", "ymax"])
+def test_fmps_mpo_fitting(from_which, get_fpeps_and_norm):
+    
+    fpeps_norm, benchmark_norm = get_fpeps_and_norm
+    print(f"Benchmark norm: {benchmark_norm}")
+    print("Boundary fMPS-MPO fitting contraction test:")
+    # contraction bond dimension
+    chi = 128
+    if from_which == "xmin":
+        print("xmin:")
+        c_xmin_0 = fpeps_norm.contract_boundary_from_xmin(
+            xrange=(0, 1),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_xmin_1 = fpeps_norm.contract_boundary_from_xmin(
+            xrange=(0, 2),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_xmin_2 = fpeps_norm.contract_boundary_from_xmin(
+            xrange=(0, 1),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        c_xmin_3 = fpeps_norm.contract_boundary_from_xmin(
+            xrange=(0, 2),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        print(c_xmin_0, np.allclose(c_xmin_0, benchmark_norm, rtol=1e-4))
+        print(c_xmin_1, np.allclose(c_xmin_1, benchmark_norm, rtol=1e-4))
+        print(c_xmin_2, np.allclose(c_xmin_2, benchmark_norm, rtol=1e-4))
+        print(c_xmin_3, np.allclose(c_xmin_3, benchmark_norm, rtol=1e-4))
+        assert np.allclose(c_xmin_0, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmin_1, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmin_2, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmin_3, benchmark_norm, rtol=1e-4)
+    elif from_which == "xmax":
+        print("xmax:")
+        c_xmax_0 = fpeps_norm.contract_boundary_from_xmax(
+            xrange=(2, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_xmax_1 = fpeps_norm.contract_boundary_from_xmax(
+            xrange=(1, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_xmax_2 = fpeps_norm.contract_boundary_from_xmax(
+            xrange=(2, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        c_xmax_3 = fpeps_norm.contract_boundary_from_xmax(
+            xrange=(1, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        print(c_xmax_0, np.allclose(c_xmax_0, benchmark_norm, rtol=1e-4))
+        print(c_xmax_1, np.allclose(c_xmax_1, benchmark_norm, rtol=1e-4))
+        print(c_xmax_2, np.allclose(c_xmax_2, benchmark_norm, rtol=1e-4))
+        print(c_xmax_3, np.allclose(c_xmax_3, benchmark_norm, rtol=1e-4))
+        assert np.allclose(c_xmax_0, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmax_1, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmax_2, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_xmax_3, benchmark_norm, rtol=1e-4)
+    elif from_which == "ymin":
+        print("ymin:")
+        c_ymin_0 = fpeps_norm.contract_boundary_from_ymin(
+            yrange=(0, 1),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_ymin_1 = fpeps_norm.contract_boundary_from_ymin(
+            yrange=(0, 2),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_ymin_2 = fpeps_norm.contract_boundary_from_ymin(
+            yrange=(0, 1),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        c_ymin_3 = fpeps_norm.contract_boundary_from_ymin(
+            yrange=(0, 2),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        print(c_ymin_0, np.allclose(c_ymin_0, benchmark_norm, rtol=1e-4))
+        print(c_ymin_1, np.allclose(c_ymin_1, benchmark_norm, rtol=1e-4))
+        print(c_ymin_2, np.allclose(c_ymin_2, benchmark_norm, rtol=1e-4))
+        print(c_ymin_3, np.allclose(c_ymin_3, benchmark_norm, rtol=1e-4))
+        assert np.allclose(c_ymin_0, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymin_1, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymin_2, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymin_3, benchmark_norm, rtol=1e-4)
+    elif from_which == "ymax":
+        print("ymax:")
+        c_ymax_0 = fpeps_norm.contract_boundary_from_ymax(
+            yrange=(2, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_ymax_1 = fpeps_norm.contract_boundary_from_ymax(
+            yrange=(1, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=6,
+        ).contract()
+        c_ymax_2 = fpeps_norm.contract_boundary_from_ymax(
+            yrange=(2, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        c_ymax_3 = fpeps_norm.contract_boundary_from_ymax(
+            yrange=(1, 3),
+            max_bond=chi,
+            cutoff=0.0,
+            mode="fit",
+            tol=1e-5,
+            tn_fit="zipup",
+            bsz=2,
+            max_iterations=5,
+        ).contract()
+        print(c_ymax_0, np.allclose(c_ymax_0, benchmark_norm, rtol=1e-4))
+        print(c_ymax_1, np.allclose(c_ymax_1, benchmark_norm, rtol=1e-4))
+        print(c_ymax_2, np.allclose(c_ymax_2, benchmark_norm, rtol=1e-4))
+        print(c_ymax_3, np.allclose(c_ymax_3, benchmark_norm, rtol=1e-4))
+        assert np.allclose(c_ymax_0, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymax_1, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymax_2, benchmark_norm, rtol=1e-4)
+        assert np.allclose(c_ymax_3, benchmark_norm, rtol=1e-4)