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

trajectree/quimb/quimb/tensor/geometry.py

@@ -0,0 +1,583 @@
+"""Some functions for generating the edges of a lattice."""
+
+import itertools
+import random
+
+
+def sort_unique(edges):
+    """Make sure there are no duplicate edges and that for each
+    ``coo_a < coo_b``.
+    """
+    return tuple(
+        sorted(tuple(sorted(edge)) for edge in set(map(frozenset, edges)))
+    )
+
+
+# ----------------------------------- 1D ------------------------------------ #
+
+
+def edges_1d_chain(L, cyclic=False):
+    """Return the graph edges of a finite 1D chain lattice.
+
+    Parameters
+    ----------
+    L : int
+        The number of cells.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+
+    Returns
+    -------
+    edges : list[(int, int)]
+    """
+    edges = []
+    for i in range(L):
+        if i < L - 1 or cyclic:
+            edges.append((i, (i + 1) % L))
+
+    return sort_unique(edges)
+
+
+# ----------------------------------- 2D ------------------------------------ #
+
+
+def check_2d(coo, Lx, Ly, cyclic):
+    """Check ``coo`` in inbounds for a maybe cyclic 2D lattice."""
+    x, y = coo
+    if (not cyclic) and not ((0 <= x < Lx) and (0 <= y < Ly)):
+        return
+    return (x % Lx, y % Ly)
+
+
+def edges_2d_square(Lx, Ly, cyclic=False, cells=None):
+    """Return the graph edges of a finite 2D square lattice. The nodes
+    (sites) are labelled like ``(i, j)``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly))``.
+
+    Returns
+    -------
+    edges : list[((int, int), (int, int))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly))
+
+    edges = []
+    for i, j in cells:
+        for coob in [(i, j + 1), (i + 1, j)]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j), coob))
+
+    return sort_unique(edges)
+
+
+def edges_2d_hexagonal(Lx, Ly, cyclic=False, cells=None):
+    """Return the graph edges of a finite 2D hexagonal lattice. There are two
+    sites per cell, and note the cells do not form a square tiling. The nodes
+    (sites) are labelled like ``(i, j, s)`` for ``s`` in ``'AB'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly))``.
+
+    Returns
+    -------
+    edges : list[((int, int, str), (int, int, str))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly))
+
+    edges = []
+    for i, j in cells:
+        for *coob, lbl in [
+            (i, j, "B"),
+            (i, j - 1, "B"),
+            (i - 1, j, "B"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, "A"),
+            (i, j + 1, "A"),
+            (i + 1, j, "A"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "B"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+def edges_2d_triangular(Lx, Ly, cyclic=False, cells=None):
+    """Return the graph edges of a finite 2D triangular lattice. There is a
+    single site per cell, and note the cells do not form a square tiling.
+    The nodes (sites) are labelled like ``(i, j)``.
+
+    Parameters
+    ----------
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly))``.
+
+    Returns
+    -------
+    edges : list[((int, int), (int, int))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly))
+
+    edges = []
+    for i, j in cells:
+        for coob in [(i, j + 1), (i + 1, j), (i + 1, j - 1)]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j), coob))
+
+    return sort_unique(edges)
+
+
+def edges_2d_triangular_rectangular(Lx, Ly, cyclic=False, cells=None):
+    """Return the graph edges of a finite 2D triangular lattice tiled in a
+    rectangular geometry. There are two sites per rectangular cell. The nodes
+    (sites) are labelled like ``(i, j, s)`` for ``s`` in ``'AB'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly))``.
+
+    Returns
+    -------
+    edges : list[((int, int, s), (int, int, s))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly))
+
+    edges = []
+    for i, j in cells:
+        for *coob, lbl in [
+            (i, j, "B"),
+            (i, j - 1, "B"),
+            (i, j + 1, "A"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i + 1, j, "A"),
+            (i, j + 1, "B"),
+            (i + 1, j + 1, "A"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "B"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+def edges_2d_kagome(Lx, Ly, cyclic=False, cells=None):
+    """Return the graph edges of a finite 2D kagome lattice. There are
+    three sites per cell, and note the cells do not form a square tiling. The
+    nodes (sites) are labelled like ``(i, j, s)`` for ``s`` in ``'ABC'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly))``.
+
+    Returns
+    -------
+    edges : list[((int, int, str), (int, int, str))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly))
+
+    edges = []
+    for i, j in cells:
+        for *coob, lbl in [
+            (i, j, "B"),
+            (i, j - 1, "B"),
+            (i, j, "C"),
+            (i - 1, j, "C"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, "C"),
+            (i - 1, j + 1, "C"),
+            (i, j, "A"),
+            (i, j + 1, "A"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "B"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, "A"),
+            (i + 1, j, "A"),
+            (i, j, "B"),
+            (i + 1, j - 1, "B"),
+        ]:
+            coob = check_2d(coob, Lx, Ly, cyclic)
+            if coob:
+                edges.append(((i, j, "C"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+# ----------------------------------- 3D ------------------------------------ #
+
+
+def check_3d(coo, Lx, Ly, Lz, cyclic):
+    """Check ``coo`` in inbounds for a maybe cyclic 3D lattice."""
+    x, y, z = coo
+    OBC = not cyclic
+    inbounds = (0 <= x < Lx) and (0 <= y < Ly) and (0 <= z < Lz)
+    if OBC and not inbounds:
+        return
+    return (x % Lx, y % Ly, z % Lz)
+
+
+def edges_3d_cubic(Lx, Ly, Lz, cyclic=False, cells=None):
+    """Return the graph edges of a finite 3D cubic lattice. The nodes
+    (sites) are labelled like ``(i, j, k)``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    Lz : int
+        The number of cells along the z-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly), range(Lz))``.
+
+    Returns
+    -------
+    edges : list[((int, int, int), (int, int, int))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly), range(Lz))
+
+    edges = []
+    for i, j, k in cells:
+        for coob in [(i, j, k + 1), (i, j + 1, k), (i + 1, j, k)]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k), coob))
+
+    return sort_unique(edges)
+
+
+def edges_3d_pyrochlore(Lx, Ly, Lz, cyclic=False, cells=None):
+    """Return the graph edges of a finite 3D pyorchlore lattice. There are
+    four sites per cell, and note the cells do not form a cubic tiling. The
+    nodes (sites) are labelled like ``(i, j, k, s)`` for ``s`` in ``'ABCD'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    Lz : int
+        The number of cells along the z-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly), range(Lz))``.
+
+    Returns
+    -------
+    edges : list[((int, int, int, str), (int, int, int, str))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly), range(Lz))
+
+    edges = []
+    for i, j, k in cells:
+        for *coob, lbl in [
+            (i, j, k, "B"),
+            (i, j - 1, k, "B"),
+            (i, j, k, "C"),
+            (i - 1, j, k, "C"),
+            (i, j, k, "D"),
+            (i, j, k - 1, "D"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "C"),
+            (i - 1, j + 1, k, "C"),
+            (i, j, k, "D"),
+            (i, j + 1, k - 1, "D"),
+            (i, j, k, "A"),
+            (i, j + 1, k, "A"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "B"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "D"),
+            (i + 1, j, k - 1, "D"),
+            (i, j, k, "A"),
+            (i + 1, j, k, "A"),
+            (i, j, k, "B"),
+            (i + 1, j - 1, k, "B"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "C"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "A"),
+            (i, j, k + 1, "A"),
+            (i, j, k, "B"),
+            (i, j - 1, k + 1, "B"),
+            (i, j, k, "C"),
+            (i - 1, j, k + 1, "C"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "D"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+def edges_3d_diamond(Lx, Ly, Lz, cyclic=False, cells=None):
+    """Return the graph edges of a finite 3D diamond lattice. There are
+    two sites per cell, and note the cells do not form a cubic tiling.  The
+    nodes (sites) are labelled like ``(i, j, k, s)`` for ``s`` in ``'AB'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    Lz : int
+        The number of cells along the z-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly), range(Lz))``.
+
+    Returns
+    -------
+    edges : list[((int, int, int, str), (int, int, int, str))]
+    """
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly), range(Lz))
+
+    edges = []
+    for i, j, k in cells:
+        for *coob, lbl in [
+            (i, j, k, "B"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k + 1, "A"),
+            (i, j + 1, k, "A"),
+            (i + 1, j, k, "A"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "B"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+def edges_3d_diamond_cubic(Lx, Ly, Lz, cyclic=False, cells=None):
+    """Return the graph edges of a finite 3D diamond lattice tiled in a cubic
+    geometry. There are eight sites per cubic cell. The nodes (sites) are
+    labelled like ``(i, j, k, s)`` for ``s`` in ``'ABCDEFGH'``.
+
+    Parameters
+    ----------
+    Lx : int
+        The number of cells along the x-direction.
+    Ly : int
+        The number of cells along the y-direction.
+    Lz : int
+        The number of cells along the z-direction.
+    cyclic : bool, optional
+        Whether to use periodic boundary conditions.
+    cells : list, optional
+        A list of cells to use. If not given the cells used are
+        ``itertools.product(range(Lx), range(Ly), range(Lz))``.
+
+    Returns
+    -------
+    edges : list[((int, int, int, str), (int, int, int, str))]
+    """
+
+    if cells is None:
+        cells = itertools.product(range(Lx), range(Ly), range(Lz))
+
+    edges = []
+    for i, j, k in cells:
+        for *coob, lbl in [
+            (i, j, k, "E"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "A"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "E"),
+            (i, j, k, "F"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "B"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "E"),
+            (i, j, k, "G"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "C"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j, k, "E"),
+            (i, j, k, "H"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "D"), (*coob, lbl)))
+
+        for *coob, lbl in []:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "E"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j + 1, k, "C"),
+            (i + 1, j, k, "D"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "F"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i + 1, j, k + 1, "A"),
+            (i, j, k + 1, "B"),
+            (i + 1, j, k, "D"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "G"), (*coob, lbl)))
+
+        for *coob, lbl in [
+            (i, j + 1, k + 1, "A"),
+            (i, j, k + 1, "B"),
+            (i, j + 1, k, "C"),
+        ]:
+            coob = check_3d(coob, Lx, Ly, Lz, cyclic)
+            if coob:
+                edges.append(((i, j, k, "H"), (*coob, lbl)))
+
+    return sort_unique(edges)
+
+
+def edges_tree_rand(n, max_degree=None, seed=None):
+    """Return a random tree with ``n`` nodes. This a convenience function for
+    testing purposes and the trees generated are not guaranteed to be uniformly
+    random (for that see ``networkx.random_labeled_tree``).
+
+    Parameters
+    ----------
+    n : int
+        The number of nodes.
+    max_degree : int, optional
+        The maximum degree of the nodes. For example ``max_degree=3`` means
+        generate a binary tree.
+    seed : int, optional
+        The random seed.
+
+    Returns
+    -------
+    edges : list[(int, int)]
+    """
+    rng = random.Random(seed)
+    edges = []
+
+    if max_degree is None:
+        nodes = [0]
+        for i in range(1, n):
+            ib = rng.choice(nodes)
+            nodes.append(i)
+            edges.append((ib, i))
+    else:
+        degrees = {0: 0}
+        for i in range(1, n):
+            ib = rng.choice(list(degrees))
+            edges.append((ib, i))
+
+            degrees[i] = 1
+            if degrees[ib] + 1 == max_degree:
+                # node is finished
+                degrees.pop(ib)
+            else:
+                degrees[ib] += 1
+
+    return edges

trajectree/quimb/quimb/tensor/interface.py

@@ -0,0 +1,114 @@
+"""Tools for interfacing the tensor and tensor network objects with other
+libraries.
+"""
+
+import functools
+
+from ..utils import tree_map
+from .tensor_core import Tensor, TensorNetwork
+
+
+class Placeholder:
+    __slots__ = ("shape",)
+
+    def __init__(self, x):
+        self.shape = getattr(x, "shape", None)
+
+    def __repr__(self):
+        return f"Placeholder(shape={self.shape})"
+
+
+def pack(obj):
+    """Take a tensor or tensor network like object and return a skeleton needed
+    to reconstruct it, and a pytree of raw parameters.
+
+    Parameters
+    ----------
+    obj : Tensor, TensorNetwork, or similar
+        Something that has ``copy``, ``set_params``, and ``get_params``
+        methods.
+
+    Returns
+    -------
+    params : pytree
+        A pytree of raw parameter arrays.
+    skeleton : Tensor, TensorNetwork, or similar
+        A copy of ``obj`` with all references to the original data removed.
+    """
+    try:
+        skeleton = obj.copy()
+        params = skeleton.get_params()
+        placeholders = tree_map(Placeholder, params)
+        skeleton.set_params(placeholders)
+    except AttributeError:
+        # assume it's a raw array
+        params = obj
+        skeleton = Placeholder(obj)
+    return params, skeleton
+
+
+def unpack(params, skeleton):
+    """Take a skeleton of a tensor or tensor network like object and a pytree
+    of raw parameters and return a new reconstructed object with those
+    parameters inserted.
+
+    Parameters
+    ----------
+    params : pytree
+        A pytree of raw parameter arrays, with the same structure as the
+        output of ``skeleton.get_params()``.
+    skeleton : Tensor, TensorNetwork, or similar
+        Something that has ``copy``, ``set_params``, and ``get_params``
+        methods.
+
+    Returns
+    -------
+    obj : Tensor, TensorNetwork, or similar
+        A copy of ``skeleton`` with parameters inserted.
+    """
+    try:
+        obj = skeleton.copy()
+        obj.set_params(params)
+    except AttributeError:
+        # assume it's a raw array
+        obj = params
+    return obj
+
+
+# -------------------------------- jax -------------------------------------- #
+
+
+_JAX_REGISTERED_TN_CLASSES = set()
+
+
+def jax_pack(obj):
+    # jax requires the top level children to be a tuple
+    params, aux = pack(obj)
+    children = (params,)
+    return children, aux
+
+
+def jax_unpack(aux, children):
+    # jax also flips the return order from above
+    (params,) = children
+    return unpack(params, aux)
+
+
+def jax_register_pytree():
+    import jax
+
+    queue = [Tensor, TensorNetwork]
+    while queue:
+        cls = queue.pop()
+        if cls not in _JAX_REGISTERED_TN_CLASSES:
+            jax.tree_util.register_pytree_node(cls, jax_pack, jax_unpack)
+            _JAX_REGISTERED_TN_CLASSES.add(cls)
+        queue.extend(cls.__subclasses__())
+
+
+@functools.lru_cache(1)
+def get_jax():
+    import jax
+
+    jax_register_pytree()
+    return jax