def kronsum(A, B, format=None):
"""Kronecker sum of sparse matrices A and B.
Kronecker sum is the sum of two Kronecker products
``kron(I_n, A) + kron(B, I_m)``, where ``I_n`` and ``I_m`` are identity
matrices.
Args:
A (cupyx.scipy.sparse.spmatrix): a sparse matrix.
B (cupyx.scipy.sparse.spmatrix): a sparse matrix.
format (str): the format of the returned sparse matrix.
Returns:
cupyx.scipy.sparse.spmatrix:
Generated sparse matrix with the specified ``format``.
.. seealso:: :func:`scipy.sparse.kronsum`
"""
A = coo.coo_matrix(A)
B = coo.coo_matrix(B)
if A.shape[0] != A.shape[1]:
raise ValueError('A is not square matrix')
if B.shape[0] != B.shape[1]:
raise ValueError('B is not square matrix')
dtype = sputils.upcast(A.dtype, B.dtype)
L = kron(eye(B.shape[0], dtype=dtype), A, format=format)
R = kron(B, eye(A.shape[0], dtype=dtype), format=format)
return (L + R).asformat(format)
def bmat(blocks, format=None, dtype=None):
"""Builds a sparse matrix from sparse sub-blocks
Args:
blocks (array_like):
Grid of sparse matrices with compatible shapes.
An entry of None implies an all-zero matrix.
format ({'bsr', 'coo', 'csc', 'csr', 'dia', 'dok', 'lil'}, optional):
The sparse format of the result (e.g. "csr"). By default an
appropriate sparse matrix format is returned.
This choice is subject to change.
dtype (dtype, optional):
The data-type of the output matrix. If not given, the dtype is
determined from that of `blocks`.
Returns:
bmat (sparse matrix)
.. seealso:: :func:`scipy.sparse.bmat`
Examples:
>>> from cupy import array
>>> from cupyx.scipy.sparse import csr_matrix, bmat
>>> A = csr_matrix(array([[1., 2.], [3., 4.]]))
>>> B = csr_matrix(array([[5.], [6.]]))
>>> C = csr_matrix(array([[7.]]))
>>> bmat([[A, B], [None, C]]).toarray()
array([[1., 2., 5.],
[3., 4., 6.],
[0., 0., 7.]])
>>> bmat([[A, None], [None, C]]).toarray()
array([[1., 2., 0.],
[3., 4., 0.],
[0., 0., 7.]])
"""
# We assume here that blocks will be 2-D so we need to look, at most,
# 2 layers deep for the shape
# TODO(Corey J. Nolet): Check this assumption and raise ValueError
# NOTE: We can't follow scipy exactly here
# since we don't have an `object` datatype
M = len(blocks)
N = len(blocks[0])
blocks_flat = []
for m in range(M):
for n in range(N):
if blocks[m][n] is not None:
blocks_flat.append(blocks[m][n])
if len(blocks_flat) == 0:
return coo.coo_matrix((0, 0), dtype=dtype)
# check for fast path cases
if (N == 1 and format in (None, 'csr') and
all(isinstance(b, csr.csr_matrix)
for b in blocks_flat)):
A = _compressed_sparse_stack(blocks_flat, 0)
if dtype is not None:
A = A.astype(dtype)
return A
elif (M == 1 and format in (None, 'csc')
and all(isinstance(b, csc.csc_matrix) for b in blocks_flat)):
A = _compressed_sparse_stack(blocks_flat, 1)
if dtype is not None:
A = A.astype(dtype)
return A
block_mask = numpy.zeros((M, N), dtype=bool)
brow_lengths = numpy.zeros(M+1, dtype=numpy.int64)
bcol_lengths = numpy.zeros(N+1, dtype=numpy.int64)
# convert everything to COO format
for i in range(M):
for j in range(N):
if blocks[i][j] is not None:
A = coo.coo_matrix(blocks[i][j])
blocks[i][j] = A
block_mask[i][j] = True
if brow_lengths[i+1] == 0:
brow_lengths[i+1] = A.shape[0]
elif brow_lengths[i+1] != A.shape[0]:
msg = ('blocks[{i},:] has incompatible row dimensions. '
'Got blocks[{i},{j}].shape[0] == {got}, '
'expected {exp}.'.format(i=i, j=j,
exp=brow_lengths[i+1],
got=A.shape[0]))
raise ValueError(msg)
if bcol_lengths[j+1] == 0:
bcol_lengths[j+1] = A.shape[1]
elif bcol_lengths[j+1] != A.shape[1]:
msg = ('blocks[:,{j}] has incompatible row dimensions. '
'Got blocks[{i},{j}].shape[1] == {got}, '
'expected {exp}.'.format(i=i, j=j,
exp=bcol_lengths[j+1],
got=A.shape[1]))
raise ValueError(msg)
nnz = sum(block.nnz for block in blocks_flat)
if dtype is None:
all_dtypes = [blk.dtype for blk in blocks_flat]
dtype = sputils.upcast(*all_dtypes) if all_dtypes else None
row_offsets = numpy.cumsum(brow_lengths)
col_offsets = numpy.cumsum(bcol_lengths)
shape = (row_offsets[-1], col_offsets[-1])
data = cupy.empty(nnz, dtype=dtype)
idx_dtype = sputils.get_index_dtype(maxval=max(shape))
row = cupy.empty(nnz, dtype=idx_dtype)
col = cupy.empty(nnz, dtype=idx_dtype)
nnz = 0
ii, jj = numpy.nonzero(block_mask)
for i, j in zip(ii, jj):
B = blocks[int(i)][int(j)]
idx = slice(nnz, nnz + B.nnz)
data[idx] = B.data
row[idx] = B.row + row_offsets[i]
col[idx] = B.col + col_offsets[j]
nnz += B.nnz
return coo.coo_matrix((data, (row, col)), shape=shape).asformat(format)