def _compressed_sparse_stack(blocks, axis):
"""Fast path for stacking CSR/CSC matrices
(i) vstack for CSR, (ii) hstack for CSC.
"""
other_axis = 1 if axis == 0 else 0
data = cupy.concatenate([b.data for b in blocks])
constant_dim = blocks[0].shape[other_axis]
idx_dtype = sputils.get_index_dtype(arrays=[b.indptr for b in blocks],
maxval=max(data.size, constant_dim))
indices = cupy.empty(data.size, dtype=idx_dtype)
indptr = cupy.empty(sum(b.shape[axis]
for b in blocks) + 1, dtype=idx_dtype)
last_indptr = idx_dtype(0)
sum_dim = 0
sum_indices = 0
for b in blocks:
if b.shape[other_axis] != constant_dim:
raise ValueError(
'incompatible dimensions for axis %d' % other_axis)
indices[sum_indices:sum_indices+b.indices.size] = b.indices
sum_indices += b.indices.size
idxs = slice(sum_dim, sum_dim + b.shape[axis])
indptr[idxs] = b.indptr[:-1]
indptr[idxs] += last_indptr
sum_dim += b.shape[axis]
last_indptr += b.indptr[-1]
indptr[-1] = last_indptr
if axis == 0:
return csr.csr_matrix((data, indices, indptr),
shape=(sum_dim, constant_dim))
else:
return csc.csc_matrix((data, indices, indptr),
shape=(constant_dim, sum_dim))
def reshape(self, *shape, order='C'):
"""Gives a new shape to a sparse matrix without changing its data.
Args:
shape (tuple):
The new shape should be compatible with the original shape.
order: {'C', 'F'} (optional)
Read the elements using this index order. 'C' means to read and
write the elements using C-like index order. 'F' means to read
and write the elements using Fortran-like index order. Default:
C.
Returns:
cupyx.scipy.sparse.coo_matrix: sparse matrix
"""
shape = sputils.check_shape(shape, self.shape)
if shape == self.shape:
return self
nrows, ncols = self.shape
if order == 'C': # C to represent matrix in row major format
dtype = sputils.get_index_dtype(maxval=(ncols * max(0, nrows - 1) +
max(0, ncols - 1)))
flat_indices = cupy.multiply(ncols, self.row,
dtype=dtype) + self.col
new_row, new_col = divmod(flat_indices, shape[1])
elif order == 'F':
dtype = sputils.get_index_dtype(maxval=(ncols * max(0, nrows - 1) +
max(0, ncols - 1)))
flat_indices = cupy.multiply(ncols, self.row,
dtype=dtype) + self.row
new_col, new_row = divmod(flat_indices, shape[0])
else:
raise ValueError("'order' must be 'C' or 'F'")
new_data = self.data
return coo_matrix((new_data, (new_row, new_col)),
shape=shape,
copy=False)
def _insert_many(self, i, j, x):
"""Inserts new nonzero at each (i, j) with value x
Here (i,j) index major and minor respectively.
i, j and x must be non-empty, 1d arrays.
Inserts each major group (e.g. all entries per row) at a time.
Maintains has_sorted_indices property.
Modifies i, j, x in place.
"""
order = cupy.argsort(i) # stable for duplicates
i = i.take(order)
j = j.take(order)
x = x.take(order)
# Update index data type
idx_dtype = sputils.get_index_dtype(
(self.indices, self.indptr), maxval=(
self.nnz + x.size))
self.indptr = self.indptr.astype(idx_dtype)
self.indices = self.indices.astype(idx_dtype)
self.data = self.data.astype(self.dtype)
indptr_inserts, indices_inserts, data_inserts = \
_index._select_last_indices(i, j, x, idx_dtype)
rows, ui_indptr = cupy.unique(indptr_inserts, return_index=True)
to_add = cupy.empty(ui_indptr.size+1, ui_indptr.dtype)
to_add[-1] = j.size
to_add[:-1] = ui_indptr
ui_indptr = to_add
# Compute the counts for each row in the insertion array
row_counts = cupy.zeros(ui_indptr.size-1, dtype=idx_dtype)
cupyx.scatter_add(
row_counts, cupy.searchsorted(rows, indptr_inserts), 1)
self._perform_insert(indices_inserts, data_inserts,
rows, row_counts, idx_dtype)
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)
