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

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

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

@@ -0,0 +1,286 @@
+"""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