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

trajectree/quimb/quimb/experimental/belief_propagation/d2bp.py

@@ -1,653 +0,0 @@
-import autoray as ar
-
-import quimb.tensor as qtn
-from quimb.utils import oset
-
-from .bp_common import (
-    BeliefPropagationCommon,
-    combine_local_contractions,
-)
-
-
-class D2BP(BeliefPropagationCommon):
-    """Dense (as in one tensor per site) 2-norm (as in for wavefunctions and
-    operators) belief propagation. Allows messages reuse. This version assumes
-    no hyper indices (i.e. a standard PEPS like tensor network).
-
-    Potential use cases for D2BP and a PEPS like tensor network are:
-
-        - globally compressing it from bond dimension ``D`` to ``D'``
-        - eagerly applying gates and locally compressing back to ``D``
-        - sampling configurations
-        - estimating the norm of the tensor network
-
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to form the 2-norm of and run BP on.
-    messages : dict[(str, int), array_like], optional
-        The initial messages to use, effectively defaults to all ones if not
-        specified.
-    output_inds : set[str], optional
-        The indices to consider as output (dangling) indices of the tn.
-        Computed automatically if not specified.
-    optimize : str or PathOptimizer, optional
-        The path optimizer to use when contracting the messages.
-    damping : float, optional
-        The damping factor to use, 0.0 means no damping.
-    update : {'parallel', 'sequential'}, optional
-        Whether to update all messages in parallel or sequentially.
-    local_convergence : bool, optional
-        Whether to allow messages to locally converge - i.e. if all their
-        input messages have converged then stop updating them.
-    contract_opts
-        Other options supplied to ``cotengra.array_contract``.
-    """
-
-    def __init__(
-        self,
-        tn,
-        messages=None,
-        output_inds=None,
-        optimize="auto-hq",
-        damping=0.0,
-        update="sequential",
-        local_convergence=True,
-        **contract_opts,
-    ):
-        from quimb.tensor.contraction import array_contract_expression
-
-        self.tn = tn
-        self.contract_opts = contract_opts
-        self.contract_opts.setdefault("optimize", optimize)
-        self.damping = damping
-        self.local_convergence = local_convergence
-        self.update = update
-
-        if output_inds is None:
-            self.output_inds = set(self.tn.outer_inds())
-        else:
-            self.output_inds = set(output_inds)
-
-        self.backend = next(t.backend for t in tn)
-        _abs = ar.get_lib_fn(self.backend, "abs")
-        _sum = ar.get_lib_fn(self.backend, "sum")
-
-        def _normalize(x):
-            return x / _sum(x)
-
-        def _distance(x, y):
-            return _sum(_abs(x - y))
-
-        self._normalize = _normalize
-        self._distance = _distance
-
-        if messages is None:
-            self.messages = {}
-        else:
-            self.messages = messages
-
-        # record which messages touch each others, for efficient updates
-        self.touch_map = {}
-        self.touched = oset()
-        self.exprs = {}
-
-        # populate any messages
-        for ix, tids in self.tn.ind_map.items():
-            if ix in self.output_inds:
-                continue
-
-            tida, tidb = tids
-            jx = ix + "*"
-            ta, tb = self.tn._tids_get(tida, tidb)
-
-            for tid, t, t_in in ((tida, ta, tb), (tidb, tb, ta)):
-                this_touchmap = []
-                for nx in t.inds:
-                    if nx in self.output_inds or nx == ix:
-                        continue
-                    # where this message will be sent on to
-                    (tidn,) = (n for n in self.tn.ind_map[nx] if n != tid)
-                    this_touchmap.append((nx, tidn))
-                self.touch_map[ix, tid] = this_touchmap
-
-                if (ix, tid) not in self.messages:
-                    tm = (t_in.reindex({ix: jx}).conj_() @ t_in).data
-                    self.messages[ix, tid] = self._normalize(tm.data)
-
-        # for efficiency setup all the contraction expressions ahead of time
-        for ix, tids in self.tn.ind_map.items():
-            if ix in self.output_inds:
-                continue
-
-            for tida, tidb in (sorted(tids), sorted(tids, reverse=True)):
-                ta = self.tn.tensor_map[tida]
-                kix = ta.inds
-                bix = tuple(
-                    i if i in self.output_inds else i + "*" for i in kix
-                )
-                inputs = [kix, bix]
-                data = [ta.data, ta.data.conj()]
-                shapes = [ta.shape, ta.shape]
-                for i in kix:
-                    if (i != ix) and i not in self.output_inds:
-                        inputs.append((i + "*", i))
-                        data.append((i, tida))
-                        shapes.append(self.messages[i, tida].shape)
-
-                expr = array_contract_expression(
-                    inputs=inputs,
-                    output=(ix + "*", ix),
-                    shapes=shapes,
-                    **self.contract_opts,
-                )
-                self.exprs[ix, tidb] = expr, data
-
-    def update_touched_from_tids(self, *tids):
-        """Specify that the messages for the given ``tids`` have changed."""
-        for tid in tids:
-            t = self.tn.tensor_map[tid]
-            for ix in t.inds:
-                if ix in self.output_inds:
-                    continue
-                (ntid,) = (n for n in self.tn.ind_map[ix] if n != tid)
-                self.touched.add((ix, ntid))
-
-    def update_touched_from_tags(self, tags, which="any"):
-        """Specify that the messages for the messages touching ``tags`` have
-        changed.
-        """
-        tids = self.tn._get_tids_from_tags(tags, which)
-        self.update_touched_from_tids(*tids)
-
-    def update_touched_from_inds(self, inds, which="any"):
-        """Specify that the messages for the messages touching ``inds`` have
-        changed.
-        """
-        tids = self.tn._get_tids_from_inds(inds, which)
-        self.update_touched_from_tids(*tids)
-
-    def iterate(self, tol=5e-6):
-        """Perform a single iteration of dense 2-norm belief propagation."""
-
-        if (not self.local_convergence) or (not self.touched):
-            # assume if asked to iterate that we want to check all messages
-            self.touched.update(self.exprs.keys())
-
-        ncheck = len(self.touched)
-        nconv = 0
-        max_mdiff = -1.0
-        new_touched = oset()
-
-        def _compute_m(key):
-            expr, data = self.exprs[key]
-            m = expr(*data[:2], *(self.messages[mkey] for mkey in data[2:]))
-            # enforce hermiticity and normalize
-            return self._normalize(m + ar.dag(m))
-
-        def _update_m(key, new_m):
-            nonlocal nconv, max_mdiff
-
-            old_m = self.messages[key]
-            if self.damping > 0.0:
-                new_m = self._normalize(
-                    self.damping * old_m + (1 - self.damping) * new_m
-                )
-            try:
-                mdiff = float(self._distance(old_m, new_m))
-            except (TypeError, ValueError):
-                # handle e.g. lazy arrays
-                mdiff = float("inf")
-            if mdiff > tol:
-                # mark touching messages for update
-                new_touched.update(self.touch_map[key])
-            else:
-                nconv += 1
-            max_mdiff = max(max_mdiff, mdiff)
-            self.messages[key] = new_m
-
-        if self.update == "parallel":
-            new_messages = {}
-            # compute all new messages
-            while self.touched:
-                key = self.touched.pop()
-                new_messages[key] = _compute_m(key)
-            # insert all new messages
-            for key, new_m in new_messages.items():
-                _update_m(key, new_m)
-
-        elif self.update == "sequential":
-            # compute each new message and immediately re-insert it
-            while self.touched:
-                key = self.touched.pop()
-                new_m = _compute_m(key)
-                _update_m(key, new_m)
-
-        self.touched = new_touched
-
-        return nconv, ncheck, max_mdiff
-
-    def compute_marginal(self, ind):
-        """Compute the marginal for the index ``ind``."""
-        (tid,) = self.tn.ind_map[ind]
-        t = self.tn.tensor_map[tid]
-
-        arrays = [t.data, ar.do("conj", t.data)]
-        k_input = []
-        b_input = []
-        m_inputs = []
-        for j, jx in enumerate(t.inds, 1):
-            k_input.append(j)
-
-            if jx == ind:
-                # output index -> take diagonal
-                output = (j,)
-                b_input.append(j)
-            else:
-                try:
-                    # partial trace with message
-                    m = self.messages[jx, tid]
-                    arrays.append(m)
-                    b_input.append(-j)
-                    m_inputs.append((-j, j))
-                except KeyError:
-                    # direct partial trace
-                    b_input.append(j)
-
-        p = qtn.array_contract(
-            arrays,
-            inputs=(tuple(k_input), tuple(b_input), *m_inputs),
-            output=output,
-            **self.contract_opts,
-        )
-        p = ar.do("real", p)
-        return p / ar.do("sum", p)
-
-    def contract(self, strip_exponent=False):
-        """Estimate the total contraction, i.e. the 2-norm.
-
-        Parameters
-        ----------
-        strip_exponent : bool, optional
-            Whether to strip the exponent from the final result. If ``True``
-            then the returned result is ``(mantissa, exponent)``.
-
-        Returns
-        -------
-        scalar or (scalar, float)
-        """
-        tvals = []
-
-        for tid, t in self.tn.tensor_map.items():
-            arrays = [t.data, ar.do("conj", t.data)]
-            k_input = []
-            b_input = []
-            m_inputs = []
-            for i, ix in enumerate(t.inds, 1):
-                k_input.append(i)
-                if ix in self.output_inds:
-                    b_input.append(i)
-                else:
-                    b_input.append(-i)
-                    m_inputs.append((-i, i))
-                    arrays.append(self.messages[ix, tid])
-
-            inputs = (tuple(k_input), tuple(b_input), *m_inputs)
-            output = ()
-            tval = qtn.array_contract(
-                arrays, inputs, output, **self.contract_opts
-            )
-            tvals.append(tval)
-
-        mvals = []
-        for ix, tids in self.tn.ind_map.items():
-            if ix in self.output_inds:
-                continue
-            tida, tidb = tids
-            ml = self.messages[ix, tidb]
-            mr = self.messages[ix, tida]
-            mval = qtn.array_contract(
-                (ml, mr), ((1, 2), (1, 2)), (), **self.contract_opts
-            )
-            mvals.append(mval)
-
-        return combine_local_contractions(
-            tvals, mvals, self.backend, strip_exponent=strip_exponent
-        )
-
-    def compress(
-        self,
-        max_bond,
-        cutoff=0.0,
-        cutoff_mode=4,
-        renorm=0,
-        inplace=False,
-    ):
-        """Compress the initial tensor network using the current messages."""
-        tn = self.tn if inplace else self.tn.copy()
-
-        for ix, tids in tn.ind_map.items():
-            if len(tids) != 2:
-                continue
-            tida, tidb = tids
-
-            # messages are left and right factors squared already
-            ta = tn.tensor_map[tida]
-            dm = ta.ind_size(ix)
-            dl = ta.size // dm
-            ml = self.messages[ix, tidb]
-            Rl = qtn.decomp.squared_op_to_reduced_factor(
-                ml, dl, dm, right=True
-            )
-
-            tb = tn.tensor_map[tidb]
-            dr = tb.size // dm
-            mr = self.messages[ix, tida].T
-            Rr = qtn.decomp.squared_op_to_reduced_factor(
-                mr, dm, dr, right=False
-            )
-
-            # compute the compressors
-            Pl, Pr = qtn.decomp.compute_oblique_projectors(
-                Rl,
-                Rr,
-                max_bond=max_bond,
-                cutoff=cutoff,
-                cutoff_mode=cutoff_mode,
-                renorm=renorm,
-            )
-
-            # contract the compressors into the tensors
-            tn.tensor_map[tida].gate_(Pl.T, ix)
-            tn.tensor_map[tidb].gate_(Pr, ix)
-
-            # update messages with projections
-            if inplace:
-                new_Ra = Rl @ Pl
-                new_Rb = Pr @ Rr
-                self.messages[ix, tidb] = ar.dag(new_Ra) @ new_Ra
-                self.messages[ix, tida] = new_Rb @ ar.dag(new_Rb)
-
-        return tn
-
-
-def contract_d2bp(
-    tn,
-    messages=None,
-    output_inds=None,
-    optimize="auto-hq",
-    damping=0.0,
-    update="sequential",
-    local_convergence=True,
-    max_iterations=1000,
-    tol=5e-6,
-    strip_exponent=False,
-    info=None,
-    progbar=False,
-    **contract_opts,
-):
-    """Estimate the norm squared of ``tn`` using dense 2-norm belief
-    propagation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to form the 2-norm of and run BP on.
-    messages : dict[(str, int), array_like], optional
-        The initial messages to use, effectively defaults to all ones if not
-        specified.
-    max_iterations : int, optional
-        The maximum number of iterations to perform.
-    tol : float, optional
-        The convergence tolerance for messages.
-    output_inds : set[str], optional
-        The indices to consider as output (dangling) indices of the tn.
-        Computed automatically if not specified.
-    optimize : str or PathOptimizer, optional
-        The path optimizer to use when contracting the messages.
-    damping : float, optional
-        The damping parameter to use, defaults to no damping.
-    update : {'parallel', 'sequential'}, optional
-        Whether to update all messages in parallel or sequentially.
-    local_convergence : bool, optional
-        Whether to allow messages to locally converge - i.e. if all their
-        input messages have converged then stop updating them.
-    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.
-    contract_opts
-        Other options supplied to ``cotengra.array_contract``.
-
-    Returns
-    -------
-    scalar or (scalar, float)
-    """
-    bp = D2BP(
-        tn,
-        messages=messages,
-        output_inds=output_inds,
-        optimize=optimize,
-        damping=damping,
-        local_convergence=local_convergence,
-        update=update,
-        **contract_opts,
-    )
-    bp.run(
-        max_iterations=max_iterations,
-        tol=tol,
-        info=info,
-        progbar=progbar,
-    )
-    return bp.contract(strip_exponent=strip_exponent)
-
-
-def compress_d2bp(
-    tn,
-    max_bond,
-    cutoff=0.0,
-    cutoff_mode="rsum2",
-    renorm=0,
-    messages=None,
-    output_inds=None,
-    optimize="auto-hq",
-    damping=0.0,
-    update="sequential",
-    local_convergence=True,
-    max_iterations=1000,
-    tol=5e-6,
-    inplace=False,
-    info=None,
-    progbar=False,
-    **contract_opts,
-):
-    """Compress the tensor network ``tn`` using dense 2-norm belief
-    propagation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to form the 2-norm of, run BP on and then compress.
-    max_bond : int
-        The maximum bond dimension to compress to.
-    cutoff : float, optional
-        The cutoff to use when compressing.
-    cutoff_mode : int, optional
-        The cutoff mode to use when compressing.
-    messages : dict[(str, int), array_like], optional
-        The initial messages to use, effectively defaults to all ones if not
-        specified.
-    max_iterations : int, optional
-        The maximum number of iterations to perform.
-    tol : float, optional
-        The convergence tolerance for messages.
-    output_inds : set[str], optional
-        The indices to consider as output (dangling) indices of the tn.
-        Computed automatically if not specified.
-    optimize : str or PathOptimizer, optional
-        The path optimizer to use when contracting the messages.
-    damping : float, optional
-        The damping parameter to use, defaults to no damping.
-    update : {'parallel', 'sequential'}, optional
-        Whether to update all messages in parallel or sequentially.
-    local_convergence : bool, optional
-        Whether to allow messages to locally converge - i.e. if all their
-        input messages have converged then stop updating them.
-    inplace : bool, optional
-        Whether to perform the compression inplace.
-    info : dict, optional
-        If specified, update this dictionary with information about the
-        belief propagation run.
-    progbar : bool, optional
-        Whether to show a progress bar.
-    contract_opts
-        Other options supplied to ``cotengra.array_contract``.
-
-    Returns
-    -------
-    TensorNetwork
-    """
-    bp = D2BP(
-        tn,
-        messages=messages,
-        output_inds=output_inds,
-        optimize=optimize,
-        damping=damping,
-        update=update,
-        local_convergence=local_convergence,
-        **contract_opts,
-    )
-    bp.run(
-        max_iterations=max_iterations,
-        tol=tol,
-        info=info,
-        progbar=progbar,
-    )
-    return bp.compress(
-        max_bond=max_bond,
-        cutoff=cutoff,
-        cutoff_mode=cutoff_mode,
-        renorm=renorm,
-        inplace=inplace,
-    )
-
-
-def sample_d2bp(
-    tn,
-    output_inds=None,
-    messages=None,
-    max_iterations=100,
-    tol=1e-2,
-    bias=None,
-    seed=None,
-    local_convergence=True,
-    progbar=False,
-    **contract_opts,
-):
-    """Sample a configuration from ``tn`` using dense 2-norm belief
-    propagation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to sample from.
-    output_inds : set[str], optional
-        Which indices to sample.
-    messages : dict[(str, int), array_like], optional
-        The initial messages to use, effectively defaults to all ones if not
-        specified.
-    max_iterations : int, optional
-        The maximum number of iterations to perform, per marginal.
-    tol : float, optional
-        The convergence tolerance for messages.
-    bias : float, optional
-        Bias the sampling towards more locally likely bit-strings. This is
-        done by raising the probability of each bit-string to this power.
-    seed : int, optional
-        A random seed for reproducibility.
-    local_convergence : bool, optional
-        Whether to allow messages to locally converge - i.e. if all their
-        input messages have converged then stop updating them.
-    progbar : bool, optional
-        Whether to show a progress bar.
-    contract_opts
-        Other options supplied to ``cotengra.array_contract``.
-
-    Returns
-    -------
-    config : dict[str, int]
-        The sampled configuration, a mapping of output indices to values.
-    tn_config : TensorNetwork
-        The tensor network with the sampled configuration applied.
-    omega : float
-        The BP probability of the sampled configuration.
-    """
-    import numpy as np
-
-    if output_inds is None:
-        output_inds = tn.outer_inds()
-
-    rng = np.random.default_rng(seed)
-    config = {}
-    omega = 1.0
-
-    tn = tn.copy()
-    bp = D2BP(
-        tn,
-        messages=messages,
-        local_convergence=local_convergence,
-        **contract_opts,
-    )
-    bp.run(max_iterations=max_iterations, tol=tol)
-
-    marginals = dict.fromkeys(output_inds)
-
-    if progbar:
-        import tqdm
-
-        pbar = tqdm.tqdm(total=len(marginals))
-    else:
-        pbar = None
-
-    while marginals:
-        for ix in marginals:
-            marginals[ix] = bp.compute_marginal(ix)
-
-        ix, p = max(marginals.items(), key=lambda x: max(x[1]))
-        p = ar.to_numpy(p)
-
-        if bias is not None:
-            # bias distribution towards more locally likely bit-strings
-            p = p**bias
-            p /= np.sum(p)
-
-        v = rng.choice([0, 1], p=p)
-        config[ix] = v
-        del marginals[ix]
-
-        tids = tuple(tn.ind_map[ix])
-        tn.isel_({ix: v})
-
-        omega *= p[v]
-        if progbar:
-            pbar.update(1)
-            pbar.set_description(f"{ix}->{v}", refresh=False)
-
-        bp = D2BP(
-            tn,
-            messages=bp.messages,
-            local_convergence=local_convergence,
-            **contract_opts,
-        )
-        bp.update_touched_from_tids(*tids)
-        bp.run(tol=tol, max_iterations=max_iterations)
-
-    if progbar:
-        pbar.close()
-
-    return config, tn, omega