From 09adc142d03a5126c20dede14169ff0da49f0a09 Mon Sep 17 00:00:00 2001
From: Johnnie Gray <johnniemcgray@gmail.com>
Date: Fri, 31 Jan 2025 13:17:14 -0800
Subject: [PATCH] TN: move .bond and .bond_size into arbgeom

---
 quimb/tensor/tensor_2d.py      | 14 -------------
 quimb/tensor/tensor_3d.py      | 12 -----------
 quimb/tensor/tensor_arbgeom.py | 37 ++++++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py
index cba0841c..bf7fa1a1 100644
--- 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
index 2de3743a..0f4c658a 100644
--- 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
index 01c4d4fb..16d82c54 100644
--- 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()