gloops, sloops + bp updates

Author: jcmgray

quimb/experimental/belief_propagation/hd1gbp.py

@@ -173,7 +173,10 @@ def iterate(self, tol=5e-6):
                 # check for convergence
                 try:
                     m_old = self.messages[parent, child]
-                    mdiff = self._distance_fn(m.data, m_old.data)
+                    # XXX: need to handle index alignment here to compare
+                    # using _distance_fn:
+                    # mdiff = self._distance_fn(m_old.data, m.data)
+                    mdiff = (m_old - m).norm()
                 except KeyError:
                     mdiff = 1.0
                 max_mdiff = max(mdiff, max_mdiff)
@@ -186,14 +189,17 @@ def iterate(self, tol=5e-6):
         #     note that the raw, undamped `new_messages` are used in the
         #     denominator of the message computations, and so kept 'as is'
         for pair in self.new_messages:
-            if self.damping == 0.0 or pair not in self.messages:
+            if pair not in self.messages:
+                # no old message yet
                 self.messages[pair] = self.new_messages[pair]
             else:
-                self.messages[pair] = self.fn_damping(
+                self.messages[pair] = self._damping_fn(
                     self.messages[pair],
                     self.new_messages[pair],
                 )
 
+        # self.new_messages.clear()
+
         return {
             "nconv": nconv,
             "ncheck": ncheck,

quimb/tensor/belief_propagation/bp_common.py

@@ -129,14 +129,21 @@ def damping(self):
     @damping.setter
     def damping(self, damping):
         if callable(damping):
-            self.fn_damping = self._damping = damping
+            self._damping_fn = self._damping = damping
         else:
             self._damping = damping
 
-            def fn_damping(old, new):
-                return damping * old + (1 - damping) * new
+            if damping == 0.0:
 
-            self.fn_damping = fn_damping
+                def _damping_fn(old, new):
+                    return new
+
+            else:
+
+                def _damping_fn(old, new):
+                    return damping * old + (1 - damping) * new
+
+            self._damping_fn = _damping_fn
 
     @property
     def normalize(self):
@@ -416,6 +423,29 @@ def mdiff(self):
         except IndexError:
             return float("nan")
 
+    def iterate(self, tol=1e-6):
+        """Perform a single iteration of belief propagation. Subclasses should
+        implement this method, returning either `max_mdiff` or a dictionary
+        containing `max_mdiff` and any other relevant information:
+
+            {
+                "nconv": nconv,
+                "ncheck": ncheck,
+                "max_mdiff": max_mdiff,
+            }
+
+        """
+        raise NotImplementedError
+
+    def contract(
+        self,
+        strip_exponent=False,
+        check_zero=True,
+        **kwargs,
+    ):
+        """Contract the tensor network and return the resulting value."""
+        raise NotImplementedError
+
     def __repr__(self):
         return f"{self.__class__.__name__}(n={self.n}, mdiff={self.mdiff:.3g})"
 
@@ -761,3 +791,18 @@ def create_lazy_community_edge_map(tn, site_tags=None, rank_simplify=True):
                 pass
 
     return edges, neighbors, local_tns, touch_map
+
+
+def auto_add_indices(tn, regions):
+    """Make sure all indices incident to any tensor in each region are
+    included in the region.
+    """
+    new_regions = []
+    for r in regions:
+        new_r = set(r)
+        tids = [x for x in new_r if isinstance(x, int)]
+        for tid in tids:
+            t = tn.tensor_map[tid]
+            new_r.update(t.inds)
+        new_regions.append(frozenset(new_r))
+    return new_regions

quimb/tensor/belief_propagation/d1bp.py

@@ -180,7 +180,7 @@ def _update_m(key, new_m):
             mdiff = self._distance_fn(old_m, new_m)
 
             if self.damping:
-                new_m = self.fn_damping(old_m, new_m)
+                new_m = self._damping_fn(old_m, new_m)
 
             # # post-damp distance
             # mdiff = self._distance_fn(old_m, new_m)
@@ -341,7 +341,7 @@ def contract(
 
         return combine_local_contractions(
             zvals,
-            self.backend,
+            backend=self.backend,
             strip_exponent=strip_exponent,
             check_zero=check_zero,
             mantissa=self.sign,
@@ -406,9 +406,9 @@ def contract_with_loops(
             exponent=self.exponent,
         )
 
-    def contract_cluster_expansion(
+    def contract_gloop_expand(
         self,
-        clusters=None,
+        gloops=None,
         autocomplete=True,
         strip_exponent=False,
         check_zero=True,
@@ -417,20 +417,18 @@ def contract_cluster_expansion(
     ):
         from .regions import RegionGraph
 
-        if isinstance(clusters, int):
-            max_cluster_size = clusters
-            clusters = None
+        if isinstance(gloops, int):
+            max_size = gloops
+            gloops = None
         else:
-            max_cluster_size = None
+            max_size = None
 
-        if clusters is None:
-            clusters = tuple(
-                self.tn.gen_regions(max_region_size=max_cluster_size)
-            )
+        if gloops is None:
+            gloops = tuple(self.tn.gen_gloops(max_size=max_size))
         else:
-            clusters = tuple(clusters)
+            gloops = tuple(gloops)
 
-        rg = RegionGraph(clusters, autocomplete=autocomplete)
+        rg = RegionGraph(gloops, autocomplete=autocomplete)
 
         zvals = []
         for r in rg.regions:

quimb/tensor/belief_propagation/d2bp.py

@@ -227,7 +227,7 @@ def _update_m(key, new_m):
             mdiff = self._distance_fn(old_m, new_m)
 
             if self.damping:
-                new_m = self.fn_damping(old_m, new_m)
+                new_m = self._damping_fn(old_m, new_m)
 
             # # post-damp distance
             # mdiff = self._distance_fn(old_m, new_m)
@@ -381,9 +381,9 @@ def contract(
             check_zero=check_zero,
         )
 
-    def contract_cluster_expansion(
+    def contract_gloop_expand(
         self,
-        clusters=None,
+        gloops=None,
         autocomplete=True,
         optimize="auto-hq",
         strip_exponent=False,
@@ -394,20 +394,20 @@ def contract_cluster_expansion(
     ):
         self.normalize_message_pairs()
 
-        if isinstance(clusters, int):
-            max_cluster_size = clusters
-            clusters = None
+        if isinstance(gloops, int):
+            max_size = gloops
+            gloops = None
         else:
-            max_cluster_size = None
+            max_size = None
 
-        if clusters is None:
-            clusters = tuple(
-                self.tn.gen_regions(max_region_size=max_cluster_size)
+        if gloops is None:
+            gloops = tuple(
+                self.tn.gen_gloops(max_size=max_size)
             )
         else:
-            clusters = tuple(clusters)
+            gloops = tuple(gloops)
 
-        rg = RegionGraph(clusters, autocomplete=autocomplete)
+        rg = RegionGraph(gloops, autocomplete=autocomplete)
 
         for tid in self.tn.tensor_map:
             rg.add_region([tid])

quimb/tensor/belief_propagation/hd1bp.py

@@ -10,10 +10,12 @@
 """
 
 import autoray as ar
+
 import quimb.tensor as qtn
 
 from .bp_common import (
     BeliefPropagationCommon,
+    combine_local_contractions,
     compute_all_index_marginals_from_messages,
     contract_hyper_messages,
     initialize_hyper_messages,
@@ -266,7 +268,7 @@ def _normalize_and_insert(key, new_m, max_mdiff):
             mdiff = self._distance_fn(old_m, new_m)
 
             if self.damping:
-                new_m = self.fn_damping(old_m, new_m)
+                new_m = self._damping_fn(old_m, new_m)
 
             # # post-damp distance
             # mdiff = self._distance_fn(old_m, new_m)
@@ -341,6 +343,101 @@ def contract(self, strip_exponent=False, check_zero=True):
             backend=self.backend,
         )
 
+    def normalize_messages(self):
+        """Normalize all messages such that the 'region contraction' of a
+        single hyper index is 1.
+        """
+        for ind, tids in self.tn.ind_map.items():
+            ms = [self.messages[tid, ind] for tid in tids]
+            overlap = qtn.array_contract(ms, [(0,) for _ in ms], [])
+            overlap **= 1 / len(ms)
+            for tid, m in zip(tids, ms):
+                self.messages[tid, ind] = m / overlap
+
+    def get_cluster(self, r, virtual=True, autocomplete=True):
+        """Get the tensor network of a region ``r``, with all boundary
+        messages attached.
+
+        Parameters
+        ----------
+        r : sequence of int or str
+            The region to get, given as a sequence of indices or tensor ids.
+        virtual : bool, optional
+            Whether the view the original tensors (`virtual=True`, the default)
+            or take copies (`virtual=False`).
+        autocomplete : bool, optional
+            Whether to automatically include all indices attached to the
+            tensors in the region, or just the ones given in ``r``.
+
+        Returns
+        -------
+        TensorNetwork
+        """
+        rtids = set()
+        rinds = set()
+        for x in r:
+            if isinstance(x, int):
+                rtids.add(x)
+                if autocomplete:
+                    rinds.update(self.tn.tensor_map[x].inds)
+            else:
+                rinds.add(x)
+
+        tnr = self.tn._select_tids(rtids, virtual=virtual)
+        for ind in rinds:
+            # attach all messages coming from tensors outside the cluster
+            for ntid in self.tn.ind_map[ind]:
+                if ntid not in rtids:
+                    tnr |= qtn.Tensor(
+                        data=self.messages[ntid, ind], inds=(ind,)
+                    )
+
+        return tnr
+
+    def contract_gloop_expand(
+        self,
+        gloops=None,
+        strip_exponent=False,
+        check_zero=True,
+        optimize="auto-hq",
+        progbar=False,
+        **contract_otps,
+    ):
+        from .regions import RegionGraph
+
+        # if we normalized messages we can ignore all index-only regions
+        self.normalize_messages()
+
+        if isinstance(gloops, int):
+            gloops = tuple(self.tn.gen_gloops(gloops))
+
+        rg = RegionGraph(gloops)
+
+        if progbar:
+            import tqdm
+
+            regions = tqdm.tqdm(rg.regions)
+        else:
+            regions = rg.regions
+
+        zvals = []
+        for r in regions:
+            # XXX: autoreduce intersecting clusters to gloops?
+            cr = rg.get_count(r)
+            # either we autocomplete above or we do it here per region
+            tnr = self.get_cluster(r, autocomplete=True)
+            zr = tnr.contract(
+                output_inds=(), optimize=optimize, **contract_otps
+            )
+            zvals.append((zr, cr))
+
+        return combine_local_contractions(
+            zvals,
+            strip_exponent=strip_exponent,
+            check_zero=check_zero,
+            backend=self.backend,
+        )
+
 
 def contract_hd1bp(
     tn,

quimb/tensor/belief_propagation/l1bp.py

@@ -178,7 +178,7 @@ def _update_m(key, data):
             mdiff = self._distance_fn(data, tm.data)
 
             if self.damping:
-                data = self.fn_damping(data, tm.data)
+                data = self._damping_fn(data, tm.data)
 
             # # post-damp distance
             # mdiff = self._distance_fn(data, tm.data)

quimb/tensor/belief_propagation/l2bp.py

@@ -246,7 +246,7 @@ def _update_m(key, data):
             mdiff = self._distance_fn(data, tm.data)
 
             if self.damping:
-                data = self.fn_damping(data, tm.data)
+                data = self._damping_fn(data, tm.data)
 
             # # post-damp distance
             # mdiff = self._distance_fn(data, tm.data)