R/matrix_functions.R

In Timothy-Barry/ravel:

#' Create example sparse matrix
#'
#' Creates an example sparse matrix.
#'
#' @param n_row number of rows in matrix
#' @param n_col number of columns in matrix
#' @param p_zero (optional, default 0.8) proportion of entries to set to zero (in expectation)
#' @param matrix_values (optional, default 1:9) a vector of numeric values from which to sample
#'
#' @return
#' A random matrix of dimension n_row by n_col.
#' @export
#'
#' @examples
#' create_example_sparse_matrix(10, 5)
#'
create_example_sparse_matrix <- function(n_row, n_col, p_zero = 0.8, matrix_values = 1:9) {
  m <- matrix(data = sample(x = matrix_values, size = n_row * n_col, replace = TRUE), nrow = n_row, ncol = n_col)
  r <- matrix(data = rbinom(n =  n_row * n_col, size = 1, prob = 1 - p_zero), nrow = n_row, ncol = n_col)
  out <- m * r
  return(out)
}


#' Get compressed matrix representation
#'
#' @param A a matrix
#' @param row_or_column_storage "row" or "column", indicating whether to store A in compressed row or compressed column format.
#'
#' @return
#' a list containing components p (pointer into the indexes and data), i (the indexes), x (the data), dim (the dimension of the matrix), row_or_column_storage ("row" or "column", depending on storage type), zero_idx (boolean indicating whether the vectors p and i are 0-index-based).
#' @export
#'
#' @examples
#' A <- create_example_sparse_matrix(10, 5)
#' get_compressed_matrix_representation(A, "column")
get_compressed_matrix_representation <- function(A, row_or_column_storage) {
  if (row_or_column_storage == "row") A <- t(A)
  sparse_m <- Matrix(data = A, sparse = TRUE)
  dim <- if (row_or_column_storage == "column") sparse_m@Dim else c(sparse_m@Dim[2], sparse_m@Dim[1])
  out <- list(p = sparse_m@p, i = sparse_m@i, x = sparse_m@x, dim = dim, row_or_column_storage = row_or_column_storage, zero_idx = TRUE)
  return(out)
}


#' Extract vector from compressed matrix
#'
#' A function to extract a vector(s) from a matrix stored in compressed column or row format
#'
#' @param vector_nos indexes of the vectors to extract
#' @param x sparse-matrix data
#' @param i sparse-matrix indexes
#' @param p sparse-matrix pointers
#' @param vector_length length of the vectors to extract
#' @param zero_based_idx (boolean; default TRUE) do i and p use zero-based indexing?
#' @param vector_subset (optional) integers indicating a subset of the vectors to return
#'
#' @return the extracted vectors
#' @export
extract_vector_from_compressed_matrix <- function(vector_nos, x, i, p, vector_length, zero_based_idx = TRUE, vector_subset = NULL) {
  if (zero_based_idx) {
    i <- i + 1
    p <- p + 1
  }
  f <- function(no) {
    curr_ptr_range <- p[no:(no + 1)] + c(0,-1)
    out <- rep(0, vector_length)
    if (curr_ptr_range[1] <= curr_ptr_range[2]) {
      x_out <- x[curr_ptr_range[1]:curr_ptr_range[2]]
      i_out <- i[curr_ptr_range[1]:curr_ptr_range[2]]
      out[i_out] <- x_out
    }
    if (!is.null(vector_subset)) out <- out[vector_subset]
    return(out)
  }
  ret_val <- if (length(vector_nos) == 1) f(vector_nos) else sapply(X = vector_nos, FUN = f)
  return(ret_val)
}


#' Extract column from compressed sparse column matrix
#'
#' @param col_nos idxs of the columns to extract
#' @param csc_mat a list representing a matrix stored in csc form containing the entries p, i, x, and dim.
#' @param zero_based_idx (boolean) do i and p use zero-based indexing
#' @param row_subset (optional) subset the extracted columns by row index
#'
#' @return a matrix containing the extracted columns
#' @export
extract_column_from_csc_matrix <- function(col_nos, csc_mat, zero_based_idx = TRUE, row_subset = NULL) {
  p <- csc_mat$p
  i <- csc_mat$i
  x <- csc_mat$x
  n_rows <- csc_mat$dim[1]
  out <- extract_vector_from_compressed_matrix(vector_nos = col_nos, x = x, i = i, p = p, vector_length = n_rows, zero_based_idx = zero_based_idx, vector_subset = row_subset)
  return(out)
}

#' Extract row from compressed sparse row matrix
#'
#' This function is similar to extract_column_from_csc_matrix; see the documentation of that function.
#'
#' @export
extract_row_from_csr_matrix <- function(row_nos, csr_mat, zero_based_idx = TRUE, col_subset = NULL) {
  x <- csr_mat$x
  i <- csr_mat$i
  p <- csr_mat$p
  n_cols <- csr_mat$dim[2]
  out <- extract_vector_from_compressed_matrix(vector_nos = row_nos, x = x, i = i, p = p, vector_length = n_cols, zero_based_idx = zero_based_idx, vector_subset = col_subset)
  if (is.matrix(out)) out <- t(out)
  return(out)
}


#' Flip a compressed matrix
#'
#' Converts compressed sparse row format to compressed sparse column format and vice versa. The unerlying matrix is the same; this function simply reorganizes the data to allow for easier column/row access.
#'
#' @param x the vector of data
#' @param row_idx the vector of row indexes
#' @param col_ptr the column pointer
#' @param n_col number of columns
#' @param n_row number of rows
#' @param zero_based_idx (boolean; default TRUE) does the input matrix use zero-based indexing?
#'
#' @return the transposed matrix; specifically, a list with the elements p, i, x, dim, and zero_based_idx.
#' @export
#'
#' @examples
#' A <- create_example_sparse_matrix(10, 5)
#' A_csc <- get_compressed_matrix_representation(A, "column")
#' x <- A_csc$x
#' row_idx <- A_csc$i # row index
#' col_ptr <- A_csc$p # column pointer
#' n_col <- A_csc$dim[2]
#' n_row <- A_csc$dim[1]
#' flip_compressed_matrix(x, row_idx, col_ptr, n_col, n_row)
flip_compressed_matrix <- function(x, row_idx, col_ptr, n_col, n_row, zero_based_idx = TRUE) {
  if (zero_based_idx) {
    row_idx <- row_idx + 1
    col_ptr <- col_ptr + 1
  }
  # first, count the number of non-zero elements in each row
  n_nonzero_per_row <- rep(0, n_row)
  for (curr_i in row_idx) {
    n_nonzero_per_row[curr_i] <- n_nonzero_per_row[curr_i] + 1
  }
  # then, define the row_ptr
  row_ptr <- c(1, cumsum(n_nonzero_per_row) + 1)
  col_idx <- integer(length = length(x))
  new_x <- integer(length = length(x))
  # finally, fill new_x and col_indx
  for (col_no in 1:n_col) {
    col_ptr_start <- col_ptr[col_no]
    col_ptr_end <- col_ptr[col_no + 1] -1
    if (col_ptr_start <= col_ptr_end) {
      for (data_idx in col_ptr_start:col_ptr_end) {
        # extract the current datum
        x_curr <- x[data_idx]
        row_curr <- row_idx[data_idx]
        # insert into new matrix
        new_data_idx <- row_ptr[row_curr]
        col_idx[new_data_idx] <- col_no
        new_x[new_data_idx] <- x_curr
        # increment row_ptr
        row_ptr[row_curr] <- row_ptr[row_curr] + 1
      }
    }
  }
  row_ptr <- c(1, row_ptr[-length(row_ptr)])
  if (zero_based_idx) {
    col_idx <- col_idx - 1
    row_ptr <- row_ptr - 1
  }
  list(p = row_ptr, i = col_idx, x = new_x, dim = c(n_row, n_col), zero_based_idx = zero_based_idx)
}


#' Flip compressed sparse column matrix
#'
#' Flip a matrix in compressed sparse column format to compressed sparse row format
#'
#' @param csc_mat a matrix in compressed sparse column format
#' @param zero_based_idx (boolean) are the i and p vectors 0-indexed?
#' @export
flip_csc_matrix <- function(csc_mat, zero_based_idx = TRUE) {
  flip_compressed_matrix(x = csc_mat$x, row_idx = csc_mat$i, col_ptr = csc_mat$p, n_col = csc_mat$dim[2], n_row = csc_mat$dim[1], zero_based_idx = csc_mat$zero_idx)
}

#' Flip compressed sparse row matrix
#'
#' Flip a matrix in compressed spares row format to compressed sparse column format
#'
#' @export
flip_csr_matrix <- function(csr_mat, zero_based_idx = TRUE) {
  out <- flip_compressed_matrix(x = csr_mat$x, row_idx = csr_mat$i, col_ptr = csr_mat$p, n_col = csr_mat$dim[1], n_row = csr_mat$dim[2], zero_based_idx = csr_mat$zero_idx)
  out$dim <- c(csr_mat$dim[1], csr_mat$dim[2])
  out
}

A <- create_example_sparse_matrix(n_row = 6, n_col = 8, p_zero = 0.8)
get_triplet_representation <- function(A) {
  triplet_m <- as(A, "dgTMatrix")
}