From c6d95a1e264160ed4f765e17c40994b882a86654 Mon Sep 17 00:00:00 2001
From: WANG Zhiwei <48282751+statwangz@users.noreply.github.com>
Date: Sat, 16 Sep 2023 00:56:47 +0800
Subject: [PATCH] Support sparse matrix mode

---
 R/mfairInitialization.R | 24 +++++++++++++++++++++---
 R/mfairObject.R         |  6 ++++--
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/R/mfairInitialization.R b/R/mfairInitialization.R
index f428442..34e79b2 100644
--- a/R/mfairInitialization.R
+++ b/R/mfairInitialization.R
@@ -20,16 +20,33 @@ createMFAIR <- function(Y, X, Y_center = TRUE, K_max = 1L, project = "MFAIR") {
     stop("The number of samples in Y and X should be consistent!")
   } # End
 
+  if(is.matrix(Y)){
+    Y_sparse <- FALSE
+  }else{
+    Y_sparse <- TRUE
+  }
+
   # Center the matrix Y
   if (Y_center) {
-    Y_mean <- mean(Y, na.rm = TRUE)
-    Y <- Y - Y_mean
+    if(!Y_sparse){
+      Y_mean <- mean(Y, na.rm = TRUE)
+      Y <- Y - Y_mean
+    }
+else{
+  Y_mean <- mean(summary(Y)$x)
+  Y@x <- Y@x - Y_mean
+}
   } else {
     Y_mean <- 0
   }
 
   # Check Y's sparsity
-  n_missing <- sum(is.na(Y))
+  if(!Y_sparse){
+    n_missing <- sum(is.na(Y))
+  }
+  else{
+    n_missing <- length(Y) - length(Y@x)
+  }
   if (n_missing >= 1) {
     if (n_missing == length(Y)) {
       stop("The main data matrix Y has no observed values!")
@@ -52,6 +69,7 @@ createMFAIR <- function(Y, X, Y_center = TRUE, K_max = 1L, project = "MFAIR") {
     Class = "MFAIR",
     Y = Y,
     X = as.data.frame(X),
+    Y_sparse = Y_sparse,
     Y_center = Y_center,
     Y_mean = Y_mean,
     Y_missing = Y_missing,
diff --git a/R/mfairObject.R b/R/mfairObject.R
index 38e386e..f8852d2 100644
--- a/R/mfairObject.R
+++ b/R/mfairObject.R
@@ -1,8 +1,9 @@
 #' Each MFAIR object has a number of slots which store information. Key slots to access are listed below.
 #'
 #'
-#' @slot Y A matrix. The main data matrix of N samples and M features.
+#' @slot Y A matrix or Matrix::dgCMatrix. The main data matrix of N samples and M features.
 #' @slot X A data.frame. The auxiliary information data frame of N samples and C covariates.
+#' @slot Y_sparse Logical. Whether the main data matrix Y is stored in the sparse mode.
 #' @slot Y_center Logical. Whether the main data matrix Y is centered.
 #' @slot Y_mean Numeric. Mean of the main data matrix Y if centered. Zero if not.
 #' @slot Y_missing Logical. Whether the main data matrix Y is partially observed.
@@ -34,8 +35,9 @@ setClass(
 
   # Define the slots
   slots = c(
-    Y = "matrix",
+    Y = "matrixORdgCMatrix",
     X = "data.frame",
+    Y_sparse = "logical",
     Y_center = "logical",
     Y_mean = "numeric",
     Y_missing = "logical",