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

trajectree/quimb/quimb/linalg/autoblock.py

@@ -1,258 +0,0 @@
-"""Numba accelerated functions for finding charge sectors and subselecting
-submatrices.
-"""
-
-import numpy as np
-import numba
-
-from ..core import njit, pnjit, qarray
-
-
-@njit
-def get_nz(A):  # pragma: no cover
-    return np.nonzero(A)
-
-
-@njit
-def compute_blocks(ix, jx, d):  # pragma: no cover
-    """Find the charge sectors (blocks in matrix terms) given element
-    coordinates ``ix`` and ``jx`` and total size ``d``.
-
-    Parameters
-    ----------
-    ix : array of int
-        The row coordinates of non-zero elements.
-    jx : array of int
-        The column coordinates of non-zero elements.
-    d : int
-        The total size of the operator.
-
-    Returns
-    -------
-    sectors : list[list[int]]
-        The list of charge sectors. Each element is itself a sorted list of the
-        basis numbers that make up that sector. The permutation that would
-        block diagonalize the operator is then ``np.concatenate(sectors)``.
-
-    Examples
-    --------
-
-    >>> H = ham_hubbard_hardcore(4, sparse=True)
-    >>> ix, jx = H.nonzero()
-    >>> d = H.shape[0]
-    >>> sectors = compute_blocks(ix, jx, d)
-    >>> sectors
-    [[0], [1, 2, 4, 8], [3, 5, 6, 9, 10, 12], [7, 11, 13, 14], [15]]
-    """
-    groups = []
-
-    # go through actual nz -> these define edges of a graph and we are
-    # looking for all connected components (disconnected subgraphs)
-    for i, j in zip(ix, jx):
-        merge = []
-        for g, group in enumerate(groups):
-            if i in group:
-                group.add(j)
-                merge.append(g)
-            elif j in group:
-                group.add(i)
-                merge.append(g)
-
-        if len(merge) == 0:
-            # new group
-            groups.append({i, j})
-
-        elif len(merge) > 1:
-            # merge groups
-            group0 = groups[merge[0]]
-            for g in merge[-1:0:-1]:
-                # XXX: just popping here causes numba big problems?
-                # so we clear and filter empty groups later
-                other_group = groups[g]
-                group0.update(other_group)
-                other_group.clear()
-
-    # make sure kernel added as subspace
-    for i in range(d):
-        for group in groups:
-            if i in group:
-                break
-        else:  # no break
-            groups.append({i})
-
-    # sort indices in each group and groups by first element
-    return sorted([sorted(g) for g in groups if g])
-
-
-@pnjit
-def subselect(A, p):  # pragma: no cover
-    """Select only the intersection of rows and columns of ``A`` matching the
-    basis indices ``p``. Faster than double numpy slicing.
-
-    Parameters
-    ----------
-    A : 2D-array
-        Dense matrix to select from.
-    p : sequence of int
-        The basis indices to select.
-
-    Returns
-    -------
-    B : 2D-array
-        The matrix, of size ``(len(p), len(p))``.
-
-    Examples
-    --------
-    >>> A = np.arange(25).reshape(5, 5)
-    >>> A
-    array([[ 0,  1,  2,  3,  4],
-           [ 5,  6,  7,  8,  9],
-           [10, 11, 12, 13, 14],
-           [15, 16, 17, 18, 19],
-           [20, 21, 22, 23, 24]])
-
-    >>> subselect(A, [1, 3])
-    array([[ 6,  8],
-           [16, 18]])
-    """
-    dp = len(p)
-    out = np.empty((dp, dp), dtype=A.dtype)
-
-    for i in numba.prange(dp):
-        for j in numba.prange(dp):
-            out[i, j] = A[p[i], p[j]]
-
-    return out
-
-
-@pnjit
-def subselect_set(A, B, p):  # pragma: no cover
-    """Set only the intersection of rows and colums of ``A`` matching the
-    basis indices ``p`` to ``B``.
-
-    Parameters
-    ----------
-    A : array with shape (d, d)
-        The matrix to set elements in.
-    B : array with shape (dp, dp)
-        The matrix to set elements from.
-    p : sequence of size dp
-        The basis indices.
-
-    Examples
-    --------
-    >>> A = np.zeros((5, 5))
-    >>> B = np.random.randn(3, 3)
-    >>> p = [0, 2, 4]
-    >>> subselect_set(A, B, p)
-    array([[-0.31888218,  0.        ,  0.39293245,  0.        ,  0.21822712],
-           [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
-           [ 0.66674486,  0.        ,  1.03388035,  0.        ,  1.7319345 ],
-           [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
-           [-0.94542733,  0.        , -0.37211882,  0.        ,  0.51951555]])
-    """
-    dp = len(p)
-
-    for i in numba.prange(dp):
-        for j in numba.prange(dp):
-            A[p[i], p[j]] = B[i, j]
-
-
-# XXX: want to cache this eventaully -> need parallel+cache numba support?
-@njit
-def _eigh_autoblocked(A, sort=True):  # pragma: no cover
-    d = A.shape[0]
-
-    # allocate output arrays
-    el = np.empty(d)
-    ev = np.zeros_like(A)
-
-    # find non-zero elements and group into charge sectors
-    ix, jx = get_nz(A)
-    gs = compute_blocks(ix, jx, d)
-    gs = [np.array(g) for g in gs]
-
-    # diagonalize each charge sector seperately
-    for g in gs:
-        ng = len(g)
-
-        # check if trivial
-        if ng == 1:
-            el[g[0]] = A[g[0], g[0]].real
-            ev[g[0], g[0]] = 1.0
-            continue
-
-        # else diagonalize just the block
-        sub_el, sub_ev = np.linalg.eigh(subselect(A, g))
-
-        # set the correct eigenpairs in the output
-        el[g] = sub_el
-        subselect_set(ev, sub_ev, g)
-
-    # sort into ascending eigenvalue order
-    if sort:
-        so = np.argsort(el)
-        el[:] = el[so]
-        ev[:, :] = ev[:, so]
-
-    return el, ev
-
-
-# XXX: want to cache this eventaully -> need parallel+cache numba support?
-@njit
-def _eigvalsh_autoblocked(A, sort=True):  # pragma: no cover
-    # as above but ignore eigenvector for extra speed
-    d = A.shape[0]
-
-    el = np.empty(d)
-
-    ix, jx = get_nz(A)
-    gs = compute_blocks(ix, jx, d)
-    gs = [np.array(g) for g in gs]
-
-    for _, g in enumerate(gs):
-        if len(g) == 1:
-            el[g[0]] = A[g[0], g[0]]
-            continue
-
-        el[g] = np.linalg.eigvalsh(subselect(A, g))
-
-    if sort:
-        return np.sort(el)
-
-    return el
-
-
-def eigensystem_autoblocked(A, sort=True, return_vecs=True, isherm=True):
-    """Perform Hermitian eigen-decomposition, automatically identifying and
-    exploiting symmetries appearing in the current basis as block diagonals
-    formed via permutation of rows and columns. The whole process is
-    accelerated using ``numba``.
-
-    Parameters
-    ----------
-    A : array_like
-        The operator to eigen-decompose.
-    sort : bool, optional
-        Whether to sort into ascending order, default True.
-    isherm : bool, optional
-        Whether ``A`` is hermitian, default True.
-    return_vecs : bool, optional
-        Whether to return the eigenvectors, default True.
-
-    Returns
-    -------
-    evals : 1D-array
-        The eigenvalues.
-    evecs : qarray
-        If ``return_vecs=True``, the eigenvectors.
-    """
-    if not isherm:
-        err_msg = "Non-hermitian autoblocking not implemented yet."
-        raise NotImplementedError(err_msg)
-
-    if not return_vecs:
-        return _eigvalsh_autoblocked(A, sort=sort)
-
-    el, ev = _eigh_autoblocked(A, sort=sort)
-    return el, qarray(ev)