Refactored scanpy.tools._sparse_nanmean to eliminate unnecessary data copying

Author: Reovirus

src/scanpy/tools/_score_genes.py

@@ -34,19 +34,19 @@ def _sparse_nanmean(X: CSBase, axis: Literal[0, 1]) -> NDArray[np.float64]:
         msg = "X must be a compressed sparse matrix"
         raise TypeError(msg)
 
-    # count the number of nan elements per row/column (dep. on axis)
     Z = X.copy()
-    Z.data = np.isnan(Z.data)
-    Z.eliminate_zeros()
-    n_elements = Z.shape[axis] - Z.sum(axis)
 
-    # set the nans to 0, so that a normal .sum() works
-    Y = X.copy()
-    Y.data[np.isnan(Y.data)] = 0
-    Y.eliminate_zeros()
+    # count the number of nonzero elements (include nans) per row/column (dep. on axis)
+    nonzeros_and_nones=Z.count_nonzero(axis=axis)
+
+    # just sum the data withput nan
+    Z.data[np.isnan(Z.data)] = 0
+    Z.eliminate_zeros()
+    s = Z.sum(axis, dtype="float64")
 
-    # the average
-    s = Y.sum(axis, dtype="float64")  # float64 for score_genes function compatibility)
+    # Z.count_nonzero(axis=axis) is now non-zero not-nan elements in X
+    #diff between nonzeros_and_nones and curr nonzero is nans
+    n_elements=(Z.shape[axis] - (nonzeros_and_nones - Z.count_nonzero(axis=axis))).reshape(s.shape, copy=False)
     m = s / n_elements
 
     return m