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

trajectree/quimb/quimb/experimental/belief_propagation/l2bp.py

@@ -1,537 +0,0 @@
-import math
-
-import autoray as ar
-
-import quimb.tensor as qtn
-from quimb.utils import oset
-
-from .bp_common import (
-    BeliefPropagationCommon,
-    combine_local_contractions,
-    create_lazy_community_edge_map,
-)
-
-
-class L2BP(BeliefPropagationCommon):
-    """Lazy (as in multiple uncontracted tensors per site) 2-norm (as in for
-    wavefunctions and operators) belief propagation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to form the 2-norm of and run BP on.
-    site_tags : sequence of str, optional
-        The tags identifying the sites in ``tn``, each tag forms a region,
-        which should not overlap. If the tensor network is structured, then
-        these are inferred automatically.
-    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.
-    optimize : str or PathOptimizer, optional
-        The path optimizer to use when contracting the messages.
-    contract_opts
-        Other options supplied to ``cotengra.array_contract``.
-    """
-
-    def __init__(
-        self,
-        tn,
-        site_tags=None,
-        damping=0.0,
-        update="sequential",
-        local_convergence=True,
-        optimize="auto-hq",
-        **contract_opts,
-    ):
-        self.backend = next(t.backend for t in tn)
-        self.damping = damping
-        self.local_convergence = local_convergence
-        self.update = update
-        self.optimize = optimize
-        self.contract_opts = contract_opts
-
-        if site_tags is None:
-            self.site_tags = tuple(tn.site_tags)
-        else:
-            self.site_tags = tuple(site_tags)
-
-        (
-            self.edges,
-            self.neighbors,
-            self.local_tns,
-            self.touch_map,
-        ) = create_lazy_community_edge_map(tn, site_tags)
-        self.touched = oset()
-
-        _abs = ar.get_lib_fn(self.backend, "abs")
-        _sum = ar.get_lib_fn(self.backend, "sum")
-        _transpose = ar.get_lib_fn(self.backend, "transpose")
-        _conj = ar.get_lib_fn(self.backend, "conj")
-
-        def _normalize(x):
-            return x / _sum(x)
-
-        def _symmetrize(x):
-            N = ar.ndim(x)
-            perm = (*range(N // 2, N), *range(0, N // 2))
-            return x + _conj(_transpose(x, perm))
-
-        def _distance(x, y):
-            return _sum(_abs(x - y))
-
-        self._normalize = _normalize
-        self._symmetrize = _symmetrize
-        self._distance = _distance
-
-        # initialize messages
-        self.messages = {}
-
-        for pair, bix in self.edges.items():
-            cix = tuple(ix + "_l2bp*" for ix in bix)
-            remapper = dict(zip(bix, cix))
-            output_inds = cix + bix
-
-            # compute leftwards and righwards messages
-            for i, j in (sorted(pair), sorted(pair, reverse=True)):
-                tn_i = self.local_tns[i]
-                tn_i2 = tn_i & tn_i.conj().reindex_(remapper)
-                tm = tn_i2.contract(
-                    all,
-                    output_inds=output_inds,
-                    optimize=self.optimize,
-                    drop_tags=True,
-                    **self.contract_opts,
-                )
-                tm.modify(apply=self._symmetrize)
-                tm.modify(apply=self._normalize)
-                self.messages[i, j] = tm
-
-        # initialize contractions
-        self.contraction_tns = {}
-        for pair, bix in self.edges.items():
-            for i, j in (sorted(pair), sorted(pair, reverse=True)):
-                # form the ket side and messages
-                tn_i_left = self.local_tns[i]
-                # get other incident nodes which aren't j
-                ks = [k for k in self.neighbors[i] if k != j]
-                tks = [self.messages[k, i] for k in ks]
-
-                # form the 'bra' side
-                tn_i_right = tn_i_left.conj()
-                # get the bonds that attach the bra to messages
-                outer_bix = {
-                    ix for k in ks for ix in self.edges[tuple(sorted((k, i)))]
-                }
-                # need to reindex to join message bonds, and create bra outputs
-                remapper = {}
-                for ix in tn_i_right.ind_map:
-                    if ix in bix:
-                        # bra outputs
-                        remapper[ix] = ix + "_l2bp**"
-                    elif ix in outer_bix:
-                        # messages connected
-                        remapper[ix] = ix + "_l2bp*"
-                    # remaining indices are either internal and will be mangled
-                    # or global outer indices and will be contracted directly
-
-                tn_i_right.reindex_(remapper)
-
-                self.contraction_tns[i, j] = qtn.TensorNetwork(
-                    (tn_i_left, *tks, tn_i_right), virtual=True
-                )
-
-    def iterate(self, tol=5e-6):
-        if (not self.local_convergence) or (not self.touched):
-            # assume if asked to iterate that we want to check all messages
-            self.touched.update(
-                pair for edge in self.edges for pair in (edge, edge[::-1])
-            )
-
-        ncheck = len(self.touched)
-        nconv = 0
-        max_mdiff = -1.0
-        new_touched = oset()
-
-        def _compute_m(key):
-            i, j = key
-            bix = self.edges[(i, j) if i < j else (j, i)]
-            cix = tuple(ix + "_l2bp**" for ix in bix)
-            output_inds = cix + bix
-
-            tn_i_to_j = self.contraction_tns[i, j]
-
-            tm_new = tn_i_to_j.contract(
-                all,
-                output_inds=output_inds,
-                drop_tags=True,
-                optimize=self.optimize,
-                **self.contract_opts,
-            )
-            tm_new.modify(apply=self._symmetrize)
-            tm_new.modify(apply=self._normalize)
-            return tm_new.data
-
-        def _update_m(key, data):
-            nonlocal nconv, max_mdiff
-
-            tm = self.messages[key]
-
-            if self.damping > 0.0:
-                data = (1 - self.damping) * data + self.damping * tm.data
-
-            try:
-                mdiff = float(self._distance(tm.data, data))
-            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)
-            tm.modify(data=data)
-
-        if self.update == "parallel":
-            new_data = {}
-            # compute all new messages
-            while self.touched:
-                key = self.touched.pop()
-                new_data[key] = _compute_m(key)
-            # insert all new messages
-            for key, data in new_data.items():
-                _update_m(key, data)
-
-        elif self.update == "sequential":
-            # compute each new message and immediately re-insert it
-            while self.touched:
-                key = self.touched.pop()
-                data = _compute_m(key)
-                _update_m(key, data)
-
-        self.touched = new_touched
-
-        return nconv, ncheck, max_mdiff
-
-    def normalize_messages(self):
-        """Normalize all messages such that for each bond `<m_i|m_j> = 1` and
-        `<m_i|m_i> = <m_j|m_j>` (but in general != 1).
-        """
-        for i, j in self.edges:
-            tmi = self.messages[i, j]
-            tmj = self.messages[j, i]
-            nij = (tmi @ tmj)**0.5
-            nii = (tmi @ tmi)**0.25
-            njj = (tmj @ tmj)**0.25
-            tmi /= (nij * nii / njj)
-            tmj /= (nij * njj / nii)
-
-    def contract(self, strip_exponent=False):
-        """Estimate the contraction of the norm squared using the current
-        messages.
-        """
-        tvals = []
-        for i, ket in self.local_tns.items():
-            # we allow missing keys here for tensors which are just
-            # disconnected but still appear in local_tns
-            ks = self.neighbors.get(i, ())
-            bix = [ix for k in ks for ix in self.edges[tuple(sorted((k, i)))]]
-            bra = ket.H.reindex_({ix: ix + "_l2bp*" for ix in bix})
-            tni = qtn.TensorNetwork(
-                (
-                    ket,
-                    *(self.messages[k, i] for k in ks),
-                    bra,
-                )
-            )
-            tvals.append(
-                tni.contract(all, optimize=self.optimize, **self.contract_opts)
-            )
-
-        mvals = []
-        for i, j in self.edges:
-            mvals.append(
-                (self.messages[i, j] & self.messages[j, i]).contract(
-                    all,
-                    optimize=self.optimize,
-                    **self.contract_opts,
-                )
-            )
-
-        return combine_local_contractions(
-            tvals, mvals, self.backend, strip_exponent=strip_exponent
-        )
-
-    def partial_trace(
-        self,
-        site,
-        normalized=True,
-        optimize="auto-hq",
-    ):
-        example_tn = next(tn for tn in self.local_tns.values())
-
-        site_tag = example_tn.site_tag(site)
-        ket_site_ind = example_tn.site_ind(site)
-
-        ks = self.neighbors[site_tag]
-        tn_rho_i = self.local_tns[site_tag].copy()
-        tn_bra_i = tn_rho_i.H
-
-        for k in ks:
-            tn_rho_i &= self.messages[k, site_tag]
-
-        outer_bix = {
-            ix for k in ks for ix in self.edges[tuple(sorted((k, site_tag)))]
-        }
-
-        ind_changes = {}
-        for ix in tn_bra_i.ind_map:
-            if ix == ket_site_ind:
-                # open up the site index
-                bra_site_ind = ix + "_l2bp**"
-                ind_changes[ix] = bra_site_ind
-            if ix in outer_bix:
-                # attach bra message indices
-                ind_changes[ix] = ix + "_l2bp*"
-        tn_bra_i.reindex_(ind_changes)
-
-        tn_rho_i &= tn_bra_i
-
-        rho_i = tn_rho_i.to_dense(
-            [ket_site_ind],
-            [bra_site_ind],
-            optimize=optimize,
-            **self.contract_opts,
-        )
-        if normalized:
-            rho_i = rho_i / ar.do("trace", rho_i)
-
-        return rho_i
-
-    def compress(
-        self,
-        tn,
-        max_bond=None,
-        cutoff=5e-6,
-        cutoff_mode="rsum2",
-        renorm=0,
-        lazy=False,
-    ):
-        """Compress the state ``tn``, assumed to matched this L2BP instance,
-        using the messages stored.
-        """
-        for (i, j), bix in self.edges.items():
-            tml = self.messages[i, j]
-            tmr = self.messages[j, i]
-
-            bix_sizes = [tml.ind_size(ix) for ix in bix]
-            dm = math.prod(bix_sizes)
-
-            ml = ar.reshape(tml.data, (dm, dm))
-            dl = self.local_tns[i].outer_size() // dm
-            Rl = qtn.decomp.squared_op_to_reduced_factor(
-                ml, dl, dm, right=True
-            )
-
-            mr = ar.reshape(tmr.data, (dm, dm)).T
-            dr = self.local_tns[j].outer_size() // dm
-            Rr = qtn.decomp.squared_op_to_reduced_factor(
-                mr, dm, dr, right=False
-            )
-
-            Pl, Pr = qtn.decomp.compute_oblique_projectors(
-                Rl,
-                Rr,
-                cutoff_mode=cutoff_mode,
-                renorm=renorm,
-                max_bond=max_bond,
-                cutoff=cutoff,
-            )
-
-            Pl = ar.do("reshape", Pl, (*bix_sizes, -1))
-            Pr = ar.do("reshape", Pr, (-1, *bix_sizes))
-
-            ltn = tn.select(i)
-            rtn = tn.select(j)
-
-            new_lix = [qtn.rand_uuid() for _ in bix]
-            new_rix = [qtn.rand_uuid() for _ in bix]
-            new_bix = [qtn.rand_uuid()]
-            ltn.reindex_(dict(zip(bix, new_lix)))
-            rtn.reindex_(dict(zip(bix, new_rix)))
-
-            # ... and insert the new projectors in place
-            tn |= qtn.Tensor(Pl, inds=new_lix + new_bix, tags=(i,))
-            tn |= qtn.Tensor(Pr, inds=new_bix + new_rix, tags=(j,))
-
-        if not lazy:
-            for st in self.site_tags:
-                try:
-                    tn.contract_tags_(
-                        st, optimize=self.optimize, **self.contract_opts
-                    )
-                except KeyError:
-                    pass
-
-        return tn
-
-
-def contract_l2bp(
-    tn,
-    site_tags=None,
-    damping=0.0,
-    update="sequential",
-    local_convergence=True,
-    optimize="auto-hq",
-    max_iterations=1000,
-    tol=5e-6,
-    strip_exponent=False,
-    info=None,
-    progbar=False,
-    **contract_opts,
-):
-    """Estimate the norm squared of ``tn`` using lazy belief propagation.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to estimate the norm squared of.
-    site_tags : sequence of str, optional
-        The tags identifying the sites in ``tn``, each tag forms a region.
-    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.
-    optimize : str or PathOptimizer, optional
-        The contraction strategy to use.
-    max_iterations : int, optional
-        The maximum number of iterations to perform.
-    tol : float, optional
-        The convergence tolerance for messages.
-    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``.
-    """
-    bp = L2BP(
-        tn,
-        site_tags=site_tags,
-        damping=damping,
-        update=update,
-        local_convergence=local_convergence,
-        optimize=optimize,
-        **contract_opts,
-    )
-    bp.run(
-        max_iterations=max_iterations,
-        tol=tol,
-        info=info,
-        progbar=progbar,
-    )
-    return bp.contract(strip_exponent=strip_exponent)
-
-
-def compress_l2bp(
-    tn,
-    max_bond,
-    cutoff=0.0,
-    cutoff_mode="rsum2",
-    max_iterations=1000,
-    tol=5e-6,
-    site_tags=None,
-    damping=0.0,
-    update="sequential",
-    local_convergence=True,
-    optimize="auto-hq",
-    lazy=False,
-    inplace=False,
-    info=None,
-    progbar=False,
-    **contract_opts,
-):
-    """Compress ``tn`` using lazy belief propagation, producing a tensor
-    network with a single tensor per site.
-
-    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.
-    max_iterations : int, optional
-        The maximum number of iterations to perform.
-    tol : float, optional
-        The convergence tolerance for messages.
-    site_tags : sequence of str, optional
-        The tags identifying the sites in ``tn``, each tag forms a region. If
-        the tensor network is structured, then these are inferred
-        automatically.
-    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.
-    optimize : str or PathOptimizer, optional
-        The path optimizer to use when contracting the messages.
-    lazy : bool, optional
-        Whether to perform the compression lazily, i.e. to leave the computed
-        compression projectors uncontracted.
-    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
-    """
-    tnc = tn if inplace else tn.copy()
-
-    bp = L2BP(
-        tnc,
-        site_tags=site_tags,
-        damping=damping,
-        update=update,
-        local_convergence=local_convergence,
-        optimize=optimize,
-        **contract_opts,
-    )
-    bp.run(
-        max_iterations=max_iterations,
-        tol=tol,
-        info=info,
-        progbar=progbar,
-    )
-    return bp.compress(
-        tnc,
-        max_bond=max_bond,
-        cutoff=cutoff,
-        cutoff_mode=cutoff_mode,
-        lazy=lazy,
-    )