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

trajectree/quimb/quimb/experimental/belief_propagation/hd1bp.py

@@ -1,571 +0,0 @@
-"""Hyper dense belief propagation for arbitrary `quimb` tensor networks. This
-is the classic 1-norm version of belief propagation, which treats the tensor
-network directly as a factor graph. Messages are processed one at a time.
-
-TODO:
-
-- [ ] implement 'touching', so that only necessary messages are updated
-- [ ] implement sequential update
-
-"""
-import autoray as ar
-import quimb.tensor as qtn
-
-from .bp_common import (
-    BeliefPropagationCommon,
-    compute_all_index_marginals_from_messages,
-    contract_hyper_messages,
-    initialize_hyper_messages,
-    prod,
-)
-
-
-def compute_all_hyperind_messages_prod(ms, smudge_factor=1e-12):
-    """Given set of messages ``ms`` incident to a single index, compute the
-    corresponding next output messages, using the 'product' implementation.
-    """
-    if len(ms) == 2:
-        # shortcut for 2 messages
-        return [ms[1], ms[0]]
-
-    x = prod(ms)
-    return [x / (m + smudge_factor) for m in ms]
-
-
-def compute_all_hyperind_messages_tree(ms):
-    """Given set of messages ``ms`` incident to a single index, compute the
-    corresponding next output messages, using the 'tree' implementation.
-    """
-    ndim = len(ms)
-    if len(ms) == 2:
-        # shortcut for 2 messages
-        return [ms[1], ms[0]]
-
-    mouts = [None for _ in range(ndim)]
-    queue = [(tuple(range(ndim)), 1, ms)]
-
-    while queue:
-        js, x, ms = queue.pop()
-
-        ndim = len(ms)
-        if ndim == 1:
-            # reached single message
-            mouts[js[0]] = x
-            continue
-        elif ndim == 2:
-            # shortcut for 2 messages left
-            mouts[js[0]] = x * ms[1]
-            mouts[js[1]] = ms[0] * x
-            continue
-
-        # else split in two and contract each half
-        k = ndim // 2
-        jl, jr = js[:k], js[k:]
-        ml, mr = ms[:k], ms[k:]
-
-        # contract the right messages to get new left array
-        xl = prod((*mr, x))
-
-        # contract the left messages to get new right array
-        xr = prod((*ml, x))
-
-        # add the queue for possible further halving
-        queue.append((jl, xl, ml))
-        queue.append((jr, xr, mr))
-
-    return mouts
-
-
-def compute_all_tensor_messages_shortcuts(x, ms, ndim):
-    if ndim == 2:
-        # shortcut for 2 messages
-        return [x @ ms[1], ms[0] @ x]
-    elif ndim == 1:
-        # shortcut for single message
-        return [x]
-    elif ndim == 0:
-        # shortcut for no messages
-        return []
-
-
-def compute_all_tensor_messages_prod(
-    x,
-    ms,
-    backend=None,
-    smudge_factor=1e-12,
-):
-    """Given set of messages ``ms`` incident to tensor with data ``x``, compute
-    the corresponding next output messages, using the 'prod' implementation.
-    """
-    ndim = len(ms)
-    if ndim <= 2:
-        return compute_all_tensor_messages_shortcuts(x, ms, ndim)
-
-    js = tuple(range(ndim))
-
-    mx = qtn.array_contract(
-        arrays=(x, *ms), inputs=(js, *((j,) for j in js)), output=js
-    )
-    mouts = []
-
-    for j, g in enumerate(ms):
-        mouts.append(
-            qtn.array_contract(
-                arrays=(mx, 1 / (g + smudge_factor)),
-                inputs=(js, (j,)),
-                output=(j,),
-                backend=backend,
-            )
-        )
-
-    return mouts
-
-
-def compute_all_tensor_messages_tree(x, ms, backend=None):
-    """Given set of messages ``ms`` incident to tensor with data ``x``, compute
-    the corresponding next output messages, using the 'tree' implementation.
-    """
-    ndim = len(ms)
-    if ndim <= 2:
-        return compute_all_tensor_messages_shortcuts(x, ms, ndim)
-
-    mouts = [None for _ in range(ndim)]
-    queue = [(tuple(range(ndim)), x, ms)]
-
-    while queue:
-        js, x, ms = queue.pop()
-
-        ndim = len(ms)
-        if ndim == 1:
-            # reached single message
-            mouts[js[0]] = x
-            continue
-        elif ndim == 2:
-            # shortcut for 2 messages left
-            mouts[js[0]] = x @ ms[1]
-            mouts[js[1]] = ms[0] @ x
-            continue
-
-        # else split in two and contract each half
-        k = ndim // 2
-        jl, jr = js[:k], js[k:]
-        ml, mr = ms[:k], ms[k:]
-
-        # contract the right messages to get new left array
-        xl = qtn.array_contract(
-            arrays=(x, *mr),
-            inputs=(js, *((j,) for j in jr)),
-            output=jl,
-            backend=backend,
-        )
-
-        # contract the left messages to get new right array
-        xr = qtn.array_contract(
-            arrays=(x, *ml),
-            inputs=(js, *((j,) for j in jl)),
-            output=jr,
-            backend=backend,
-        )
-
-        # add the queue for possible further halving
-        queue.append((jl, xl, ml))
-        queue.append((jr, xr, mr))
-
-    return mouts
-
-
-def iterate_belief_propagation_basic(
-    tn,
-    messages,
-    damping=None,
-    smudge_factor=1e-12,
-    tol=None,
-):
-    """Run a single iteration of belief propagation. This is the basic version
-    that does not vectorize contractions.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to run BP on.
-    messages : dict
-        The current messages. For every index and tensor id pair, there should
-        be a message to and from with keys ``(ix, tid)`` and ``(tid, ix)``.
-    smudge_factor : float, optional
-        A small number to add to the denominator of messages to avoid division
-        by zero. Note when this happens the numerator will also be zero.
-
-    Returns
-    -------
-    new_messages : dict
-        The new messages.
-    """
-    backend = ar.infer_backend(next(iter(messages.values())))
-
-    # _sum = ar.get_lib_fn(backend, "sum")
-    # n.b. at small sizes python sum is faster than numpy sum
-    _sum = ar.get_lib_fn(backend, "sum")
-    # _max = ar.get_lib_fn(backend, "max")
-    _abs = ar.get_lib_fn(backend, "abs")
-
-    def _normalize_and_insert(k, m, max_dm):
-        # normalize and insert
-        m = m / _sum(m)
-
-        old_m = messages[k]
-
-        if damping is not None:
-            # mix old and new
-            m = damping * old_m + (1 - damping) * m
-
-        # compare to the old messages
-        dm = _sum(_abs(m - old_m))
-        max_dm = max(dm, max_dm)
-
-        # set and return the max diff so far
-        messages[k] = m
-        return max_dm
-
-    max_dm = 0.0
-
-    # hyper index messages
-    for ix, tids in tn.ind_map.items():
-        ms = compute_all_hyperind_messages_prod(
-            [messages[tid, ix] for tid in tids], smudge_factor
-        )
-        for tid, m in zip(tids, ms):
-            max_dm = _normalize_and_insert((ix, tid), m, max_dm)
-
-    # tensor messages
-    for tid, t in tn.tensor_map.items():
-        inds = t.inds
-        ms = compute_all_tensor_messages_tree(
-            t.data,
-            [messages[ix, tid] for ix in inds],
-        )
-        for ix, m in zip(inds, ms):
-            max_dm = _normalize_and_insert((tid, ix), m, max_dm)
-
-    return messages, max_dm
-
-
-class HD1BP(BeliefPropagationCommon):
-    """Object interface for hyper, dense, 1-norm belief propagation. This is
-    standard belief propagation in tensor network form.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to run BP on.
-    messages : dict, optional
-        Initial messages to use, if not given then uniform messages are used.
-    smudge_factor : float, optional
-        A small number to add to the denominator of messages to avoid division
-        by zero. Note when this happens the numerator will also be zero.
-    """
-
-    def __init__(
-        self,
-        tn,
-        messages=None,
-        damping=None,
-        smudge_factor=1e-12,
-    ):
-        self.tn = tn
-        self.backend = next(t.backend for t in tn)
-        self.smudge_factor = smudge_factor
-        self.damping = damping
-        if messages is None:
-            messages = initialize_hyper_messages(
-                tn, smudge_factor=smudge_factor
-            )
-        self.messages = messages
-
-    def iterate(self, **kwargs):
-        self.messages, max_dm = iterate_belief_propagation_basic(
-            self.tn,
-            self.messages,
-            damping=self.damping,
-            smudge_factor=self.smudge_factor,
-            **kwargs,
-        )
-        return None, None, max_dm
-
-    def get_gauged_tn(self):
-        """Assuming the supplied tensor network has no hyper or dangling
-        indices, gauge it by inserting the BP-approximated transfer matrix
-        eigenvectors, which may be complex. The BP-contraction of this gauged
-        network is then simply the product of zeroth entries of each tensor.
-        """
-        tng = self.tn.copy()
-        for ind, tids in self.tn.ind_map.items():
-            tida, tidb = tids
-            ka = (ind, tida)
-            kb = (ind, tidb)
-            ma = self.messages[ka]
-            mb = self.messages[kb]
-
-            el, ev = ar.do('linalg.eig', ar.do('outer', ma, mb))
-            k = ar.do('argsort', -ar.do('abs', el))
-            ev = ev[:, k]
-            Uinv = ev
-            U = ar.do('linalg.inv', ev)
-            tng._insert_gauge_tids(U, tida, tidb, Uinv)
-        return tng
-
-    def contract(self, strip_exponent=False):
-        """Estimate the total contraction, i.e. the exponential of the 'Bethe
-        free entropy'.
-        """
-        return contract_hyper_messages(
-            self.tn,
-            self.messages,
-            strip_exponent=strip_exponent,
-            backend=self.backend,
-        )
-
-
-def contract_hd1bp(
-    tn,
-    messages=None,
-    max_iterations=1000,
-    tol=5e-6,
-    damping=0.0,
-    smudge_factor=1e-12,
-    strip_exponent=False,
-    info=None,
-    progbar=False,
-):
-    """Estimate the contraction of ``tn`` with hyper, vectorized, 1-norm
-    belief propagation, via the exponential of the Bethe free entropy.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to run BP on, can have hyper indices.
-    messages : dict, optional
-        Initial messages to use, if not given then uniform messages are used.
-    max_iterations : int, optional
-        The maximum number of iterations to perform.
-    tol : float, optional
-        The convergence tolerance for messages.
-    damping : float, optional
-        The damping factor to use, 0.0 means no damping.
-    smudge_factor : float, optional
-        A small number to add to the denominator of messages to avoid division
-        by zero. Note when this happens the numerator will also be zero.
-    strip_exponent : bool, optional
-        Whether to strip the exponent from the final result. If ``True``
-        then the returned result is ``(mantissa, exponent)``.
-    info : dict, optional
-        If specified, update this dictionary with information about the
-        belief propagation run.
-    progbar : bool, optional
-        Whether to show a progress bar.
-
-    Returns
-    -------
-    scalar or (scalar, float)
-    """
-    bp = HD1BP(
-        tn,
-        messages=messages,
-        damping=damping,
-        smudge_factor=smudge_factor,
-    )
-    bp.run(
-        max_iterations=max_iterations,
-        tol=tol,
-        info=info,
-        progbar=progbar,
-    )
-    return bp.contract(strip_exponent=strip_exponent)
-
-
-def run_belief_propagation_hd1bp(
-    tn,
-    messages=None,
-    max_iterations=1000,
-    tol=5e-6,
-    damping=0.0,
-    smudge_factor=1e-12,
-    info=None,
-    progbar=False,
-):
-    """Run belief propagation on a tensor network until it converges. This
-    is the basic version that does not vectorize contractions.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to run BP on.
-    messages : dict, optional
-        The current messages. For every index and tensor id pair, there should
-        be a message to and from with keys ``(ix, tid)`` and ``(tid, ix)``.
-        If not given, then messages are initialized as uniform.
-    max_iterations : int, optional
-        The maximum number of iterations to run for.
-    tol : float, optional
-        The convergence tolerance.
-    smudge_factor : float, optional
-        A small number to add to the denominator of messages to avoid division
-        by zero. Note when this happens the numerator will also be zero.
-    info : dict, optional
-        If specified, update this dictionary with information about the
-        belief propagation run.
-    progbar : bool, optional
-        Whether to show a progress bar.
-
-    Returns
-    -------
-    messages : dict
-        The final messages.
-    converged : bool
-        Whether the algorithm converged.
-    """
-    bp = HD1BP(
-        tn, messages=messages, damping=damping, smudge_factor=smudge_factor
-    )
-    bp.run(max_iterations=max_iterations, tol=tol, info=info, progbar=progbar)
-    return bp.messages, bp.converged
-
-
-def sample_hd1bp(
-    tn,
-    messages=None,
-    output_inds=None,
-    max_iterations=1000,
-    tol=1e-2,
-    damping=0.0,
-    smudge_factor=1e-12,
-    bias=False,
-    seed=None,
-    progbar=False,
-):
-    """Sample all indices of a tensor network using repeated belief propagation
-    runs and decimation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to sample.
-    messages : dict, optional
-        The current messages. For every index and tensor id pair, there should
-        be a message to and from with keys ``(ix, tid)`` and ``(tid, ix)``.
-        If not given, then messages are initialized as uniform.
-    output_inds : sequence of str, optional
-        The indices to sample. If not given, then all indices are sampled.
-    max_iterations : int, optional
-        The maximum number of iterations for each message passing run.
-    tol : float, optional
-        The convergence tolerance for each message passing run.
-    smudge_factor : float, optional
-        A small number to add to each message to avoid zeros. Making this large
-        is similar to adding a temperature, which can aid convergence but
-        likely produces less accurate marginals.
-    bias : bool or float, optional
-        Whether to bias the sampling towards the largest marginal. If ``False``
-        (the default), then indices are sampled proportional to their
-        marginals. If ``True``, then each index is 'sampled' to be its largest
-        weight value always. If a float, then the local probability
-        distribution is raised to this power before sampling.
-    thread_pool : bool, int or ThreadPoolExecutor, optional
-        Whether to use a thread pool for parallelization. If an integer, then
-        this is the number of threads to use. If ``True``, then the number of
-        threads is set to the number of cores. If a ``ThreadPoolExecutor``,
-        then this is used directly.
-    seed : int, optional
-        A random seed to use for the sampling.
-    progbar : bool, optional
-        Whether to show a progress bar.
-
-    Returns
-    -------
-    config : dict[str, int]
-        The sample configuration, mapping indices to values.
-    tn_config : TensorNetwork
-        The tensor network with all index values (or just those in
-        `output_inds` if supllied) selected. Contracting this tensor network
-        (which will just be a sequence of scalars if all index values have been
-        sampled) gives the weight of the sample, e.g. should be 1 for a SAT
-        problem and valid assignment.
-    omega : float
-        The probability of choosing this sample (i.e. product of marginal
-        values). Useful possibly for importance sampling.
-    """
-    import numpy as np
-
-    rng = np.random.default_rng(seed)
-
-    tn_config = tn.copy()
-
-    if messages is None:
-        messages = initialize_hyper_messages(tn_config)
-
-    if output_inds is None:
-        output_inds = tn_config.ind_map.keys()
-    output_inds = set(output_inds)
-
-    config = {}
-    omega = 1.0
-
-    if progbar:
-        import tqdm
-
-        pbar = tqdm.tqdm(total=len(output_inds))
-    else:
-        pbar = None
-
-    while output_inds:
-        messages, _ = run_belief_propagation_hd1bp(
-            tn_config,
-            messages,
-            max_iterations=max_iterations,
-            tol=tol,
-            damping=damping,
-            smudge_factor=smudge_factor,
-            progbar=True,
-        )
-
-        marginals = compute_all_index_marginals_from_messages(
-            tn_config, messages
-        )
-
-        # choose most peaked marginal
-        ix, p = max(
-            (m for m in marginals.items() if m[0] in output_inds),
-            key=lambda ix_p: max(ix_p[1]),
-        )
-
-        if bias is False:
-            # sample the value according to the marginal
-            v = rng.choice(np.arange(p.size), p=p)
-        elif bias is True:
-            v = np.argmax(p)
-            # in some sense omega is really 1.0 here
-        else:
-            # bias towards larger marginals by raising to a power
-            p = p**bias
-            p /= np.sum(p)
-            v = np.random.choice(np.arange(p.size), p=p)
-
-        omega *= p[v]
-        config[ix] = v
-
-        # clean up messages
-        for tid in tn_config.ind_map[ix]:
-            del messages[ix, tid]
-            del messages[tid, ix]
-
-        # remove index
-        tn_config.isel_({ix: v})
-        output_inds.remove(ix)
-
-        if progbar:
-            pbar.update(1)
-            pbar.set_description(f"{ix}->{v}", refresh=False)
-
-    if progbar:
-        pbar.close()
-
-    return config, tn_config, omega