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

trajectree/quimb/quimb/linalg/autoblock.py

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