Repository: owkin/PyDESeq2
Pull request #324 (2024-10-17 19:30:13)

ENH: Memory Optimizations & low_memory Flag
#### What does your PR implement? Be specific.

The `replace_cooks` data was stored as a pandas DataFrame. This is memory inefficient for large data. It has been refactored from a DataFrame to a numpy array.

A new `low_memory` argument to `DeseqDataSet` is available. When set to True, large data arrays saved into the AnnData elements `.obsm` and `.layers` are deleted if there is no further use for them in the standard deseq workflow. This reduces peak memory consumption dramatically. 



asistradition: Yes, this is part of my standard single-cell workflow, and memory limitations are the main bottleneck. 60k x 35k won't be a problem with this PR.

Patch:
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index b421a50..2b9f653 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -115,6 +115,11 @@ class DeseqDataSet(ad.AnnData):
     quiet : bool
         Suppress deseq2 status updates during fit.
 
+    low_memory : bool
+        Remove intermediate data structures from .layers and from .obsm that are no
+        longer necessary after they are used during deseq2 run, such as Cook's
+        distances. (default: False)
+
     Attributes
     ----------
     X
@@ -193,6 +198,7 @@ class DeseqDataSet(ad.AnnData):
         n_cpus: Optional[int] = None,
         inference: Optional[Inference] = None,
         quiet: bool = False,
+        low_memory: bool = False,
     ) -> None:
         # Initialize the AnnData part
         if adata is not None:
@@ -281,6 +287,7 @@ class DeseqDataSet(ad.AnnData):
         self.min_replicates = min_replicates
         self.beta_tol = beta_tol
         self.quiet = quiet
+        self.low_memory = low_memory
         self.logmeans = None
         self.filtered_genes = None
 
@@ -500,6 +507,9 @@ class DeseqDataSet(ad.AnnData):
             # for genes that had outliers replaced
             self.refit()
 
+        # Compute gene mask for cooks outliers
+        self.cooks_outlier()
+
     def fit_size_factors(
         self,
         fit_type: Literal["ratio", "poscounts", "iterative"] = "ratio",
@@ -667,6 +677,8 @@ class DeseqDataSet(ad.AnnData):
         self.layers[mu_param_name] = np.full((self.n_obs, self.n_vars), np.nan)
         self.layers[mu_param_name][:, self.varm["non_zero"]] = mu_hat_
 
+        del mu_hat_
+
         if not self.quiet:
             print("Fitting dispersions...", file=sys.stderr)
         start = time.time()
@@ -826,6 +838,9 @@ class DeseqDataSet(ad.AnnData):
             "genewise_dispersions"
         ][self.varm["_outlier_genes"]]
 
+        if self.low_memory:
+            del self.layers["_mu_hat"]
+
     def fit_LFC(self) -> None:
         """Fit log fold change (LFC) coefficients.
 
@@ -921,6 +936,10 @@ class DeseqDataSet(ad.AnnData):
 
         del diag_mul
 
+        if self.low_memory:
+            del self.obsm["_mu_LFC"]
+            del self.obsm["_hat_diagonals"]
+
         self.layers["cooks"] = np.full((self.n_obs, self.n_vars), np.nan)
         self.layers["cooks"][:, self.varm["non_zero"]] = squared_pearson_res
 
@@ -951,6 +970,52 @@ class DeseqDataSet(ad.AnnData):
                 False,
             )
 
+    def cooks_outlier(self):
+        """Filter p-values based on Cooks outliers."""
+        if "_pvalue_cooks_outlier" in self.varm.keys():
+            return self.varm["_pvalue_cooks_outlier"]
+
+        num_samples = self.n_obs
+        num_vars = self.obsm["design_matrix"].shape[-1]
+        cooks_cutoff = f.ppf(0.99, num_vars, num_samples - num_vars)
+
+        # As in DESeq2, only take samples with 3 or more replicates when looking for
+        # max cooks.
+        use_for_max = n_or_more_replicates(self.obsm["design_matrix"], 3)
+
+        # If for a gene there are 3 samples or more that have more counts than the
+        # maximum cooks sample, don't count this gene as an outlier.
+
+        # Take into account whether we already replaced outliers
+        if (
+            self.refit_cooks
+            and (self.varm["refitted"].sum() > 0)
+            and "replace_cooks" in self.layers.keys()
+        ):
+            cooks_outlier = (
+                self.layers["replace_cooks"][use_for_max, :] > cooks_cutoff
+            ).any(axis=0)
+
+        else:
+            cooks_outlier = (self.layers["cooks"][use_for_max, :] > cooks_cutoff).any(
+                axis=0
+            )
+
+        pos = self.layers["cooks"][:, cooks_outlier].argmax(0)
+
+        cooks_outlier[cooks_outlier] = (
+            self.X[:, cooks_outlier] > self.X[:, cooks_outlier][pos, np.arange(len(pos))]
+        ).sum(0) < 3
+
+        if self.low_memory:
+            del self.layers["cooks"]
+
+        if self.low_memory and "replace_cooks" in self.layers.keys():
+            del self.layers["replace_cooks"]
+
+        self.varm["_pvalue_cooks_outlier"] = cooks_outlier
+        return self.varm["_pvalue_cooks_outlier"]
+
     def _fit_MoM_dispersions(self) -> None:
         """Rough method of moments initial dispersions fit.
 
@@ -1266,10 +1331,10 @@ class DeseqDataSet(ad.AnnData):
         ]
         self.varm["dispersions"][self.varm["refitted"]] = sub_dds.varm["dispersions"]
 
-        replace_cooks = pd.DataFrame(self.layers["cooks"].copy())
-        replace_cooks.loc[self.obsm["replaceable"], self.varm["refitted"]] = 0.0
+        self.layers["replace_cooks"] = self.layers["cooks"].copy()
 
-        self.layers["replace_cooks"] = replace_cooks
+        for col in np.where(self.varm["refitted"])[0]:
+            self.layers["replace_cooks"][self.obsm["replaceable"], col] = 0.0
 
     def _fit_iterate_size_factors(self, niter: int = 10, quant: float = 0.95) -> None:
         """
diff --git a/pydeseq2/ds.py b/pydeseq2/ds.py
index 186548d..49f7c20 100644
--- a/pydeseq2/ds.py
+++ b/pydeseq2/ds.py
@@ -8,7 +8,6 @@ from typing import Optional
 import numpy as np
 import pandas as pd
 from scipy.optimize import root_scalar  # type: ignore
-from scipy.stats import f  # type: ignore
 from scipy.stats import false_discovery_control  # type: ignore
 
 from pydeseq2.dds import DeseqDataSet
@@ -16,7 +15,6 @@ from pydeseq2.default_inference import DefaultInference
 from pydeseq2.inference import Inference
 from pydeseq2.utils import lowess
 from pydeseq2.utils import make_MA_plot
-from pydeseq2.utils import n_or_more_replicates
 
 
 class DeseqStats:
@@ -570,40 +568,7 @@ class DeseqStats:
         if not hasattr(self, "p_values"):
             self.run_wald_test()
 
-        num_samples = self.dds.n_obs
-        num_vars = self.design_matrix.shape[-1]
-        cooks_cutoff = f.ppf(0.99, num_vars, num_samples - num_vars)
-
-        # As in DESeq2, only take samples with 3 or more replicates when looking for
-        # max cooks.
-        use_for_max = n_or_more_replicates(self.design_matrix, 3)
-
-        # If for a gene there are 3 samples or more that have more counts than the
-        # maximum cooks sample, don't count this gene as an outlier.
-
-        # Take into account whether we already replaced outliers
-        if self.dds.refit_cooks and self.dds.varm["refitted"].sum() > 0:
-            cooks_outlier = (
-                (self.dds[use_for_max, :].layers["replace_cooks"] > cooks_cutoff)
-                .any(axis=0)
-                .copy()
-            )
-
-        else:
-            cooks_outlier = (
-                (self.dds[use_for_max, :].layers["cooks"] > cooks_cutoff)
-                .any(axis=0)
-                .copy()
-            )
-
-        pos = self.dds[:, cooks_outlier].layers["cooks"].argmax(0)
-
-        cooks_outlier[cooks_outlier] = (
-            self.dds[:, cooks_outlier].X
-            > self.dds[:, cooks_outlier].X[pos, np.arange(len(pos))]
-        ).sum(0) < 3
-
-        self.p_values[cooks_outlier] = np.nan
+        self.p_values[self.dds.cooks_outlier()] = np.nan
 
     def _fit_prior_var(
         self, coeff_idx: str, min_var: float = 1e-6, max_var: float = 400.0


Test patch:
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 18e5701..e5657be 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -564,7 +564,9 @@ def test_continuous_lfc_shrinkage(tol=0.02):
     ).max() < tol
 
 
+@pytest.mark.parametrize("low_memory", [True, False])
 def test_wide_deseq(
+    low_memory,
     tol=0.02,
 ):
     """Test that the outputs of the DESeq2 function match those of the original R
@@ -589,6 +591,7 @@ def test_wide_deseq(
         counts=counts_df,
         metadata=metadata,
         design_factors=["group", "condition"],
+        low_memory=low_memory,
     )
     dds.deseq2()