add tn.contract(strip_exponent=True)

Author: jcmgray

ci/requirements/py-base.yml

@@ -2,7 +2,7 @@ channels:
   - conda-forge
 dependencies:
   - autoray>=0.6.12
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cytoolz
   - joblib

ci/requirements/py-jax.yml

@@ -2,7 +2,7 @@ channels:
   - conda-forge
 dependencies:
   - autoray>=0.6.12
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cytoolz
   - joblib

ci/requirements/py-openblas.yml

@@ -3,7 +3,7 @@ channels:
 dependencies:
   - autoray>=0.6.12
   - blas=*=openblas
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cytoolz
   - joblib

ci/requirements/py-slepc.yml

@@ -2,7 +2,7 @@ channels:
   - conda-forge
 dependencies:
   - autoray>=0.6.12
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cytoolz
   - joblib

ci/requirements/py-tensorflow.yml

@@ -2,7 +2,7 @@ channels:
   - conda-forge
 dependencies:
   - autoray>=0.7.0
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cytoolz
   - joblib

ci/requirements/py-torch.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - autoray>=0.6.12
-  - cotengra>=0.6.1
+  - cotengra>=0.7.1
   - coverage
   - cpuonly
   - cytoolz

docs/changelog.md

@@ -8,6 +8,7 @@ Release notes for `quimb`.
 **Breaking Changes**
 
 - move belief propagation to `quimb.tensor.belief_propagation`
+- calling `tn.contract()` when an non-zero value has been accrued into `tn.exponent` now automatically re-absorbs that exponent.
 
 **Enhancements:**
 
@@ -19,8 +20,9 @@ Release notes for `quimb`.
 - belief propagation, add a `contract_every` option.
 - HV1BP: vectorize both contraction and message initialization
 - 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
+- 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 diffs
 - [`Tensor.gate`](quimb.tensor.tensor_core.Tensor.gate): add `transposed` option
+- [`TensorNetwork.contract`](quimb.tensor.tensor_core.TensorNetwork.contract): add `strip_exponent` option for return the mantissa and exponent (log10) separately.
 
 **Bug fixes:**
 

pyproject.toml

@@ -27,14 +27,14 @@ keywords = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "autoray >=0.6.12",
-    "cotengra >=0.6.1",
-    "cytoolz >=0.8.0",
-    "numba >=0.39",
-    "numpy >=1.17",
-    "psutil >=4.3.1",
-    "scipy >=1.0.0",
-    "tqdm >=4",
+    "autoray>=0.6.12",
+    "cotengra>=0.7.1",
+    "cytoolz>=0.8.0",
+    "numba>=0.39",
+    "numpy>=1.17",
+    "psutil>=4.3.1",
+    "scipy>=1.0.0",
+    "tqdm>=4",
 ]
 
 
@@ -56,7 +56,7 @@ tests = [
 docs = [
     "astroid<3.0.0",
     "autoray>=0.6.12",
-    "cotengra>=0.6.1",
+    "cotengra>=0.7.1",
     "doc2dash>=2.4.1",
     "furo",
     "ipython!=8.7.0",

quimb/tensor/contraction.py

@@ -284,7 +284,12 @@ def array_contract(
     if backend is None:
         backend = get_contract_backend()
     return ctg.array_contract(
-        arrays, inputs, output, optimize=optimize, backend=backend, **kwargs
+        arrays,
+        inputs,
+        output,
+        optimize=optimize,
+        backend=backend,
+        **kwargs,
     )
 
 

quimb/tensor/tensor_core.py

@@ -214,6 +214,8 @@ def tensor_contract(
     backend=None,
     preserve_tensor=False,
     drop_tags=False,
+    strip_exponent=False,
+    exponent=None,
     **contract_opts,
 ):
     """Contract a collection of tensors into a scalar or tensor, automatically
@@ -268,6 +270,11 @@ def tensor_contract(
     drop_tags : bool, optional
         Whether to drop all tags from the output tensor. By default the output
         tensor will keep the union of all tags from the input tensors.
+    strip_exponent : bool, optional
+        If `True`, return the exponent of the result, log10, as well as the
+        rescaled 'mantissa'. Useful for very large or small values.
+    exponent : float, optional
+        If supplied, a base exponent to add to the result exponent.
     contract_opts
         Passed to ``cotengra.array_contract``.
 
@@ -300,20 +307,40 @@ def tensor_contract(
         inds,
         inds_out,
         optimize=optimize,
+        strip_exponent=strip_exponent,
         backend=backend,
         **contract_opts,
     )
 
-    if not inds_out and not preserve_tensor:
-        return maybe_realify_scalar(data_out)
+    if strip_exponent:
+        # mantissa and exponent returned separately
+        data_out, result_exponent = data_out
+
+        if exponent is not None:
+            # custom base exponent supplied
+            result_exponent = result_exponent + exponent
 
-    if drop_tags:
-        tags_out = None
+    elif exponent is not None:
+        # custom exponent but not stripping, so we need to scale the result
+        data_out = data_out * 10**exponent
+
+    if not inds_out and not preserve_tensor:
+        # return a scalar, possibly casting to real
+        # but only if numpy with v. small imag part
+        result = maybe_realify_scalar(data_out)
     else:
-        # union of all tags
-        tags_out = oset_union(t.tags for t in tensors)
+        if drop_tags:
+            tags_out = None
+        else:
+            # union of all tags
+            tags_out = oset_union(t.tags for t in tensors)
+
+        result = Tensor(data=data_out, inds=inds_out, tags=tags_out)
 
-    return Tensor(data=data_out, inds=inds_out, tags=tags_out)
+    if strip_exponent:
+        return result, result_exponent
+
+    return result
 
 
 # generate a random base to avoid collisions on difference processes ...
@@ -8736,6 +8763,8 @@ def contract(
         backend=None,
         preserve_tensor=False,
         max_bond=None,
+        strip_exponent=False,
+        exponent=True,
         inplace=False,
         **opts,
     ):
@@ -8788,6 +8817,14 @@ def contract(
         preserve_tensor : bool, optional
             Whether to return a tensor regardless of whether the output object
             is a scalar (has no indices) or not.
+        strip_exponent : bool, optional
+            If contracting the entire tensor network, whether to strip a log10
+            exponent and return it separately. This is useful for very large or
+            small values.
+        exponent : float, optional
+            The current exponent to scale the whole contraction by. If ``True``
+            this taken from `tn.exponent`. If `False` then this is ignored.
+            If a float, this is the exponent to use.
         inplace : bool, optional
             Whether to perform the contraction inplace. This is only valid
             if not all tensors are contracted (which doesn't produce a TN).
@@ -8806,6 +8843,7 @@ def contract(
         --------
         contract_tags, contract_cumulative
         """
+        # for visibility we put these in the function signature
         opts["output_inds"] = output_inds
         opts["optimize"] = optimize
         opts["get"] = get
@@ -8823,21 +8861,42 @@ def contract(
                 raise NotImplementedError
 
             return self.contract_compressed(
-                max_bond=max_bond, inplace=inplace, **opts
+                max_bond=max_bond,
+                inplace=inplace,
+                **opts,
             )
 
         # this checks whether certain TN classes have a manually specified
         #     contraction pattern (e.g. 1D along the line)
         if self._CONTRACT_STRUCTURED:
             if (tags is ...) or isinstance(tags, slice):
-                return self.contract_structured(tags, inplace=inplace, **opts)
+                return self.contract_structured(
+                    tags,
+                    inplace=inplace,
+                    **opts,
+                )
 
         # contracting everything to single output
         if all_tags and not inplace:
-            return tensor_contract(*self.tensor_map.values(), **opts)
+
+            if exponent is True:
+                exponent = self.exponent
+            elif exponent is False:
+                exponent = 0.0
+
+            return tensor_contract(
+                *self.tensor_map.values(),
+                strip_exponent=strip_exponent,
+                exponent=exponent,
+                **opts
+            )
 
         # contract some or all tensors, but keeping tensor network
-        return self.contract_tags(tags, inplace=inplace, **opts)
+        return self.contract_tags(
+            tags,
+            inplace=inplace,
+            **opts
+        )
 
     contract_ = functools.partialmethod(contract, inplace=True)
 

tests/test_tensor/test_belief_propagation/test_l2bp.py

@@ -61,7 +61,7 @@ def test_contract_double_layer_tree_exact(dtype):
     norm_bp = qbp.contract_l2bp(tn, info=info, progbar=True)
     assert info["converged"]
 
-    assert norm_bp == pytest.approx(norm_ex, rel=1e-6)
+    assert norm_bp == pytest.approx(norm_ex, rel=5e-6)
 
 
 @pytest.mark.parametrize("dtype", ["float32", "complex64"])

tests/test_tensor/test_contract.py

@@ -1,10 +1,36 @@
-import pytest
 import numpy as np
+import pytest
 
 import quimb.tensor as qtn
 from quimb.tensor.contraction import _CONTRACT_BACKEND, _TENSOR_LINOP_BACKEND
 
 
+def test_tensor_contract_strip_exponent():
+    tn = qtn.TN_rand_reg(10, 3, 3, dtype=complex)
+    z0 = tn.contract()
+    m1, e1 = qtn.tensor_contract(*tn, strip_exponent=True)
+    assert m1 * 10**e1 == pytest.approx(z0)
+    # test tn.exponent is reinserted
+    tn.equalize_norms_(value=1.0)
+    z2 = tn.contract()
+    assert z2 == pytest.approx(z0)
+    # test tn.exponent is reinserted with strip exponent
+    m3, e3 = tn.contract(strip_exponent=True)
+    assert m3 * 10**e3 == pytest.approx(z0)
+    # test tn.exponent is not reinserted when specified
+    z4 = tn.contract(exponent=False)
+    assert z4 != pytest.approx(z0)
+    # test tn.exponent is not reinserted when specified with strip exponent
+    m5, e5 = tn.contract(strip_exponent=True, exponent=False)
+    assert m5 * 10**e5 != pytest.approx(z0)
+    # test explicit exponent
+    z6 = tn.contract(exponent=tn.exponent)
+    assert z6 == pytest.approx(z0)
+    # test explicit exponent with strip exponent
+    m7, e7 = tn.contract(strip_exponent=True, exponent=tn.exponent)
+    assert m7 * 10**e7 == pytest.approx(z0)
+
+
 class TestContractOpts:
     def test_contract_strategy(self):
         assert qtn.get_contract_strategy() == "greedy"