def filter_sparse_rows(sp_mat, max_cols):
  """Filter rows of a CSR sparse matrix to retain upto max_cols."""
  all_remove_idx = []
  num_removed = np.zeros((sp_mat.indptr.shape[0] - 1,), dtype=np.int64)
  for ii in range(sp_mat.shape[0]):
    row_st = sp_mat.indptr[ii]
    row_en = sp_mat.indptr[ii + 1]
    my_scores = sp_mat.data[row_st:row_en]
    if len(my_scores) > max_cols:
      remove = np.argpartition(-my_scores, max_cols)[max_cols:]
      all_remove_idx.extend(remove + row_st)
    num_removed[ii] = len(all_remove_idx)
  new_data = np.delete(sp_mat.data, all_remove_idx)
  new_indices = np.delete(sp_mat.indices, all_remove_idx)
  new_indptr = np.copy(sp_mat.indptr)
  new_indptr[1:] -= num_removed
  return sp.csr_matrix((new_data, new_indices, new_indptr), shape=sp_mat.shape)
Filter rows of a CSR sparse matrix to retain upto max_cols.
language/labs/drkit/preprocessing/preprocess_utils.py