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

trajectree/quimb/quimb/tensor/tensor_mera.py

@@ -0,0 +1,204 @@
+"""Tools for building and computing with MERA.
+"""
+
+from math import log2
+import itertools
+
+import numpy as np
+
+import quimb as qu
+from .tensor_core import rand_uuid, IsoTensor, TensorNetwork
+from .tensor_1d import TensorNetwork1D, TensorNetwork1DVector
+
+
+def is_power_of_2(x):
+    return ((x & (x - 1)) == 0) and x > 0
+
+
+class MERA(TensorNetwork1DVector,
+           TensorNetwork1D,
+           TensorNetwork):
+    r"""The Multi-scale Entanglement Renormalization Ansatz (MERA) state::
+
+            ...     ...     ...     ...     ...     ...
+         |       |       |       |       |       |       |
+        ISO     ISO     ISO     ISO     ISO     ISO     ISO   :
+           \   /   \   /   \   /   \   /   \   /   \   /      : '_LAYER1'
+            UNI     UNI     UNI     UNI     UNI     UNI       :
+           /   \   /   \   /   \   /   \   /   \   /   \
+        O ISO ISO ISO ISO ISO ISO ISO ISO ISO ISO ISO ISO I   :
+        | | | | | | | | | | | | | | | | | | | | | | | | | |   : '_LAYER0'
+        UNI UNI UNI UNI UNI UNI UNI UNI UNI UNI UNI UNI UNI   :
+        | | | | | | | | | | | | | | | | | | | | | | | | | |  <-- phys_dim
+        0 1 2 3 4 ....                            ... L-2 L-1
+
+    Parameters
+    ----------
+    L : int
+        The number of phyiscal sites. Shoule be a power of 2.
+    uni : array or sequence of arrays of shape (d, d, d, d).
+        The unitary operator(s). These will be cycled over and placed from
+        bottom left to top right in diagram above.
+    iso : array or sequence of arrays of shape (d, d, d)
+        The isometry operator(s). These will be cycled over and placed from
+        bottom left to top right in diagram above.
+    phys_dim : int, optional
+        The dimension of the local hilbert space.
+    dangle : bool, optional
+        Whether to leave a dangling index on the final isometry, in order to
+        maintain perfect scale invariance, else join the final unitaries just
+        with an indentity.
+    """
+
+    _EXTRA_PROPS = ('_site_ind_id', '_site_tag_id', 'cyclic', '_L')
+    _CONTRACT_STRUCTURED = False
+
+    def __init__(self, L, uni=None, iso=None, phys_dim=2, dangle=False,
+                 site_ind_id="k{}", site_tag_id="I{}", **tn_opts):
+
+        # short-circuit for copying MERA
+        if isinstance(L, MERA):
+            super().__init__(L)
+            for ep in MERA._EXTRA_PROPS:
+                setattr(self, ep, getattr(L, ep))
+            return
+
+        self._site_ind_id = site_ind_id
+        self._site_tag_id = site_tag_id
+        self.cyclic = True
+        self._L = L
+
+        if not is_power_of_2(L):
+            raise ValueError("``L`` should be a power of 2.")
+
+        nlayers = round(log2(L))
+
+        if hasattr(uni, 'shape'):
+            uni = (uni,)
+
+        if hasattr(iso, 'shape'):
+            iso = (iso,)
+
+        unis = itertools.cycle(uni)
+        isos = itertools.cycle(iso)
+
+        def gen_mera_tensors():
+            u_ind_id = site_ind_id
+
+            for i in range(nlayers):
+
+                # index id connecting to layer below
+                l_ind_id = u_ind_id
+                # index id connecting to isos to unis
+                m_ind_id = rand_uuid() + "_{}"
+                # index id connecting to layer above
+                u_ind_id = rand_uuid() + "_{}"
+
+                # number of tensor sites in this layer
+                eff_L = L // 2**i
+
+                for j in range(0, eff_L, 2):
+
+                    # generate the unitary:
+                    #  ul | | ur
+                    #     UNI
+                    #  ll | | lr
+                    #     j j+1
+                    ll, lr = map(l_ind_id.format, (j, (j + 1) % eff_L))
+                    ul, ur = map(m_ind_id.format, (j, (j + 1) % eff_L))
+                    inds = (ll, lr, ul, ur)
+
+                    tags = ("_UNI", f"_LAYER{i}")
+                    if i == 0:
+                        tags += (site_tag_id.format(j),
+                                 site_tag_id.format(j + 1))
+
+                    yield IsoTensor(next(unis), inds=inds,
+                                    tags=tags, left_inds=(ll, lr))
+
+                    # generate the isometry (offset by one effective site):
+                    #      | ui
+                    #     ISO
+                    #  ll | | lr
+                    #   j+1 j+2
+                    ll, lr = map(m_ind_id.format, (j + 1, (j + 2) % eff_L))
+                    ui = u_ind_id.format(j // 2)
+                    inds = (ll, lr, ui)
+                    tags = ("_ISO", f"_LAYER{i}")
+
+                    if i < nlayers - 1 or dangle:
+                        yield IsoTensor(next(isos), inds=inds,
+                                        tags=tags, left_inds=(ll, lr))
+                    else:
+                        # don't leave dangling index at top
+                        iso_f = next(isos)
+                        yield IsoTensor(
+                            np.eye(iso_f.shape[0], dtype=iso_f.dtype) / 2**0.5,
+                            inds=inds[:-1], tags=tags, left_inds=(ll, lr)
+                        )
+
+        super().__init__(gen_mera_tensors(), virtual=True)
+
+        # tag the MERA with the 'causal-cone' of each site
+        for i in range(nlayers):
+            for j in range(L):
+                # get isometries in the same layer
+                for t in self.select_neighbors(j):
+                    if f'_LAYER{i}' in t.tags:
+                        t.add_tag(f'I{j}')
+
+                # get unitaries in layer above
+                for t in self.select_neighbors(j):
+                    if f'_LAYER{i + 1}' in t.tags:
+                        t.add_tag(f'I{j}')
+
+    @classmethod
+    def rand(cls, L, max_bond=None, phys_dim=2, dtype=float, **mera_opts):
+
+        d = phys_dim
+        if max_bond is None:
+            max_bond = d
+
+        def gen_unis():
+            D = d
+            m = L // 2
+
+            while True:
+                for _ in range(m):
+                    uni = qu.rand_iso(D**2, D**2, dtype=dtype)
+                    uni.shape = (D, D, D, D)
+                    yield uni
+
+                D = min(D**2, max_bond)
+                m //= 2
+
+        def gen_isos():
+            Dl = d
+            Du = min(Dl**2, max_bond)
+            m = L // 2
+
+            while True:
+                for _ in range(m):
+                    iso = qu.rand_iso(Dl**2, Du, dtype=dtype)
+                    iso.shape = (Dl, Dl, Du)
+                    yield iso
+
+                Dl = Du
+                Du = min(Dl**2, max_bond)
+                m //= 2
+
+        return cls(L, gen_unis(), gen_isos(), phys_dim=d, **mera_opts)
+
+    @classmethod
+    def rand_invar(cls, L, phys_dim=2, dtype=float, **mera_opts):
+        """Generate a random translational and scale invariant MERA.
+        """
+        d = phys_dim
+
+        iso = qu.rand_iso(d**2, d, dtype=dtype)
+        iso.shape = (d, d, d)
+
+        uni = qu.rand_iso(d**2, d**2, dtype=dtype)
+        uni.shape = (d, d, d, d)
+
+        return cls(L, uni, iso, phys_dim=d, **mera_opts)