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

trajectree/quimb/quimb/experimental/belief_propagation/regions.py

@@ -1,194 +0,0 @@
-class RegionGraph:
-    def __init__(self, regions=(), autocomplete=True):
-        self.lookup = {}
-        self.parents = {}
-        self.children = {}
-        self.counts = {}
-        for region in regions:
-            self.add_region(region)
-        if autocomplete:
-            self.autocomplete()
-
-    @property
-    def regions(self):
-        return tuple(self.children)
-
-    def neighbor_regions(self, region):
-        """Get all regions that intersect with the given region."""
-        region = frozenset(region)
-
-        other_regions = set.union(*(self.lookup[node] for node in region))
-        other_regions.discard(region)
-        return other_regions
-
-    def add_region(self, region):
-        """Add a new region and update parent-child relationships.
-
-        Parameters
-        ----------
-        region : Sequence[Hashable]
-            The new region to add.
-        """
-        region = frozenset(region)
-
-        if region in self.parents:
-            # already added
-            return
-
-        # populate data structures
-        self.parents[region] = set()
-        self.children[region] = set()
-        for node in region:
-            # collect regions that contain nodes for fast neighbor lookup
-            self.lookup.setdefault(node, set()).add(region)
-
-        # add parent-child relationships
-        for other in self.neighbor_regions(region):
-            if region.issubset(other):
-                self.parents[region].add(other)
-                self.children[other].add(region)
-            elif other.issubset(region):
-                self.children[region].add(other)
-                self.parents[other].add(region)
-
-        # prune redundant parents and children
-        children = sorted(self.children[region], key=len)
-        for i, c in enumerate(children):
-            if any(c.issubset(cc) for cc in children[i + 1 :]):
-                # child is a subset of larger child -> remove
-                self.children[region].remove(c)
-                self.parents[c].remove(region)
-
-        parents = sorted(self.parents[region], key=len, reverse=True)
-        for i, p in enumerate(parents):
-            if any(p.issuperset(pp) for pp in parents[i + 1 :]):
-                # parent is a superset of smaller parent -> remove
-                self.parents[region].remove(p)
-                self.children[p].remove(region)
-
-        self.counts.clear()
-
-    def autocomplete(self):
-        """Add all missing intersecting sub-regions."""
-        for r in self.regions:
-            for other in self.neighbor_regions(r):
-                self.add_region(r & other)
-
-    def autoextend(self, regions=None):
-        """Extend this region graph upwards by adding in all pairwise unions of
-        regions. If regions is specified, take this as one set of pairs.
-        """
-        if regions is None:
-            regions = self.regions
-
-        neighbors = {}
-        for r in regions:
-            for other in self.neighbor_regions(r):
-                neighbors.setdefault(r, []).append(other)
-
-        for r, others in neighbors.items():
-            for other in others:
-                self.add_region(r | other)
-
-    def get_parents(self, region):
-        """Get all ancestors that contain the given region, but do not contain
-        any other regions that themselves contain the given region.
-        """
-        return self.parents[region]
-
-    def get_children(self, region):
-        """Get all regions that are contained by the given region, but are not
-        contained by any other descendents of the given region.
-        """
-        return self.children[region]
-
-    def get_ancestors(self, region):
-        """Get all regions that contain the given region, not just direct
-        parents.
-        """
-        seen = set()
-        queue = [region]
-        while queue:
-            r = queue.pop()
-            for rp in self.parents[r]:
-                if rp not in seen:
-                    seen.add(rp)
-                    queue.append(rp)
-        return seen
-
-    def get_descendents(self, region):
-        """Get all regions that are contained by the given region, not just
-        direct children.
-        """
-        seen = set()
-        queue = [region]
-        while queue:
-            r = queue.pop()
-            for rc in self.children[r]:
-                if rc not in seen:
-                    seen.add(rc)
-                    queue.append(rc)
-        return seen
-
-    def get_count(self, region):
-        """Get the count of the given region, i.e. the correct weighting to
-        apply when summing over all regions to avoid overcounting.
-        """
-        try:
-            C = self.counts[region]
-        except KeyError:
-            # n.b. cache is cleared when any new region is added
-            C = self.counts[region] = 1 - sum(
-                self.get_count(a) for a in self.get_ancestors(region)
-            )
-        return C
-
-    def get_total_count(self):
-        return sum(map(self.get_count, self.regions))
-
-    def get_level(self, region):
-        """Get the level of the given region, i.e. the distance to an ancestor
-        with no parents.
-        """
-        if not self.parents[region]:
-            return 0
-        else:
-            return min(self.get_level(p) for p in self.parents[region]) - 1
-
-    def draw(self, pos=None, a=20, scale=1.0, radius=0.1, **drawing_opts):
-        from quimb.schematic import Drawing, hash_to_color
-
-        if pos is None:
-            pos = {node: node for node in self.lookup}
-
-        def get_draw_pos(coo):
-            return tuple(scale * s for s in pos[coo])
-
-        sizes = {len(r) for r in self.regions}
-        levelmap = {s: i for i, s in enumerate(sorted(sizes))}
-
-        d = Drawing(a=a, **drawing_opts)
-        for region in sorted(self.regions, key=len, reverse=True):
-            # level = self.get_level(region)
-            # level = len(region)
-            level = levelmap[len(region)]
-
-            coos = [(*get_draw_pos(coo), 2.0 * level) for coo in region]
-
-            d.patch_around(
-                coos,
-                radius=radius,
-                # edgecolor=hash_to_color(str(region)),
-                facecolor=hash_to_color(str(region)),
-                alpha=1 / 3,
-                linestyle="",
-                linewidth=3,
-            )
-
-        return d.fig, d.ax
-
-    def __repr__(self):
-        return (
-            f"<RegionGraph(regions={len(self.regions)}, "
-            f"total_count={self.get_total_count()})>"
-        )

trajectree/quimb/quimb/experimental/cluster_update.py

@@ -1,286 +0,0 @@
-"""Implementation of arbitrary geometry wavefunction cluster update.
-"""
-import functools
-
-from quimb.tensor.tensor_core import ensure_dict, bonds, bonds_size, do
-from quimb.tensor.tensor_arbgeom_tebd import SimpleUpdateGen
-
-
-def gate_inds_nn_fit(
-    self,
-    G,
-    ind1,
-    ind2,
-    max_bond=None,
-    method="als",
-    pregauge=2,
-    init_simple_guess=True,
-    steps=10,
-    fit_opts=None,
-    contract_opts=None,
-    inplace=False,
-):
-    """Gate two nearest neighbor outer indices, using full fitting of
-    reduced tensors with respect to the environment. This is more accurate
-    than a simple reduced gate when restricting the bond dimension.
-
-    Parameters
-    ----------
-    G : array_like
-        The gate to fit.
-    ind1, ind2 : str
-        The indices to gate.
-    max_bond : int, optional
-        The maximum bond dimension to use. If ``None``, use the maximum
-        bond dimension that the tensors currently share.
-    method : {'als', 'autodiff'}, optional
-        The method to use for fitting.
-    pregauge : int, optional
-        How many times to locally canonize from the purified environment
-        tensor to both the left and right reduced tensors.
-    init_simple_guess : bool, optional
-        Whether to use a 'simple update' guess for the initial guess. This
-        can be quite high quality already if pregauging is used.
-    steps : int, optional
-        The number of steps to use for fitting, can be ``0`` in which case
-        the initial guess is used, which in conjuction with the
-        envinronment pregauging can still be quite high quality.
-    inplace : bool, optional
-        Whether to update the tensor network in place.
-    contract_opts
-        Supplied to
-        :meth:`~quimb.tensor.tensor_core.TensorNetwork.contract`.
-    """
-    fit_opts = ensure_dict(fit_opts)
-    contract_opts = ensure_dict(contract_opts)
-
-    ket = self.copy()
-
-    # move indices onto shared bond so environment can be contracted
-    ket.reduce_inds_onto_bond(
-        ind1, ind2, combine=False, ndim_cutoff=0, tags="__REDUCED__"
-    )
-
-    ket.add_tag("__KET__")
-    bra = ket.conj().retag_({"__KET__": "__BRA__"})
-    norm = ket | bra
-
-    # contract environment -> all but reduced tensors
-    norm.contract_tags_("__REDUCED__", "!any", **contract_opts)
-
-    # get tensors and bond names
-    (tide,) = norm._get_tids_from_tags("__REDUCED__", "!any")
-    (te,) = norm._tids_get(tide)
-    (tid_k1,) = norm._get_tids_from_inds(ind1) & norm._get_tids_from_tags(
-        "__KET__"
-    )
-    (tid_k2,) = norm._get_tids_from_inds(ind2) & norm._get_tids_from_tags(
-        "__KET__"
-    )
-    (tid_b1,) = norm._get_tids_from_inds(ind1) & norm._get_tids_from_tags(
-        "__BRA__"
-    )
-    (tid_b2,) = norm._get_tids_from_inds(ind2) & norm._get_tids_from_tags(
-        "__BRA__"
-    )
-    tk1 = norm.tensor_map[tid_k1]
-    tk2 = norm.tensor_map[tid_k2]
-    tb1 = norm.tensor_map[tid_b1]
-    tb2 = norm.tensor_map[tid_b2]
-
-    (ix_ek1,) = bonds(te, tk1)
-    (ix_ek2,) = bonds(te, tk2)
-    (ix_eb1,) = bonds(te, tb1)
-    (ix_eb2,) = bonds(te, tb2)
-
-    if max_bond is None:
-        max_bond = bonds_size(tk1, tk2)
-
-    # split environment to get purification
-    _, tek = te.split(
-        left_inds=[ix_eb1, ix_eb2],
-        right_inds=[ix_ek1, ix_ek2],
-        method="svd",
-    )
-
-    #   ┌────────────────┐
-    # ┌─┴┐      ┌──┐     │   :
-    # │k1├──────┤k2├───┐ │   : this is the purification
-    # └─┬┘      └─┬┘   ├─┴┐  : we'll fit to gated version of itself
-    #   │         │    │ek│  :
-    #   │ind1     │    └─┬┘
-    #   │     ind2│    ┌─┴┐
-    #   │         │    │eb│
-    # ┌─┴┐      ┌─┴┐   ├─┬┘
-    # │b1├──────┤b2├───┘ │
-    # └─┬┘      └──┘     │
-    #   └────────────────┘
-
-    for _ in range(int(pregauge)):
-        # perform some conditioning: local gauging from Q -> R tensors
-        for ind, ix_ek, tk in [
-            (ind1, ix_ek1, tk1),
-            (ind2, ix_ek2, tk2),
-        ]:
-            R = tek.split(
-                left_inds=None,
-                right_inds=[ix_ek],
-                method="qr",
-                get="arrays",
-            )[1]
-
-            Rinv = do("linalg.inv", R)
-            # get Q tensor in uncontracted TN
-            (tidkq,) = ket._get_tids_from_inds(
-                ix_ek
-            ) - ket._get_tids_from_inds(ind)
-            tkq = ket.tensor_map[tidkq]
-            # need to keep environment and bare tensors in sync
-            tkq.gate_(Rinv.T, ix_ek)
-            tek.gate_(Rinv.T, ix_ek)
-            tk.gate_(R, ix_ek)
-
-    # form purification
-    tnl = tek | tk1 | tk2
-    # form gated purification
-    tnl_target = tnl.gate_inds(G, (ind1, ind2), contract=True)
-
-    if init_simple_guess:
-        # maybe initialize with simple guess
-        tnl.gate_inds_(
-            G, (ind1, ind2), contract="split", max_bond=max_bond, cutoff=0.0
-        )
-
-    if steps:
-        # perform the actual fitting, specifying only the reduced tensors
-        tnl.fit_(
-            tnl_target,
-            tags=["__REDUCED__"],
-            method=method,
-            steps=steps,
-            **fit_opts,
-        )
-
-    # re-absorb reduced factors
-    ket.contract_ind(ix_ek1)
-    ket.contract_ind(ix_ek2)
-
-    # TN we will return
-    new = self if inplace else self.copy()
-    (tn1,) = new._inds_get(ind1)
-    (tn2,) = new._inds_get(ind2)
-
-    # permute to match original tensors
-    (t1,) = ket._inds_get(ind1)
-    (t2,) = ket._inds_get(ind2)
-    t1.transpose_like_(tn1)
-    t2.transpose_like_(tn2)
-
-    tn1.modify(data=t1.data)
-    tn2.modify(data=t2.data)
-
-    return new
-
-
-gate_inds_nn_fit_ = functools.partialmethod(gate_inds_nn_fit, inplace=True)
-
-# TensorNetwork.gate_inds_nn_fit =gate_inds_nn_fit
-# TensorNetwork.gate_inds_nn_fit_ = gate_inds_nn_fit
-
-
-class ClusterUpdateNNGen(SimpleUpdateGen):
-    """Cluster update for arbitrary geometry nearest neighbor hamiltonians.
-    This keeps track of simple update style gauges, in order to approximately
-    partial trace beyond ``cluster_radius`` and form an approximate environment
-    for two nearest neighbor sites that be used to fit the gate with higher
-    quality than simple update only.
-    """
-
-    def __init__(
-        self,
-        psi0,
-        ham,
-        tau=0.01,
-        D=None,
-        cluster_radius=1,
-        cluster_fillin=0,
-        gauge_smudge=1e-6,
-        imag=True,
-        gate_opts=None,
-        ordering=None,
-        second_order_reflect=False,
-        compute_energy_every=None,
-        compute_energy_final=True,
-        compute_energy_opts=None,
-        compute_energy_fn=None,
-        compute_energy_per_site=False,
-        callback=None,
-        keep_best=False,
-        progbar=True,
-    ):
-        super().__init__(
-            psi0=psi0,
-            ham=ham,
-            tau=tau,
-            D=D,
-            imag=imag,
-            gate_opts=gate_opts,
-            ordering=ordering,
-            second_order_reflect=second_order_reflect,
-            compute_energy_every=compute_energy_every,
-            compute_energy_final=compute_energy_final,
-            compute_energy_opts=compute_energy_opts,
-            compute_energy_fn=compute_energy_fn,
-            compute_energy_per_site=compute_energy_per_site,
-            callback=callback,
-            keep_best=keep_best,
-            progbar=progbar,
-        )
-        self.cluster_radius = cluster_radius
-        self.cluster_fillin = cluster_fillin
-        self.gauge_smudge = gauge_smudge
-
-        # override some default TEBDGen gate_opts
-        self.gate_opts.pop("cutoff")
-        self.gate_opts.pop("contract")
-
-    def gate(self, U, where):
-        taga, tagb = self._psi.gen_tags_from_coos(where)
-        inda, indb = self._psi.gen_inds_from_coos(where)
-
-        # get the local cluster
-        psi_local = self._psi.select_local(
-            (taga, tagb),
-            "any",
-            max_distance=self.cluster_radius,
-            fillin=self.cluster_fillin,
-            virtual=True,
-        )
-
-        # temporarily gauge it with 'simple' gauges
-        with psi_local.gauge_simple_temp(
-            self.gauges,
-            smudge=self.gauge_smudge,
-        ):
-            # fit the gate to the gauged local cluster
-            gate_inds_nn_fit(
-                psi_local,
-                U,
-                inda,
-                indb,
-                **self.gate_opts,
-            )
-
-        # update nearest gauges for the modified tensors
-        self._psi.gauge_local_(
-            (taga, tagb),
-            "any",
-            max_distance=1,
-            method="simple",
-            gauges=self.gauges,
-            smudge=self.gauge_smudge,
-        )
-
-        # perform some conditioning
-        self._psi.equalize_norms_(1.0)
-        self._psi.exponent = 0.0