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

trajectree/quimb/quimb/tensor/array_ops.py

@@ -0,0 +1,610 @@
+"""Backend agnostic array operations."""
+
+import functools
+import itertools
+
+import numpy
+from autoray import (
+    compose,
+    do,
+    get_dtype_name,
+    get_lib_fn,
+    infer_backend,
+    reshape,
+)
+
+from ..core import njit, qarray
+from ..linalg.base_linalg import norm_fro_dense
+from ..utils import compose as fn_compose
+
+
+def asarray(array):
+    """Maybe convert data for a tensor to use. If ``array`` already has a
+    ``.shape`` attribute, i.e. looks like an array, it is left as-is. Else the
+    elements are inspected to see which libraries' array constructor should be
+    used, defaulting to ``numpy`` if everything is builtin or numpy numbers.
+    """
+    if isinstance(array, (numpy.matrix, qarray)):
+        # if numpy make sure array not subclass
+        return numpy.asarray(array)
+
+    if hasattr(array, "shape"):
+        # otherwise don't touch things which are already array like
+        return array
+
+    # else we some kind of possibly nested python iterable -> inspect items
+    backends = set()
+
+    def _nd_py_iter(x):
+        if isinstance(x, str):
+            # handle recursion error
+            return x
+
+        backend = infer_backend(x)
+        if backend != "builtins":
+            # don't iterate any non-builtin containers
+            backends.add(backend)
+            return x
+
+        # is some kind of python container or element -> iterate or return
+        try:
+            return list(_nd_py_iter(sub) for sub in x)
+        except TypeError:
+            return x
+
+    nested_tup = _nd_py_iter(array)
+
+    # numpy and builtin elements treat as basic
+    backends -= {"builtins", "numpy"}
+    if not backends:
+        backend = "numpy"
+    else:
+        (backend,) = backends
+
+    return do("array", nested_tup, like=backend)
+
+
+_blocksparselookup = {}
+
+
+def isblocksparse(x):
+    """Check if `x` is a block-sparse array. Cached on class for speed."""
+    try:
+        return _blocksparselookup[x.__class__]
+    except KeyError:
+        # XXX: make this a more established interface
+        isbs = hasattr(x, "align_axes")
+        _blocksparselookup[x.__class__] = isbs
+        return isbs
+
+
+_fermioniclookup = {}
+
+
+def isfermionic(x):
+    """Check if `x` is a fermionic array. Cached on class for speed."""
+    try:
+        return _fermioniclookup[x.__class__]
+    except KeyError:
+        # XXX: make this a more established interface
+        isf = hasattr(x, "phase_flip")
+        _fermioniclookup[x.__class__] = isf
+        return isf
+
+
+@functools.lru_cache(2**14)
+def calc_fuse_perm_and_shape(shape, axes_groups):
+    ndim = len(shape)
+
+    # which group does each axis appear in, if any
+    num_groups = len(axes_groups)
+    ax2group = {ax: g for g, axes in enumerate(axes_groups) for ax in axes}
+
+    # the permutation will be the same for every block: precalculate
+    # n.b. all new groups will be inserted at the *first fused axis*
+    position = min(g for gax in axes_groups for g in gax)
+    axes_before = tuple(
+        ax for ax in range(position) if ax2group.setdefault(ax, None) is None
+    )
+    axes_after = tuple(
+        ax
+        for ax in range(position, ndim)
+        if ax2group.setdefault(ax, None) is None
+    )
+    perm = (*axes_before, *(ax for g in axes_groups for ax in g), *axes_after)
+
+    # track where each axis will be in the new array
+    new_axes = {ax: ax for ax in axes_before}
+    for i, g in enumerate(axes_groups):
+        for ax in g:
+            new_axes[ax] = position + i
+    for i, ax in enumerate(axes_after):
+        new_axes[ax] = position + num_groups + i
+    new_ndim = len(axes_before) + num_groups + len(axes_after)
+
+    new_shape = [1] * new_ndim
+    for i, d in enumerate(shape):
+        g = ax2group[i]
+        new_ax = new_axes[i]
+        if g is None:
+            # not fusing, new value is just copied
+            new_shape[new_ax] = d
+        else:
+            # fusing: need to accumulate
+            new_shape[new_ax] *= d
+
+    if all(i == ax for i, ax in enumerate(perm)):
+        # no need to transpose
+        perm = None
+
+    new_shape = tuple(new_shape)
+    if shape == new_shape:
+        # no need to reshape
+        new_shape = None
+
+    return perm, new_shape
+
+
+@compose
+def fuse(x, *axes_groups, backend=None):
+    """Fuse the give group or groups of axes. The new fused axes will be
+    inserted at the minimum index of any fused axis (even if it is not in
+    the first group). For example, ``fuse(x, [5, 3], [7, 2, 6])`` will
+    produce an array with axes like::
+
+        groups inserted at axis 2, removed beyond that.
+                ......<--
+        (0, 1, g0, g1, 4, 8, ...)
+                |   |
+                |   g1=(7, 2, 6)
+                g0=(5, 3)
+
+    Parameters
+    ----------
+    axes_groups : sequence of sequences of int
+        The axes to fuse. Each group of axes will be fused into a single
+        axis.
+    """
+    if backend is None:
+        backend = infer_backend(x)
+    _transpose = get_lib_fn(backend, "transpose")
+    _reshape = get_lib_fn(backend, "reshape")
+
+    axes_groups = tuple(map(tuple, axes_groups))
+    if not any(axes_groups):
+        return x
+
+    shape = tuple(map(int, x.shape))
+    perm, new_shape = calc_fuse_perm_and_shape(shape, axes_groups)
+
+    if perm is not None:
+        x = _transpose(x, perm)
+    if new_shape is not None:
+        x = _reshape(x, new_shape)
+
+    return x
+
+
+def ndim(array):
+    """The number of dimensions of an array."""
+    try:
+        return array.ndim
+    except AttributeError:
+        return len(array.shape)
+
+
+@compose
+def multiply_diagonal(x, v, axis, backend=None):
+    """Multiply v into x as if contracting in a diagonal matrix."""
+    newshape = tuple((-1 if i == axis else 1) for i in range(ndim(x)))
+    v_broadcast = do("reshape", v, newshape, like=backend)
+    return x * v_broadcast
+
+
+@compose
+def align_axes(*arrays, axes, backend=None):
+    """Prepare a set of arrays that should be contractible along ``axes``.
+
+    For example, block symmetric arrays need to have aligned sectors prior to
+    fusing.
+    """
+    # default implementation is nothing
+    return arrays
+
+
+# ------------- miscelleneous other backend agnostic functions -------------- #
+
+
+def iscomplex(x):
+    """Does ``x`` have a complex dtype?"""
+    if infer_backend(x) == "builtins":
+        return isinstance(x, complex)
+    return "complex" in get_dtype_name(x)
+
+
+@compose
+def norm_fro(x):
+    """The frobenius norm of an array."""
+    try:
+        return do("linalg.norm", reshape(x, (-1,)))
+    except AttributeError:
+        return do("sum", do("abs", x) ** 2) ** 0.5
+
+
+norm_fro.register("numpy", norm_fro_dense)
+
+
+def sensibly_scale(x):
+    """Take an array and scale it *very* roughly such that random tensor
+    networks consisting of such arrays do not have gigantic norms.
+    """
+    return x / norm_fro(x) ** (1.5 / ndim(x))
+
+
+@njit
+def _numba_find_diag_axes(x, atol=1e-12):  # pragma: no cover
+    """Numba-compiled array diagonal axis finder.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        The array to search for diagonal axes.
+    atol : float
+        The tolerance with which to compare to zero.
+
+    Returns
+    -------
+    diag_axes : set[tuple[int]]
+        The set of pairs of axes which are diagonal.
+    """
+
+    # create the set of pairs of matching size axes
+    diag_axes = set()
+    for d1 in range(x.ndim - 1):
+        for d2 in range(d1 + 1, x.ndim):
+            if x.shape[d1] == x.shape[d2]:
+                diag_axes.add((d1, d2))
+
+    # enumerate through every array entry, eagerly invalidating axis pairs
+    for index, val in numpy.ndenumerate(x):
+        for d1, d2 in list(diag_axes):
+            if (index[d1] != index[d2]) and (abs(val) > atol):
+                diag_axes.remove((d1, d2))
+
+        # all pairs invalid, nothing left to do
+        if len(diag_axes) == 0:
+            break
+
+    return diag_axes
+
+
+def find_diag_axes(x, atol=1e-12):
+    """Try and find a pair of axes of ``x`` in which it is diagonal.
+
+    Parameters
+    ----------
+    x : array-like
+        The array to search.
+    atol : float, optional
+        Tolerance with which to compare to zero.
+
+    Returns
+    -------
+    tuple[int] or None
+        The two axes if found else None.
+
+    Examples
+    --------
+
+        >>> x = np.array([[[1, 0], [0, 2]],
+        ...               [[3, 0], [0, 4]]])
+        >>> find_diag_axes(x)
+        (1, 2)
+
+    Which means we can reduce ``x`` without loss of information to:
+
+        >>> np.einsum('abb->ab', x)
+        array([[1, 2],
+               [3, 4]])
+
+    """
+    shape = x.shape
+    if len(shape) < 2:
+        return None
+
+    backend = infer_backend(x)
+    zero = do("zeros", (), like=x)
+
+    # use numba-accelerated version for numpy arrays
+    if backend == "numpy":
+        diag_axes = _numba_find_diag_axes(x, atol=atol)
+        if diag_axes:
+            # make it determinstic
+            return min(diag_axes)
+        return None
+    indxrs = do("indices", shape, like=backend)
+
+    for i, j in itertools.combinations(range(len(shape)), 2):
+        if shape[i] != shape[j]:
+            continue
+        if do(
+            "allclose",
+            x[indxrs[i] != indxrs[j]],
+            zero,
+            atol=atol,
+            like=backend,
+        ):
+            return (i, j)
+    return None
+
+
+@njit
+def _numba_find_antidiag_axes(x, atol=1e-12):  # pragma: no cover
+    """Numba-compiled array antidiagonal axis finder.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        The array to search for anti-diagonal axes.
+    atol : float
+        The tolerance with which to compare to zero.
+
+    Returns
+    -------
+    antidiag_axes : set[tuple[int]]
+        The set of pairs of axes which are anti-diagonal.
+    """
+
+    # create the set of pairs of matching size axes
+    antidiag_axes = set()
+    for i in range(x.ndim - 1):
+        for j in range(i + 1, x.ndim):
+            if x.shape[i] == x.shape[j]:
+                antidiag_axes.add((i, j))
+
+    # enumerate through every array entry, eagerly invalidating axis pairs
+    for index, val in numpy.ndenumerate(x):
+        for i, j in list(antidiag_axes):
+            d = x.shape[i]
+            if (index[i] != d - 1 - index[j]) and (abs(val) > atol):
+                antidiag_axes.remove((i, j))
+
+        # all pairs invalid, nothing left to do
+        if len(antidiag_axes) == 0:
+            break
+
+    return antidiag_axes
+
+
+def find_antidiag_axes(x, atol=1e-12):
+    """Try and find a pair of axes of ``x`` in which it is anti-diagonal.
+
+    Parameters
+    ----------
+    x : array-like
+        The array to search.
+    atol : float, optional
+        Tolerance with which to compare to zero.
+
+    Returns
+    -------
+    tuple[int] or None
+        The two axes if found else None.
+
+    Examples
+    --------
+
+        >>> x = np.array([[[0, 1], [0, 2]],
+        ...               [[3, 0], [4, 0]]])
+        >>> find_antidiag_axes(x)
+        (0, 2)
+
+    Which means we can reduce ``x`` without loss of information to:
+
+        >>> np.einsum('aba->ab', x[::-1, :, :])
+        array([[3, 4],
+               [1, 2]])
+
+    as long as we flip the order of dimensions on other tensors corresponding
+    to the the same index.
+    """
+    shape = x.shape
+    if len(shape) < 2:
+        return None
+
+    backend = infer_backend(x)
+
+    # use numba-accelerated version for numpy arrays
+    if backend == "numpy":
+        antidiag_axes = _numba_find_antidiag_axes(x, atol=atol)
+        if antidiag_axes:
+            # make it determinstic
+            return min(antidiag_axes)
+        return None
+
+    indxrs = do("indices", shape, like=backend)
+    zero = do("zeros", (), like=x)
+
+    for i, j in itertools.combinations(range(len(shape)), 2):
+        di, dj = shape[i], shape[j]
+        if di != dj:
+            continue
+        if do(
+            "allclose",
+            x[indxrs[i] != dj - 1 - indxrs[j]],
+            zero,
+            atol=atol,
+            like=backend,
+        ):
+            return (i, j)
+    return None
+
+
+@njit
+def _numba_find_columns(x, atol=1e-12):  # pragma: no cover
+    """Numba-compiled single non-zero column axis finder.
+
+    Parameters
+    ----------
+    x :  array
+        The array to search.
+    atol : float, optional
+        Absolute tolerance to compare to zero with.
+
+    Returns
+    -------
+    set[tuple[int]]
+        Set of pairs (axis, index) defining lone non-zero columns.
+    """
+
+    # possible pairings of axis + index
+    column_pairs = set()
+    for ax, d in enumerate(x.shape):
+        for i in range(d):
+            column_pairs.add((ax, i))
+
+    # enumerate over all array entries, invalidating potential column pairs
+    for index, val in numpy.ndenumerate(x):
+        if abs(val) > atol:
+            for ax, i in enumerate(index):
+                for pax, pi in list(column_pairs):
+                    if ax == pax and pi != i:
+                        column_pairs.remove((pax, pi))
+
+        # all potential pairs invalidated
+        if not len(column_pairs):
+            break
+
+    return column_pairs
+
+
+def find_columns(x, atol=1e-12):
+    """Try and find columns of axes which are zero apart from a single index.
+
+    Parameters
+    ----------
+    x : array-like
+        The array to search.
+    atol : float, optional
+        Tolerance with which to compare to zero.
+
+    Returns
+    -------
+    tuple[int] or None
+        If found, the first integer is which axis, and the second is which
+        column of that axis, else None.
+
+    Examples
+    --------
+
+        >>> x = np.array([[[0, 1], [0, 2]],
+        ...               [[0, 3], [0, 4]]])
+        >>> find_columns(x)
+        (2, 1)
+
+    Which means we can happily slice ``x`` without loss of information to:
+
+        >>> x[:, :, 1]
+        array([[1, 2],
+               [3, 4]])
+
+    """
+    shape = x.shape
+    if len(shape) < 1:
+        return None
+
+    backend = infer_backend(x)
+
+    # use numba-accelerated version for numpy arrays
+    if backend == "numpy":
+        columns_pairs = _numba_find_columns(x, atol)
+        if columns_pairs:
+            return min(columns_pairs)
+        return None
+
+    indxrs = do("indices", shape, like=backend)
+    zero = do("zeros", (), like=x)
+
+    for i in range(len(shape)):
+        for j in range(shape[i]):
+            if do(
+                "allclose", x[indxrs[i] != j], zero, atol=atol, like=backend
+            ):
+                return (i, j)
+
+    return None
+
+
+class PArray:
+    """Simple array-like object that lazily generates the actual array by
+    calling a function with a set of parameters.
+
+    Parameters
+    ----------
+    fn : callable
+        The function that generates the tensor data from ``params``.
+    params : sequence of numbers
+        The initial parameters supplied to the generating function like
+        ``fn(params)``.
+
+    See Also
+    --------
+    PTensor
+    """
+
+    __slots__ = ("_fn", "_params", "_data", "_shape", "_shape_fn_id")
+
+    def __init__(self, fn, params, shape=None):
+        self.fn = fn
+        self.params = params
+        self._shape = shape
+        self._shape_fn_id = id(fn)
+
+    def copy(self):
+        new = PArray(self.fn, self.params, self.shape)
+        new._data = self._data  # for efficiency
+        return new
+
+    @property
+    def fn(self):
+        return self._fn
+
+    @fn.setter
+    def fn(self, x):
+        self._fn = x
+        self._data = None
+
+    @property
+    def params(self):
+        return self._params
+
+    @params.setter
+    def params(self, x):
+        self._params = asarray(x)
+        self._data = None
+
+    @property
+    def data(self):
+        if self._data is None:
+            self._data = self._fn(self._params)
+        return self._data
+
+    @property
+    def shape(self):
+        # if we haven't calculated shape or have updated function, get shape
+        _shape_fn_id = id(self.fn)
+        if (self._shape is None) or (self._shape_fn_id != _shape_fn_id):
+            self._shape = self.data.shape
+            self._shape_fn_id = _shape_fn_id
+        return self._shape
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    def add_function(self, g):
+        """Chain the new function ``g`` on top of current function ``f`` like
+        ``g(f(params))``.
+        """
+        f = self.fn
+        self.fn = fn_compose(g, f)