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

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

@@ -0,0 +1,537 @@
+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,
+    )