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)