more support for PBC and AG operations like summing

jcmgray · Apr 6, 2024 · 197c2e1 · 197c2e1
1 parent 8cbcc38
commit 197c2e1
Show file tree

Hide file tree

Showing 11 changed files with 375 additions and 466 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -72,6 +72,10 @@ Release notes for `quimb`.
   method for applying 'operator' TNs to the lower indices of 'operator' TNs like $B \rightarrow B A$.
 - add [`tnop.gate_sandwich_with_op_lazy`](quimb.tensor.tensor_arbgeom.TensorNetworkGenOperator.gate_sandwich_with_op_lazy)
   method for applying 'operator' TNs to the upper and lower indices of 'operator' TNs like $B \rightarrow A B A^\dagger$.
+- unify all TN summing routines into
+  [`tensor_network_ag_sum](quimb.tensor.tensor_arbgeom.tensor_network_ag_sum),
+  which allows summing any two tensor networks with matching site tags and
+  outer indices, replacing specific MPS, MPO, PEPS, PEPO, etc. summing routines.
 - add [`rand_symmetric_array`](quimb.tensor.tensor_builder.rand_symmetric_array),
   [`rand_tensor_symmetric`](quimb.tensor.tensor_builder.rand_tensor_symmetric)
   [`TN2D_rand_symmetric`](quimb.tensor.tensor_builder.TN2D_rand_symmetric)

diff --git a/quimb/tensor/__init__.py b/quimb/tensor/__init__.py
@@ -129,7 +129,7 @@
     NNI_ham_XY,
     SpinHam,
     SpinHam1D,
-    TN1D_matching,
+    TN_matching,
     TN2D_classical_ising_partition_function,
     TN2D_corner_double_line,
     TN2D_embedded_classical_ising_partition_function,
@@ -363,7 +363,7 @@
     "TN_rand_from_edges",
     "TN_rand_reg",
     "TN_rand_tree",
-    "TN1D_matching",
+    "TN_matching",
     "TN2D_classical_ising_partition_function",
     "TN2D_corner_double_line",
     "TN2D_embedded_classical_ising_partition_function",

diff --git a/quimb/tensor/tensor_1d.py b/quimb/tensor/tensor_1d.py
@@ -1,35 +1,37 @@
 """Classes and algorithms related to 1D tensor networks."""
 
-import operator
+import collections
 import functools
 import itertools
-import collections
+import operator
 from math import log2
 from numbers import Integral
 
 import scipy.sparse.linalg as spla
-from autoray import do, dag, reshape, conj, get_dtype_name, transpose
+from autoray import conj, dag, do, get_dtype_name, reshape, transpose
 
-from ..utils import print_multi_line, ensure_dict, partition_all, deprecated
 import quimb as qu
+
+from ..linalg.base_linalg import norm_trace_dense
+from ..utils import deprecated, ensure_dict, partition_all, print_multi_line
+from . import array_ops as ops
+from .tensor_arbgeom import (
+    TensorNetworkGen,
+    TensorNetworkGenOperator,
+    TensorNetworkGenVector,
+    tensor_network_ag_sum,
+    tensor_network_align,
+    tensor_network_apply_op_op,
+    tensor_network_apply_op_vec,
+)
 from .tensor_core import (
     Tensor,
     TensorNetwork,
-    rand_uuid,
     bonds,
     oset,
+    rand_uuid,
     tags_to_oset,
 )
-from .tensor_arbgeom import (
-    TensorNetworkGen,
-    TensorNetworkGenVector,
-    TensorNetworkGenOperator,
-    tensor_network_align,
-    tensor_network_apply_op_vec,
-)
-from ..linalg.base_linalg import norm_trace_dense
-from . import array_ops as ops
-
 
 align_TN_1D = deprecated(
     tensor_network_align, "align_TN_1D", "tensor_network_align"
@@ -1580,36 +1582,9 @@ def gen_tensors():
             site_tag_id=site_tag_id,
         )
 
-    def add_MPS(self, other, inplace=False, compress=False, **compress_opts):
+    def add_MPS(self, other, inplace=False, **kwargs):
         """Add another MatrixProductState to this one."""
-        if self.L != other.L:
-            raise ValueError("Can't add MPS with another of different length.")
-
-        new_mps = self if inplace else self.copy()
-
-        for i in new_mps.gen_sites_present():
-            t1, t2 = new_mps[i], other[i]
-
-            if set(t1.inds) != set(t2.inds):
-                # Need to use bonds to match indices
-                reindex_map = {}
-
-                if i > 0 or self.cyclic:
-                    pair = ((i - 1) % self.L, i)
-                    reindex_map[other.bond(*pair)] = new_mps.bond(*pair)
-
-                if i < new_mps.L - 1 or self.cyclic:
-                    pair = (i, (i + 1) % self.L)
-                    reindex_map[other.bond(*pair)] = new_mps.bond(*pair)
-
-                t2 = t2.reindex(reindex_map)
-
-            t1.direct_product_(t2, sum_inds=new_mps.site_ind(i))
-
-        if compress:
-            new_mps.compress(**compress_opts)
-
-        return new_mps
+        return tensor_network_ag_sum(self, other, inplace=inplace, **kwargs)
 
     add_MPS_ = functools.partialmethod(add_MPS, inplace=True)
 
@@ -1634,22 +1609,6 @@ def permute_arrays(self, shape="lrp"):
             inds = [inds[s] for s in shape if s in inds]
             self[i].transpose_(*inds)
 
-    def __add__(self, other):
-        """MPS addition."""
-        return self.add_MPS(other, inplace=False)
-
-    def __iadd__(self, other):
-        """In-place MPS addition."""
-        return self.add_MPS(other, inplace=True)
-
-    def __sub__(self, other):
-        """MPS subtraction."""
-        return self.add_MPS(other * -1, inplace=False)
-
-    def __isub__(self, other):
-        """In-place MPS subtraction."""
-        return self.add_MPS(other * -1, inplace=True)
-
     def normalize(self, bra=None, eps=1e-15, insert=None):
         """Normalize this MPS, optional with co-vector ``bra``. For periodic
         MPS this uses transfer matrix SVD approximation with precision ``eps``
@@ -2931,40 +2890,8 @@ def from_fill_fn(
 
         return mpo
 
-    def add_MPO(self, other, inplace=False, compress=False, **compress_opts):
-        """Add another MatrixProductState to this one."""
-        if self.L != other.L:
-            raise ValueError(
-                "Can't add MPO with another of different length."
-                f"Got lengths {self.L} and {other.L}"
-            )
-
-        summed = self if inplace else self.copy()
-
-        for i in summed.gen_sites_present():
-            t1, t2 = summed[i], other[i]
-
-            if set(t1.inds) != set(t2.inds):
-                # Need to use bonds to match indices
-                reindex_map = {}
-
-                if i > 0 or self.cyclic:
-                    pair = ((i - 1) % self.L, i)
-                    reindex_map[other.bond(*pair)] = summed.bond(*pair)
-
-                if i < summed.L - 1 or self.cyclic:
-                    pair = (i, (i + 1) % self.L)
-                    reindex_map[other.bond(*pair)] = summed.bond(*pair)
-
-                t2 = t2.reindex(reindex_map)
-
-            sum_inds = (summed.upper_ind(i), summed.lower_ind(i))
-            t1.direct_product_(t2, sum_inds=sum_inds)
-
-        if compress:
-            summed.compress(**compress_opts)
-
-        return summed
+    def add_MPO(self, other, inplace=False, **kwargs):
+        return tensor_network_ag_sum(self, other, inplace=inplace, **kwargs)
 
     add_MPO_ = functools.partialmethod(add_MPO, inplace=True)
 
@@ -2979,35 +2906,17 @@ def _apply_mps(
             **compress_opts,
         )
 
-    def _apply_mpo(self, other, compress=False, **compress_opts):
-        A, B = self.copy(), other.copy()
-
-        # align the indices and combine into a ladder
-        A.upper_ind_id = B.upper_ind_id
-        B.upper_ind_id = "__tmp{}__"
-        A.lower_ind_id = "__tmp{}__"
-        AB = A | B
-
-        # contract each pair of tensors at each site
-        for i in range(A.L):
-            AB ^= A.site_tag(i)
-
-        # convert back to MPO and fuse the double bonds
-        AB.view_as_(
-            MatrixProductOperator,
-            upper_ind_id=A.upper_ind_id,
-            lower_ind_id=B.lower_ind_id,
-            cyclic=self.cyclic,
+    def _apply_mpo(
+        self, other, compress=False, contract=True, **compress_opts
+    ):
+        return tensor_network_apply_op_op(
+            A=self,
+            B=other,
+            contract=contract,
+            compress=compress,
+            **compress_opts,
         )
 
-        AB.fuse_multibonds_()
-
-        # optionally compress
-        if compress:
-            AB.compress(**compress_opts)
-
-        return AB
-
     def apply(self, other, compress=False, **compress_opts):
         r"""Act with this MPO on another MPO or MPS, such that the resulting
         object has the same tensor network structure/indices as ``other``.
@@ -3124,22 +3033,6 @@ def partial_transpose(self, sysa, inplace=False):
         tn.reindex_({tmp_ind_id.format(i): tn.lower_ind(i) for i in sysa})
         return tn
 
-    def __add__(self, other):
-        """MPO addition."""
-        return self.add_MPO(other, inplace=False)
-
-    def __iadd__(self, other):
-        """In-place MPO addition."""
-        return self.add_MPO(other, inplace=True)
-
-    def __sub__(self, other):
-        """MPO subtraction."""
-        return self.add_MPO(-1 * other, inplace=False)
-
-    def __isub__(self, other):
-        """In-place MPO subtraction."""
-        return self.add_MPO(-1 * other, inplace=True)
-
     def rand_state(self, bond_dim, **mps_opts):
         """Get a random vector matching this MPO."""
         return qu.tensor.MPS_rand_state(