TN: move .bond and .bond_size into arbgeom

jcmgray · Jan 31, 2025 · 09adc14 · 09adc14
1 parent b6259ba
commit 09adc14
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 26 deletions.
diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py
@@ -4696,20 +4696,6 @@ class TensorNetwork2DFlat(TensorNetwork2D):
         "_Ly",
     )
 
-    def bond(self, coo1, coo2):
-        """Get the name of the index defining the bond between sites at
-        ``coo1`` and ``coo2``.
-        """
-        (b_ix,) = self[coo1].bonds(self[coo2])
-        return b_ix
-
-    def bond_size(self, coo1, coo2):
-        """Return the (combined) size of the bond(s) between sites at ``coo1``
-        and ``coo2``.
-        """
-        b_ix = self.bond(coo1, coo2)
-        return self[coo1].ind_size(b_ix)
-
     def expand_bond_dimension(
         self, new_bond_dim, inplace=True, bra=None, rand_strength=0.0
     ):

diff --git a/quimb/tensor/tensor_3d.py b/quimb/tensor/tensor_3d.py
@@ -2228,18 +2228,6 @@ class TensorNetwork3DFlat(TensorNetwork3D):
         "_Lz",
     )
 
-    def bond(self, coo1, coo2):
-        """Get the name of the index defining the bond between sites at
-        ``coo1`` and ``coo2``.
-        """
-        (b_ix,) = self[coo1].bonds(self[coo2])
-        return b_ix
-
-    def bond_size(self, coo1, coo2):
-        """Return the size of the bond between sites at ``coo1`` and ``coo2``."""
-        b_ix = self.bond(coo1, coo2)
-        return self[coo1].ind_size(b_ix)
-
 
 class PEPS3D(TensorNetwork3DVector, TensorNetwork3DFlat):
     r"""Projected Entangled Pair States object (3D).

diff --git a/quimb/tensor/tensor_arbgeom.py b/quimb/tensor/tensor_arbgeom.py
@@ -617,6 +617,43 @@ def _get_tid_to_site_map(self):
             tid2site[tid] = site
         return tid2site
 
+    def bond(self, coo1, coo2):
+        """Get the name of the index defining the bond between sites at
+        ``coo1`` and ``coo2``. This will error if there is not exactly one bond
+        between the sites.
+
+        Parameters
+        ----------
+        coo1 : hashable or str
+            The first site, or site tag.
+        coo2 : hashable or str
+            The second site, or site tag.
+
+        Returns
+        -------
+        str
+        """
+        (b_ix,) = self[coo1].bonds(self[coo2])
+        return b_ix
+
+    def bond_size(self, coo1, coo2):
+        """Return the (combined) size of the bond(s) between sites at ``coo1``
+        and ``coo2``.
+
+        Parameters
+        ----------
+        coo1 : hashable or str
+            The first site, or site tag.
+        coo2 : hashable or str
+            The second site, or site tag.
+
+        Returns
+        -------
+        int
+        """
+        b_ix = self.bond(coo1, coo2)
+        return self[coo1].ind_size(b_ix)
+
     def gen_bond_coos(self):
         """Generate the coordinates (pairs of sites) of all bonds."""
         tid2site = self._get_tid_to_site_map()