Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/experimental/autojittn.py

@@ -0,0 +1,129 @@
+"""Decorator for automatically just in time compiling tensor network functions.
+
+TODO::
+
+    - [ ] go via an intermediate pytree / array function, that could be shared
+          e.g. with the TNOptimizer class.
+
+"""
+import functools
+
+import autoray as ar
+
+
+
+class AutojittedTN:
+    """Class to hold the ``autojit_tn`` decorated function callable.
+    """
+
+    def __init__(
+        self,
+        fn,
+        decorator=ar.autojit,
+        check_inputs=True,
+        **decorator_opts
+    ):
+        self.fn = fn
+        self.fn_store = {}
+        self.decorator_opts = decorator_opts
+        self.check_inputs = check_inputs
+        self.decorator = decorator
+
+    def setup_fn(self, tn, *args, **kwargs):
+        from quimb.tensor import TensorNetwork
+
+        @self.decorator(**self.decorator_opts)
+        def fn_jit(arrays):
+            # use separate TN to trace through function
+            jtn = tn.copy()
+
+            # insert the tracing arrays
+            for t, array in zip(jtn, arrays):
+                t.modify(data=array)
+
+            # run function on TN with tracing arrays
+            result = self.fn(jtn, *args, **kwargs)
+
+            # check for a inplace tn function
+            if isinstance(result, TensorNetwork):
+                if result is not jtn:
+                    raise ValueError(
+                        "If you are compiling a function that returns a"
+                        " tensor network it needs to be inplace.")
+                self.inplace = True
+                return tuple(t.data for t in jtn)
+            else:
+                # function returns raw scalar/array(s)
+                self.inplace = False
+                return result
+
+        return fn_jit
+
+    def __call__(self, tn, *args, **kwargs):
+
+        # do we need to generate a new function for these inputs
+        if self.check_inputs:
+            key = (
+                tn.geometry_hash(strict_index_order=True),
+                tuple(args),
+                tuple(sorted(kwargs.items())),
+            )
+        else:
+            # always use the same function
+            key = None
+
+        if key not in self.fn_store:
+            self.fn_store[key] = self.setup_fn(tn, *args, **kwargs)
+        fn_jit = self.fn_store[key]
+
+        # run the compiled function
+        arrays = tuple(t.data for t in tn)
+        out = fn_jit(arrays)
+
+        if self.inplace:
+            # reinsert output arrays into input TN structure
+            for t, array in zip(tn, out):
+                t.modify(data=array)
+            return tn
+
+        return out
+
+
+def autojit_tn(
+    fn=None,
+    decorator=ar.autojit,
+    check_inputs=True,
+    **decorator_opts,
+):
+    """Decorate a tensor network function to be just in time compiled / traced.
+    This traces solely array operations resulting in a completely static
+    computational graph with no side-effects. The resulting function can be
+    much faster if called repeatedly with only numeric changes, or hardware
+    accelerated if a library such as ``jax`` is used.
+
+    Parameters
+    ----------
+    fn : callable
+        The function to be decorated. It should take as its first argument a
+        :class:`~quimb.tensor.tensor_core.TensorNetwork` and return either act
+        inplace on it or return a raw scalar or array.
+    decorator : callable
+        The decorator to use to wrap the underlying array function. For example
+        ``jax.jit``. Defaults to ``autoray.autojit``.
+    check_inputs : bool, optional
+        Whether to check the inputs to the function every call to see if a new
+        compiled function needs to be generated. If ``False`` the same compiled
+        function will be used for all inputs which might be incorrect. Defaults
+        to ``True``.
+    decorator_opts
+        Options to pass to the decorator, e.g. ``backend`` for
+        ``autoray.autojit``.
+    """
+    kwargs = {
+        'decorator': decorator,
+        'check_inputs': check_inputs,
+        **decorator_opts,
+    }
+    if fn is None:
+        return functools.partial(autojit_tn, **kwargs)
+    return functools.wraps(fn)(AutojittedTN(fn, **kwargs))

trajectree/quimb/quimb/experimental/belief_propagation/__init__.py

@@ -0,0 +1,109 @@
+"""Belief propagation (BP) routines. There are three potential categorizations
+of BP and each combination of them is potentially valid specific algorithm.
+
+1-norm vs 2-norm BP
+-------------------
+
+- 1-norm (normal): BP runs directly on the tensor network, messages have size
+  ``d`` where ``d`` is the size of the bond(s) connecting two tensors or
+  regions.
+- 2-norm (quantum): BP runs on the squared tensor network, messages have size
+  ``d^2`` where ``d`` is the size of the bond(s) connecting two tensors or
+  regions. Each local tensor or region is partially traced (over dangling
+  indices) with its conjugate to create a single node.
+
+
+Graph vs Hypergraph BP
+----------------------
+
+- Graph (simple): the tensor network lives on a graph, where indices either
+  appear on two tensors (a bond), or appear on a single tensor (are outputs).
+  In this case, messages are exchanged directly between tensors.
+- Hypergraph: the tensor network lives on a hypergraph, where indices can
+  appear on any number of tensors. In this case, the update procedure is two
+  parts, first all 'tensor' messages are computed, these are then used in the
+  second step to compute all the 'index' messages, which are then fed back into
+  the 'tensor' message update and so forth. For 2-norm BP one likely needs to
+  specify which indices are outputs and should be traced over.
+
+The hypergraph case of course includes the graph case, but since the 'index'
+message update is simply the identity, it is convenient to have a separate
+simpler implementation, where the standard TN bond vs physical index
+definitions hold.
+
+
+Dense vs Vectorized vs Lazy BP
+------------------------------
+
+- Dense: each node is a single tensor, or pair of tensors for 2-norm BP. If all
+  multibonds have been fused, then each message is a vector (1-norm case) or
+  matrix (2-norm case).
+- Vectorized: the same as the above, but all matching tensor update and message
+  updates are stacked and performed simultaneously. This can be enormously more
+  efficient for large numbers of small tensors.
+- Lazy: each node is potentially a tensor network itself with arbitrary inner
+  structure and number of bonds connecting to other nodes. The message are
+  generally tensors and each update is a lazy contraction, which is potentially
+  much cheaper / requires less memory than forming the 'dense' node for large
+  tensors.
+
+(There is also the MPS flavor where each node has a 1D structure and the
+messages are matrix product states, with updates involving compression.)
+
+
+Overall that gives 12 possible BP flavors, some implemented here:
+
+- [x] (HD1BP) hyper, dense, 1-norm - this is the standard BP algorithm
+- [x] (HD2BP) hyper, dense, 2-norm
+- [x] (HV1BP) hyper, vectorized, 1-norm
+- [ ] (HV2BP) hyper, vectorized, 2-norm
+- [ ] (HL1BP) hyper, lazy, 1-norm
+- [ ] (HL2BP) hyper, lazy, 2-norm
+- [x] (D1BP) simple, dense, 1-norm - simple BP for simple tensor networks
+- [x] (D2BP) simple, dense, 2-norm - this is the standard PEPS BP algorithm
+- [ ] (V1BP) simple, vectorized, 1-norm
+- [ ] (V2BP) simple, vectorized, 2-norm
+- [x] (L1BP) simple, lazy, 1-norm
+- [x] (L2BP) simple, lazy, 2-norm
+
+The 2-norm methods can be used to compress bonds or estimate the 2-norm.
+The 1-norm methods can be used to estimate the 1-norm, i.e. contracted value.
+Both methods can be used to compute index marginals and thus perform sampling.
+
+The vectorized methods can be extremely fast for large numbers of small
+tensors, but do currently require all dimensions to match.
+
+The dense and lazy methods can can converge messages *locally*, i.e. only
+update messages adjacent to messages which have changed.
+"""
+
+from .bp_common import initialize_hyper_messages
+from .d1bp import D1BP, contract_d1bp
+from .d2bp import D2BP, compress_d2bp, contract_d2bp, sample_d2bp
+from .hd1bp import HD1BP, contract_hd1bp, sample_hd1bp
+from .hv1bp import HV1BP, contract_hv1bp, sample_hv1bp
+from .l1bp import L1BP, contract_l1bp
+from .l2bp import L2BP, compress_l2bp, contract_l2bp
+from .regions import RegionGraph
+
+__all__ = (
+    "compress_d2bp",
+    "compress_l2bp",
+    "contract_d1bp",
+    "contract_d2bp",
+    "contract_hd1bp",
+    "contract_hv1bp",
+    "contract_l1bp",
+    "contract_l2bp",
+    "D1BP",
+    "D2BP",
+    "HD1BP",
+    "HV1BP",
+    "initialize_hyper_messages",
+    "L1BP",
+    "L2BP",
+    "RegionGraph",
+    "sample_d2bp",
+    "sample_hd1bp",
+    "sample_hv1bp",
+)

trajectree/quimb/quimb/experimental/belief_propagation/bp_common.py

@@ -0,0 +1,397 @@
+import functools
+import operator
+
+import autoray as ar
+
+import quimb.tensor as qtn
+
+
+def prod(xs):
+    """Product of all elements in ``xs``."""
+    return functools.reduce(operator.mul, xs)
+
+
+class RollingDiffMean:
+    """Tracker for the absolute rolling mean of diffs between values, to
+    assess effective convergence of BP above actual message tolerance.
+    """
+
+    def __init__(self, size=16):
+        self.size = size
+        self.diffs = []
+        self.last_x = None
+        self.dxsum = 0.0
+
+    def update(self, x):
+        if self.last_x is not None:
+            dx = x - self.last_x
+            self.diffs.append(dx)
+            self.dxsum += dx / self.size
+        if len(self.diffs) > self.size:
+            dx = self.diffs.pop(0)
+            self.dxsum -= dx / self.size
+        self.last_x = x
+
+    def absmeandiff(self):
+        if len(self.diffs) < self.size:
+            return float("inf")
+        return abs(self.dxsum)
+
+
+class BeliefPropagationCommon:
+    """Common interfaces for belief propagation algorithms.
+
+    Parameters
+    ----------
+    max_iterations : int, optional
+        The maximum number of iterations to perform.
+    tol : float, optional
+        The convergence tolerance for messages.
+    progbar : bool, optional
+        Whether to show a progress bar.
+    """
+
+    def run(self, max_iterations=1000, tol=5e-6, info=None, progbar=False):
+        if progbar:
+            import tqdm
+
+            pbar = tqdm.tqdm()
+        else:
+            pbar = None
+
+        try:
+            it = 0
+            rdm = RollingDiffMean()
+            self.converged = False
+            while not self.converged and it < max_iterations:
+                # perform a single iteration of BP
+                # we supply tol here for use with local convergence
+                nconv, ncheck, max_mdiff = self.iterate(tol=tol)
+                it += 1
+
+                # check rolling mean convergence
+                rdm.update(max_mdiff)
+                self.converged = (max_mdiff < tol) or (rdm.absmeandiff() < tol)
+
+                if pbar is not None:
+                    pbar.set_description(
+                        f"nconv: {nconv}/{ncheck} max|dM|={max_mdiff:.2e}",
+                        refresh=False,
+                    )
+                    pbar.update()
+
+        finally:
+            if pbar is not None:
+                pbar.close()
+
+        if tol != 0.0 and not self.converged:
+            import warnings
+
+            warnings.warn(
+                f"Belief propagation did not converge after {max_iterations} "
+                f"iterations, tol={tol:.2e}, max|dM|={max_mdiff:.2e}."
+            )
+
+        if info is not None:
+            info["converged"] = self.converged
+            info["iterations"] = it
+            info["max_mdiff"] = max_mdiff
+            info["rolling_abs_mean_diff"] = rdm.absmeandiff()
+
+
+def initialize_hyper_messages(
+    tn,
+    fill_fn=None,
+    smudge_factor=1e-12,
+):
+    """Initialize messages for belief propagation, this is equivalent to doing
+    a single round of belief propagation with uniform messages.
+
+    Parameters
+    ----------
+    tn : TensorNetwork
+        The tensor network to initialize messages for.
+    fill_fn : callable, optional
+        A function to fill the messages with, of signature ``fill_fn(shape)``.
+    smudge_factor : float, optional
+        A small number to add to the messages to avoid numerical issues.
+
+    Returns
+    -------
+    messages : dict
+        The initial messages. For every index and tensor id pair, there will
+        be a message to and from with keys ``(ix, tid)`` and ``(tid, ix)``.
+    """
+    from quimb.tensor.contraction import array_contract
+
+    backend = ar.infer_backend(next(t.data for t in tn))
+    _sum = ar.get_lib_fn(backend, "sum")
+
+    messages = {}
+
+    # compute first messages from tensors to indices
+    for tid, t in tn.tensor_map.items():
+        k_inputs = tuple(range(t.ndim))
+        for i, ix in enumerate(t.inds):
+            if fill_fn is None:
+                # sum over all other indices to get initial message
+                m = array_contract(
+                    arrays=(t.data,),
+                    inputs=(k_inputs,),
+                    output=(i,),
+                )
+                # normalize and insert
+                messages[tid, ix] = m / _sum(m)
+            else:
+                d = t.ind_size(ix)
+                m = fill_fn((d,))
+                messages[tid, ix] = m / _sum(m)
+
+    # compute first messages from indices to tensors
+    for ix, tids in tn.ind_map.items():
+        ms = [messages[tid, ix] for tid in tids]
+        mp = prod(ms)
+        for mi, tid in zip(ms, tids):
+            m = mp / (mi + smudge_factor)
+            # normalize and insert
+            messages[ix, tid] = m / _sum(m)
+
+    return messages
+
+
+def combine_local_contractions(
+    tvals,
+    mvals,
+    backend,
+    strip_exponent=False,
+    check_for_zero=True,
+):
+    _abs = ar.get_lib_fn(backend, "abs")
+    _log10 = ar.get_lib_fn(backend, "log10")
+
+    mantissa = 1
+    exponent = 0
+    for vt in tvals:
+        avt = _abs(vt)
+
+        if check_for_zero and (avt == 0.0):
+            if strip_exponent:
+                return 0.0, 0.0
+            else:
+                return 0.0
+
+        mantissa = mantissa * (vt / avt)
+        exponent = exponent + _log10(avt)
+    for mt in mvals:
+        amt = _abs(mt)
+        mantissa = mantissa / (mt / amt)
+        exponent = exponent - _log10(amt)
+
+    if strip_exponent:
+        return mantissa, exponent
+    else:
+        return mantissa * 10**exponent
+
+
+def contract_hyper_messages(
+    tn,
+    messages,
+    strip_exponent=False,
+    backend=None,
+):
+    """Estimate the contraction of ``tn`` given ``messages``, via the
+    exponential of the Bethe free entropy.
+    """
+    tvals = []
+    mvals = []
+
+    for tid, t in tn.tensor_map.items():
+        if backend is None:
+            backend = ar.infer_backend(t.data)
+
+        arrays = [t.data]
+        inputs = [range(t.ndim)]
+        for i, ix in enumerate(t.inds):
+            arrays.append(messages[ix, tid])
+            inputs.append((i,))
+
+            # local message overlap correction
+            mvals.append(
+                qtn.array_contract(
+                    (messages[tid, ix], messages[ix, tid]),
+                    inputs=((0,), (0,)),
+                    output=(),
+                )
+            )
+
+        # local factor free entropy
+        tvals.append(qtn.array_contract(arrays, inputs, output=()))
+
+    for ix, tids in tn.ind_map.items():
+        arrays = tuple(messages[tid, ix] for tid in tids)
+        inputs = tuple((0,) for _ in tids)
+        # local variable free entropy
+        tvals.append(qtn.array_contract(arrays, inputs, output=()))
+
+    return combine_local_contractions(
+        tvals, mvals, backend, strip_exponent=strip_exponent
+    )
+
+
+def compute_index_marginal(tn, ind, messages):
+    """Compute the marginal for a single index given ``messages``.
+
+    Parameters
+    ----------
+    tn : TensorNetwork
+        The tensor network to compute the marginal for.
+    ind : int
+        The index to compute the marginal for.
+    messages : dict
+        The messages to use, which should match ``tn``.
+
+    Returns
+    -------
+    marginal : array_like
+        The marginal probability distribution for the index ``ind``.
+    """
+    m = prod(messages[tid, ind] for tid in tn.ind_map[ind])
+    return m / ar.do("sum", m)
+
+
+def compute_tensor_marginal(tn, tid, messages):
+    """Compute the marginal for the region surrounding a single tensor/factor
+    given ``messages``.
+
+    Parameters
+    ----------
+    tn : TensorNetwork
+        The tensor network to compute the marginal for.
+    tid : int
+        The tensor id to compute the marginal for.
+    messages : dict
+        The messages to use, which should match ``tn``.
+
+    Returns
+    -------
+    marginal : array_like
+        The marginal probability distribution for the tensor/factor ``tid``.
+    """
+    t = tn.tensor_map[tid]
+
+    output = tuple(range(t.ndim))
+    inputs = [output]
+    arrays = [t.data]
+
+    for i, ix in enumerate(t.inds):
+        mix = prod(
+            messages[otid, ix] for otid in tn.ind_map[ix] if otid != tid
+        )
+        inputs.append((i,))
+        arrays.append(mix)
+
+    m = qtn.array_contract(
+        arrays=arrays,
+        inputs=inputs,
+        output=output,
+    )
+
+    return m / ar.do("sum", m)
+
+
+def compute_all_index_marginals_from_messages(tn, messages):
+    """Compute all index marginals from belief propagation messages.
+
+    Parameters
+    ----------
+    tn : TensorNetwork
+        The tensor network to compute marginals for.
+    messages : dict
+        The belief propagation messages.
+
+    Returns
+    -------
+    marginals : dict
+        The marginals for each index.
+    """
+    return {ix: compute_index_marginal(tn, ix, messages) for ix in tn.ind_map}
+
+
+def maybe_get_thread_pool(thread_pool):
+    """Get a thread pool if requested."""
+    if thread_pool is False:
+        return None
+
+    if thread_pool is True:
+        import quimb as qu
+
+        return qu.get_thread_pool()
+
+    if isinstance(thread_pool, int):
+        import quimb as qu
+
+        return qu.get_thread_pool(thread_pool)
+
+    return thread_pool
+
+
+def create_lazy_community_edge_map(tn, site_tags=None, rank_simplify=True):
+    """For lazy BP algorithms, create the data structures describing the
+    effective graph of the lazily grouped 'sites' given by ``site_tags``.
+    """
+    if site_tags is None:
+        site_tags = set(tn.site_tags)
+    else:
+        site_tags = set(site_tags)
+
+    edges = {}
+    neighbors = {}
+    local_tns = {}
+    touch_map = {}
+
+    for ix in tn.ind_map:
+        ts = tn._inds_get(ix)
+        tags = {tag for t in ts for tag in t.tags if tag in site_tags}
+        if len(tags) >= 2:
+            i, j = tuple(sorted(tags))
+
+            if (i, j) in edges:
+                # already processed this edge
+                continue
+
+            # add to neighbor map
+            neighbors.setdefault(i, []).append(j)
+            neighbors.setdefault(j, []).append(i)
+
+            # get local TNs and compute bonds between them,
+            # rank simplify here also to prepare for contractions
+            try:
+                tn_i = local_tns[i]
+            except KeyError:
+                tn_i = local_tns[i] = tn.select(i, virtual=False)
+                if rank_simplify:
+                    tn_i.rank_simplify_()
+            try:
+                tn_j = local_tns[j]
+            except KeyError:
+                tn_j = local_tns[j] = tn.select(j, virtual=False)
+                if rank_simplify:
+                    tn_j.rank_simplify_()
+
+            edges[i, j] = tuple(qtn.bonds(tn_i, tn_j))
+
+    for i, j in edges:
+        touch_map[(i, j)] = tuple((j, k) for k in neighbors[j] if k != i)
+        touch_map[(j, i)] = tuple((i, k) for k in neighbors[i] if k != j)
+
+    if len(local_tns) != len(site_tags):
+        # handle potentially disconnected sites
+        for i in sorted(site_tags):
+            try:
+                tn_i = local_tns[i] = tn.select(i, virtual=False)
+                if rank_simplify:
+                    tn_i.rank_simplify_()
+            except KeyError:
+                pass
+
+    return edges, neighbors, local_tns, touch_map