Source code for africanus.linalg.kronecker_tools

# -*- coding: utf-8 -*-


import numpy as np


def kron_N(x):
    """
    Computes N = N_1 x N_2 x ... x N_D i.e.
    the total number of rows in a kronecker matrix

    Parameters
    ----------
    x : :class:`numpy.ndarray`
        An array of arrays holding matrices/vectors [x1, x2, ..., xD]

    Returns
    -------
    N : int
        The total number of rows in a kronecker matrix or vector
    """
    D = x.shape[0]
    dims = np.zeros(D)
    for i in range(D):
        dims[i] = x[i].shape[0]
    return int(np.prod(dims))


[docs] def kron_matvec(A, b): """ Computes the matrix vector product of a kronecker matrix in linear time. Assumes A consists of kronecker product of square matrices. Parameters ---------- A : :class:`numpy.ndarray` An array of arrays holding matrices [K0, K1, ...] where :math:`A = K_0 \\otimes K_1 \\otimes \\cdots` b : :class:`numpy.ndarray` The right hand side vector Returns ------- x : :class:`numpy.ndarray` The result of :code:`A.dot(b)` """ D = A.shape[0] N = b.size x = b for d in range(D): Gd = A[d].shape[0] X = np.reshape(x, (Gd, N // Gd)) Z = np.einsum("ab,bc->ac", A[d], X) Z = np.einsum("ab -> ba", Z) x = Z.flatten() return x
def kron_tensorvec(A, b): """ Matrix vector product of kronecker matrix A with vector b. A can be made up of an arbitrary kronecker product. Parameters ---------- A : :class:`numpy.ndarray` An array of arrays holding matrices [K0, K1, ...] where :math:`A = K_0 \\otimes K_1 \\otimes \\cdots` b : :class:`numpy.ndarray` The right hand side vector Returns ------- x : :class:`numpy.ndarray` The result of :code:`A.dot(b)` """ D = A.shape[0] # get shape of sub-matrices G = np.zeros(D, dtype=np.int8) M = np.zeros(D, dtype=np.int8) for d in range(D): M[d], G[d] = A[d].shape x = b for d in range(D): Gd = G[d] rem = np.prod(np.delete(G, d)) X = np.reshape(x, (Gd, rem)) Z = np.einsum("ab,bc->ac", A[d], X) Z = np.einsum("ab -> ba", Z) x = Z.flatten() # replace with new dimension G[d] = M[d] return x def kron_matmat(A, B): """ Computes the product between a kronecker matrix A and some RHS matrix B Parameters ---------- A : :class:`numpy.ndarray` An array of arrays holding matrices [K0, K1, ...] where :math:`A = K_0 \\otimes K_1 \\otimes \\cdots` B : :class:`numpy.ndarray` The RHS matrix Returns ------- x : :class:`numpy.ndarray` The result of :code:`A.dot(B)` """ M = B.shape[1] # the product of Np_1 x Np_2 x ... x Np_3 N = kron_N(A) C = np.zeros([N, M]) for i in range(M): C[:, i] = kron_matvec(A, B[:, i]) return C def kron_tensormat(A, B): """ Computes the matrix product between A kronecker matrix A and some RHS matrix B. Does not assume A to consist of a kronecker product of square matrices. Parameters ---------- A : :class:`numpy.ndarray` An array of arrays holding matrices [K0, K1, ...] where :math:`A = K_0 \\otimes K_1 \\otimes \\cdots` B : :class:`numpy.ndarray` The RHS matrix Returns ------- x : :class:`numpy.ndarray` The result of :code:`A.dot(B)` """ M = B.shape[1] # the product of Np_1 x Np_2 x ... x Np_3 N = kron_N(A) C = np.zeros([N, M]) for i in range(M): C[:, i] = kron_tensorvec(A, B[:, i]) return C
[docs] def kron_cholesky(A): """ Computes the Cholesky decomposition of a kronecker matrix as a kronecker matrix of Cholesky factors. Parameters ---------- A : :class:`numpy.ndarray` An array of arrays holding matrices [K0, K1, ...] where :math:`A = K_0 \\otimes K_1 \\otimes \\cdots` Returns ------- L : :class:`numpy.ndarray` An array of arrays holding matrices [L0, L1, ...] where :math:`L = L_0 \\otimes L_1 \\otimes \\cdots` and each :code:`Li = cholesky(Ki)` """ D = A.shape[0] L = np.zeros_like(A) for i in range(D): try: L[i] = np.linalg.cholesky(A[i]) except Exception: # add jitter L[i] = np.linalg.cholesky(A[i] + 1e-13 * np.eye(A[i].shape[0])) return L