Sparse matrices

import numpy as np
import scipy.sparse as sps

def make_full_csr_matrix(shape, fill_value):
  """
  Creates a full csr matrix.
  Parameters:
    shape ((int,int)) : shape of the matrix
    fill_value (int) : matrix elements are set ot this value
  Returns:
    (scipy.sparse.csr_matrix()) : sparse matrix filled with the specified value
  """
# pass shape parameters
  nrow, ncol = shape
# create uniform data of the correct number
  data = np.full((nrow * ncol), fill_value = 1, dtype = np.int)
# column indices
  indices = np.tile(np.arange(ncol), nrow)
# number of nonzero elements in the rows
  indptr = np.arange(nrow + 1, dtype = np.int) * ncol
# create matrix
  return sps.csr_matrix((data, indices, indptr))

def make_full_csr_matrix_list(shapes, fill_value):
  """
  Generates a list of csr sparse matrices. The matrices are full, but in sparse format, 
  so that they can easily be processed by sparse matrix constructors.
  Parameters:
    shapes (sequence of tuples) : the sequence of shapes
    fill_value (int) : all matrix elements will have this value
  Returns:
    matrices [scipy.sparse.csr_matrix]: list of csr matrices
  """
  matrices = []

# --- iterate through shapes
  for shape in shapes:
    matrices.append(make_full_csr_matrix(shape, fill_value))

  return matrices

class FullCsrMatrixFactory(object):
  """
  Creates an instance of a generator object which generates a sequence of csr matrices.
  The matrices are full, but in sparse format, hence they can easily be processed by sparse matrix constructors.
  Attributes:
    shapes ([(,)]) : list of tuples containing the shapes of the matrices
    fill_value (int) : each element will be set to this value
  """
# -------
  def __init__(self, shapes, fill_value):
    """
    shapes ([(,)]) : list of tuples containing the shapes of the matrices
    fill_value (int) : each element will be set to this value
    """
# set shapes of matrices
    try:
      _ = len(shapes)
    except Exception as err:
      print("parameter 'shapes' must implement the '__len__()' method")
      raise err

    self._shapes = shapes

# set value for elements
    self._fill_value = fill_value

# set counter to zero
    self.__ispent = 0

  @property
  def fill_value(self):
    return self._fill_value

  @property
  def shapes(self):
    return self._shapes

  @property
  def ispent(self):
    return self.__ispent

# -------
  def __iter__(self):
    """
    Returns a sparse matrix
    """
    while self.ispent < len(self):
    
      _shape = self.shapes[self.ispent]
      yield make_full_csr_matrix(_shape, self.fill_value)
      self.__ispent += 1

# -------
  def __len__(self):
    """
    Define a len method, so that the number of matrices can be known in advance.
    """
    return len(self.shapes)

def create_block_diagonal_csr_matrix_sps(block_sizes, factory,
                                         fill_value = 1, 
                                         keep_density = 0.1,
                                         keep_each_block = True):
  """
  Creates a block diagonal csr matrix.
  Parameters:
    block_sizes ([int]) : sizes of the blocks.
    factory (callable) : should generate a list of csr matrices.
      signature ~([(int,int)], float) --> iter(scipy.sparse.csr_matrix)

    fill_value (int) : the value of the elements. Default 1.
    keep_density (float) : the proportion of rows to be kept. Default 0.01.
    keep_each_block (bool) : whether to keep at least one row from each block. Default True
  """
# number of rows to be kept
  block_sizes_ = np.array(block_sizes, dtype = np.int)
  n_keep_rows = np.rint(block_sizes_ * keep_density).astype(np.int)
# keep one row from each block all blocks
  if keep_each_block:
    n_keep_rows[n_keep_rows == 0] = 1
# create shapes for blocks with more than zero rows
  shapes = list(filter(lambda x: x[0] > 0, zip(n_keep_rows, block_sizes_)))
# set up generator for the sequence of matrices
  mats = factory(shapes, fill_value)

# create a blockdiagonal matrix by concatenating the blocks
  adj_mat = sps.block_diag(mats, format = 'csr')

  return adj_mat

block_size = [10, 50, 100, 500]
shape_list = [[x] * x for x in block_size]
keep_density = 0.1

from time import perf_counter
performance = []

for shapes in shape_list:
    tstart = perf_counter()
    matrix = create_block_diagonal_csr_matrix_sps(shapes, make_full_csr_matrix_list,
                                                  keep_density = keep_density)
    tdelta = perf_counter() - tstart
    res_dict = {'matrix_size' : matrix.size}
    res_dict.update({'time' : tdelta})
    performance.append(res_dict)
    
for res_dict in performance:
   print("\t".join(["{0} : {1}".format(_k, _v) for _k, _v in res_dict.items()]))

performance = []

for shapes in shape_list:
    tstart = perf_counter()
    matrix = create_block_diagonal_csr_matrix_sps(shapes, FullCsrMatrixFactory,
                                                  keep_density = keep_density)
    tdelta = perf_counter() - tstart
    res_dict = {'matrix_size' : matrix.size}
    res_dict.update({'time' : tdelta})
    performance.append(res_dict)
    
for res_dict in performance:
   print("\t".join(["{0} : {1}".format(_k, _v) for _k, _v in res_dict.items()]))

% timeit -n 3 np.full((500*500*50), 1)

%prun -l 15 -s "cumulative" create_block_diagonal_csr_matrix_sps(shape_list[-1], make_full_csr_matrix_list, keep_density = keep_density)

%prun -l 15 -s "cumulative" create_block_diagonal_csr_matrix_sps(shape_list[-1], FullCsrMatrixFactory, keep_density = keep_density)

def create_block_diagonal_csr_matrix_np(block_sizes, fill_value = 1, keep_density = 0.1, keep_each_block = True):
  """
  Creates a block diagonal csr matrix.
  Parameters:
    block_sizes ([int]) : sizes of the blocks.
    fill_value (int) : the value of the elements. Default 1.
    keep_density (float) : the proportion of rows to be kept. Default 0.01.
    keep_each_block (bool) : whether to keep at least one row from each block. Default True
  """
  block_sizes_ = np.array(block_sizes, dtype = np.int)
# number of columns in each block
  n_cols = np.array(block_sizes, dtype = np.int)
# number of rows in each block
  n_rows = np.rint(block_sizes_ * keep_density).astype(np.int)

# keep one row from each block all blocks
  if keep_each_block:
    n_rows[n_rows == 0] = 1

# discard empty blocks
  n_cols = n_cols[n_rows > 0]
  n_rows = n_rows[n_rows > 0]

# total number of rows, columns and nonzero elements
  num_tot_rows = np.sum(n_rows)
  num_tot_cols = np.sum(n_cols)
  num_tot_nnz = np.sum(n_cols * n_rows)

# create empty matrix with the right dimensions
  mat = sps.csr_matrix((num_tot_rows, num_tot_cols), dtype = np.int)

# set cumulative number of nnz elements for each row
  mat.indptr[1:] = np.cumsum(np.repeat(n_cols, n_rows))

# set column indices in each row
  mat.indices = np.zeros(num_tot_nnz, dtype = np.int)
  offset = 0 # offset of column indices for each block
  ilow = 0  # position of start index in indices for each row

# fill in column indices for each row successively
  for ncol, nrow in zip(n_cols, n_rows):
    ihgh = ilow + nrow * ncol
    mat.indices[ilow:ihgh] = np.tile(np.arange(offset, offset + ncol), nrow)
    offset += ncol
    ilow = ihgh + 0

# set data
  mat.data = np.full(num_tot_nnz, fill_value)

  return mat

performance = []

shape_list.extend([[x] * x for x in [1000, 2000]])

for shapes in shape_list:
    tstart = perf_counter()
    matrix = create_block_diagonal_csr_matrix_np(shapes, keep_density = keep_density)
    tdelta = perf_counter() - tstart
    res_dict = {'matrix_size' : matrix.size}
    res_dict.update({'time' : tdelta})
    performance.append(res_dict)
    
for res_dict in performance:
   print("\t".join(["{0} : {1}".format(_k, _v) for _k, _v in res_dict.items()]))