Source code for quimb.linalg.autoblock

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(['void(int32, int32, List(Set(int32)))',
       'void(int64, int64, List(Set(int64)))'])
def _add_to_groups(i, j, groups):  # pragma: no cover
    for group in groups:
        if i in group:
            group.add(j)
            return
        if j in group:
            group.add(i)
            return
    # pair is not in a sector yet - create new one
    groups.append({i, j})


[docs]@njit(['List(List(int32))(int32[:], int32[:], int_)', 'List(List(int64))(int64[:], int64[:], int_)']) 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 = [{ix[0], jx[0]}] # go through actual nz for i, j in zip(ix, jx): _add_to_groups(i, j, groups) # make sure kernel added as subspace for i in range(d): _add_to_groups(i, i, groups) # sort indices in each group and groups by first element return sorted([sorted(g) for g in groups])
[docs]@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
[docs]@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 i, g in enumerate(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 i, 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: el[:] = np.sort(el) return el
[docs]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)