From 43fccb96b986f9da2c3a4320fe58693ca660193b Mon Sep 17 00:00:00 2001
From: Yichen Wang <wayichen@umich.edu>
Date: Thu, 4 Apr 2024 10:11:22 -0400
Subject: [PATCH] update c-iNMF strategy, with bug fix

---
 R/cINMF.R          | 33 ++++++++++++++++++---------------
 src/cinmf_util.cpp |  3 +++
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/R/cINMF.R b/R/cINMF.R
index dc504e6e..26299ab9 100644
--- a/R/cINMF.R
+++ b/R/cINMF.R
@@ -368,13 +368,16 @@ runCINMF.Seurat <- function(
     # The first is for W, the rest are for Vs
     selection <- factor_cluster_sel(geneLoadings, nNeighbor = nNeighbor,
                                     minWeight = 0.6, k = k, resolution = 0.2)
-    W <- geneLoadings[[1]][, selection$idx]
+    W <- Reduce(cbind, Ws)[, selection$idx]
+    # W <- geneLoadings[[1]][, selection$idx]
     W <- colAggregateMedian_dense_cpp(W, group = selection$cluster, n = k)
-    W <- colNormalize_dense_cpp(W, L = 1)
+    # W <- colNormalize_dense_cpp(W, L = 1)
     for (i in seq_along(Vs)) {
-        V <- geneLoadings[[i + 1]][, selection$idx]
+        V <- Reduce(cbind, Vs[[i]])[, selection$idx]
+        # V <- geneLoadings[[i + 1]][, selection$idx]
         V <- colAggregateMedian_dense_cpp(V, group = selection$cluster, n = k)
-        Vs[[i]] <- colNormalize_dense_cpp(V, L = 1)
+        # Vs[[i]] <- colNormalize_dense_cpp(V, L = 1)
+        Vs[[i]] <- V
     }
     # Vs <- lapply(seq_along(object), function(i) {
     #     matrix(stats::runif(nrow(W) * k, 0, 2), nrow(W), k)
@@ -486,17 +489,17 @@ inmf_solveV <- function(H, W, V, E, lambda, nCores = 2L) {
     return(t(V))
 }
 
-# inmf_solveW <- function(Hs, W, Vs, Es, lambda, nCores = 2L) {
-#     CtC <- matrix(0, ncol(Vs[[1]]), ncol(Vs[[1]]))
-#     CtB <- matrix(0, ncol(Vs[[1]]), nrow(Vs[[1]]))
-#     for (i in seq_along(Es)) {
-#         HtH <- t(Hs[[i]]) %*% Hs[[i]]
-#         CtC <- CtC + HtH
-#         CtB <- CtB + as.matrix(t(Hs[[i]]) %*% t(Es[[i]])) - HtH %*% t(Vs[[i]])
-#     }
-#     W <- RcppPlanc::bppnnls_prod(CtC, CtB, nCores = nCores)
-#     return(t(W))
-# }
+inmf_solveW <- function(Hs, W, Vs, Es, lambda, nCores = 2L) {
+    CtC <- matrix(0, ncol(Vs[[1]]), ncol(Vs[[1]]))
+    CtB <- matrix(0, ncol(Vs[[1]]), nrow(Vs[[1]]))
+    for (i in seq_along(Es)) {
+        HtH <- t(Hs[[i]]) %*% Hs[[i]]
+        CtC <- CtC + HtH
+        CtB <- CtB + as.matrix(t(Hs[[i]]) %*% t(Es[[i]])) - HtH %*% t(Vs[[i]])
+    }
+    W <- RcppPlanc::bppnnls_prod(CtC, CtB, nCores = nCores)
+    return(t(W))
+}
 
 inmf_objErr_i <- function(H, W, V, E, lambda) {
     # Objective error function was originally stated as:
diff --git a/src/cinmf_util.cpp b/src/cinmf_util.cpp
index d11b73ee..e8345175 100644
--- a/src/cinmf_util.cpp
+++ b/src/cinmf_util.cpp
@@ -9,6 +9,9 @@ arma::mat colNormalize_dense_cpp(arma::mat& x, const arma::uword L) {
     arma::mat result(x);
     for (int j = 0; j < x.n_cols; ++j) {
         double norm = arma::norm(x.col(j), L);
+        if (norm == 0) {
+            continue;
+        }
         for (int i = 0; i < x.n_rows; ++i) {
             result(i, j) /= norm;
         }