tn.gauge_all_simple: add info option

docs/changelog.md

@@ -9,19 +9,16 @@ Release notes for `quimb`.
 
 - move belief propagation to `quimb.tensor.belief_propagation`
 
-**Breaking Changes**
-
-- move belief propagation to `quimb.tensor.belief_propagation`
-
 **Enhancements:**
 
 - [`MatrixProductState.measure`](quimb.tensor.tensor_1d.MatrixProductState.measure): add a `seed` kwarg
 - belief propagation, implement DIIS (direct inversion in the iterative subspace)
 - belief propagation, unify various aspects such as message normalization and distance.
-- belief propagation, add a `plot` method.
+- belief propagation, add a [`plot`](quimb.tensor.belief_propagation.bp_common.BeliefPropagationCommon.plot) method.
 - belief propagation, add a `contract_every` option.
 - HV1BP: vectorize both contraction and message initialization
-- add `qu.plot_multi_series_zoom` for plotting multiple series with a zoomed inset, useful for various convergence plots such as BP
+- add [`qu.plot_multi_series_zoom`](quimb.utils_plot.plot_multi_series_zoom) for plotting multiple series with a zoomed inset, useful for various convergence plots such as BP
+- add `info` option to [`tn.gauge_all_simple`](quimb.tensor.tensor_core.TensorNetwork.gauge_all_simple) for tracking extra information such as number of iterations and max gauge difffi
 
 **Bug fixes:**
 

quimb/tensor/tensor_core.py

@@ -7029,6 +7029,7 @@ def gauge_all_simple(
         gauges=None,
         equalize_norms=False,
         touched_tids=None,
+        info=None,
         progbar=False,
         inplace=False,
     ):
@@ -7058,6 +7059,13 @@ def gauge_all_simple(
             Whether to equalize the norms of the tensors after each update.
         touched_tids : sequence of int, optional
             The tensor identifiers to start the gauge sweep from.
+        info : dict, optional
+            Store extra information about the gauging process in this dict.
+            If supplied, the following keys are filled:
+
+            - 'iterations': the number of iterations performed.
+            - 'max_sdiff': the maximum singular value difference.
+
         progbar : bool, optional
             Whether to show a progress bar.
         inplace : bool, optional
@@ -7093,7 +7101,7 @@ def gauge_all_simple(
         }[(power == 1.0, smudge == 0.0)]
 
         # for retrieving singular values
-        info = {}
+        sub_info = {}
 
         # accrue scaling to avoid numerical blow-ups
         nfact = 0.0
@@ -7176,10 +7184,10 @@ def gauge_all_simple(
 
                 # perform SVD to get new bond gauge
                 tensor_compress_bond(
-                    t1, t2, absorb=None, info=info, cutoff=0.0
+                    t1, t2, absorb=None, info=sub_info, cutoff=0.0
                 )
 
-                s = info["singular_values"]
+                s = sub_info["singular_values"]
                 snorm = do("linalg.norm", s)
                 new_gauge = s / snorm
                 nfact = do("log10", snorm) + nfact
@@ -7249,6 +7257,10 @@ def gauge_all_simple(
                 t1.multiply_index_diagonal_(ix, s_1_2)
                 t2.multiply_index_diagonal_(ix, s_1_2)
 
+        if info is not None:
+            info["iterations"] = it
+            info["max_sdiff"] = max_sdiff
+
         return tn
 
     gauge_all_simple_ = functools.partialmethod(gauge_all_simple, inplace=True)

tests/test_tensor/test_tensor_arbgeom.py

@@ -121,9 +121,14 @@ def test_gate_sandwich_with_op():
 def test_normalize_simple():
     psi = qtn.PEPS.rand(3, 3, 2, dtype=complex)
     gauges = {}
-    psi.gauge_all_simple_(100, 5e-6, gauges=gauges)
+    info = {}
+    psi.gauge_all_simple_(100, 5e-6, gauges=gauges, info=info)
     psi.normalize_simple(gauges)
 
+    assert info["iterations"] <= 100
+    if info["iterations"] < 100:
+        assert info["max_sdiff"] < 5e-6
+
     for where in [
         [(0, 0)],
         [(1, 1), (1, 2)],