turn sparsified matrix into sparse format

dpeerlab · Nov 26, 2024 · b0f949e · b0f949e
1 parent 898c8fa
commit b0f949e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/palantir/utils.py b/src/palantir/utils.py
@@ -652,7 +652,12 @@ def run_magic_imputation(
 
     # Set small values to zero if returning sparse matrix 
     if sparse:
-        imputed_data[imputed_data < clip_threshold] = 0
+        if issparse(X):
+            imputed_data.data[imputed_data.data < clip_threshold] = 0
+            imputed_data.eliminate_zeros()
+        else:
+            imputed_data = np.where(imputed_data < clip_threshold, 0, imputed_data)
+            imputed_data = csr_matrix(imputed_data)
 
     # Clean up
     gc.collect()