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

trajectree/quimb/quimb/tensor/geometry.py

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

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