Repository: owkin/PyDESeq2
Pull request #292 (2024-06-11 15:22:54)

ENH Remove count DataFrame from calculate_cooks
#### What does your PR implement? Be specific.

`calculate_cooks` casts `normed_counts` into a pandas DataFrame for `robust_method_of_moments_disp`. This is memory inefficient for large data.

`robust_method_of_moments_disp` has been refactored to accept an ndarray directly and the DataFrame has been removed. There is no numerical change as a result.

asistradition: The main advantage to using `.obsm` is that column slicing a row-major array requires a copy, and there is a considerable amount of overhead when calling `.layers[key][:, filtered_genes]` repeatedly.

As those keys are only used in `cooks_distances` it is a considerable optimization (for large data, e.g. 50k x 30k) to move them to `.obsm` remove those copies.
asistradition: Also includes an optional `control_genes` argument to `fit_size_factors` which has the same behavior as the `controlGenes` argument to `estimateSizeFactors` and the associated unit test

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e4af33d..d431605 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,12 +7,12 @@ repos:
       args: # arguments to configure black
        - --line-length=89
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.10
+    rev: v0.5.0
     hooks:
     - id: ruff
       args: [ --fix, --exit-non-zero-on-fix]
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.0
+    rev: v1.10.1
     hooks:
     - id: mypy
       exclude: ^(tests/|docs/source/conf.py)
\ No newline at end of file
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index 80c3606..a3d5dc8 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -499,7 +499,11 @@ class DeseqDataSet(ad.AnnData):
             self.refit()
 
     def fit_size_factors(
-        self, fit_type: Literal["ratio", "poscounts", "iterative"] = "ratio"
+        self,
+        fit_type: Literal["ratio", "poscounts", "iterative"] = "ratio",
+        control_genes: Optional[
+            Union[np.ndarray, List[str], List[int], pd.Index]
+        ] = None,
     ) -> None:
         """Fit sample-wise deseq2 normalization (size) factors.
 
@@ -512,21 +516,44 @@ class DeseqDataSet(ad.AnnData):
         have no zero values. In this situation, size factors can depend on a very small
         number of features (or only one feature) leading to incorrect inference. This
         method for calculating size factors will only exclude genes which have all-0
-        values (and are not amenable to inference anyway)
+        values (and are not amenable to inference anyway).
 
         The "poscounts" method calculates the n-th root of the product of the non-zero
-        (positive) counts
+        (positive) counts.
+
+        Control genes can be optionally provided; if so, size factors will be fit to
+        only the genes in this argument. This is the same functionality as controlGenes
+        in R DESeq2. Any valid AnnData indexer (bool, int position, var_name string) is
+        accepted.
 
         Parameters
         ----------
         fit_type : str
-            The normalization method to use (default: ``"ratio"``).
+            The normalization method to use: "ratio", "poscounts" or "iterative".
+            (default: ``"ratio"``).
+        control_genes : ndarray, list, pandas.Index, or None
+            Genes to use as control genes for size factor fitting. If None, all genes
+            are used. (default: ``None``).
         """
         if not self.quiet:
             print("Fitting size factors...", file=sys.stderr)
 
         start = time.time()
 
+        # If control genes are provided, set a mask where those genes are True
+        if control_genes is not None:
+            _control_mask = np.zeros(self.X.shape[1], dtype=bool)
+
+            # Use AnnData internal indexing to get gene index array
+            # Allows bool/int/var_name to be provided
+            _control_mask[self._normalize_indices((slice(None), control_genes))[1]] = (
+                True
+            )
+
+        # Otherwise mask all genes to be True
+        else:
+            _control_mask = np.ones(self.X.shape[1], dtype=bool)
+
         if fit_type == "iterative":
             self._fit_iterate_size_factors()
 
@@ -539,10 +566,11 @@ class DeseqDataSet(ad.AnnData):
 
             # Determine which genes are usable (finite logmeans)
             self.filtered_genes = (~np.isinf(logmeans)) & (logmeans > 0)
+            _control_mask &= self.filtered_genes
 
             # Calculate size factor per sample
             def sizeFactor(x):
-                _mask = np.logical_and(self.filtered_genes, x > 0)
+                _mask = np.logical_and(_control_mask, x > 0)
                 return np.exp(np.median(np.log(x[_mask]) - logmeans[_mask]))
 
             sf = np.apply_along_axis(sizeFactor, 1, self.X)
@@ -566,10 +594,12 @@ class DeseqDataSet(ad.AnnData):
 
         else:
             self.logmeans, self.filtered_genes = deseq2_norm_fit(self.X)
+            _control_mask &= self.filtered_genes
+
             (
                 self.layers["normed_counts"],
                 self.obsm["size_factors"],
-            ) = deseq2_norm_transform(self.X, self.logmeans, self.filtered_genes)
+            ) = deseq2_norm_transform(self.X, self.logmeans, _control_mask)
 
         end = time.time()
         self.varm["_normed_means"] = self.layers["normed_counts"].mean(0)
@@ -838,11 +868,8 @@ class DeseqDataSet(ad.AnnData):
             )
         )
 
-        self.layers["_mu_LFC"] = np.full((self.n_obs, self.n_vars), np.nan)
-        self.layers["_mu_LFC"][:, self.varm["non_zero"]] = mu_
-
-        self.layers["_hat_diagonals"] = np.full((self.n_obs, self.n_vars), np.nan)
-        self.layers["_hat_diagonals"][:, self.varm["non_zero"]] = hat_diagonals_
+        self.obsm["_mu_LFC"] = mu_
+        self.obsm["_hat_diagonals"] = hat_diagonals_
 
         self.varm["_LFC_converged"] = np.full(self.n_vars, np.nan)
         self.varm["_LFC_converged"][self.varm["non_zero"]] = converged_
@@ -856,34 +883,48 @@ class DeseqDataSet(ad.AnnData):
         if "dispersions" not in self.varm:
             self.fit_MAP_dispersions()
 
-        num_vars = self.obsm["design_matrix"].shape[-1]
+        if not self.quiet:
+            print("Calculating cook's distance...", file=sys.stderr)
 
-        # Keep only non-zero genes
-        nonzero_data = self[:, self.non_zero_genes]
-        normed_counts = pd.DataFrame(
-            nonzero_data.X / self.obsm["size_factors"][:, None],
-            index=self.obs_names,
-            columns=self.non_zero_genes,
-        )
+        start = time.time()
+        num_vars = self.obsm["design_matrix"].shape[-1]
 
+        # Calculate dispersion
         dispersions = robust_method_of_moments_disp(
-            normed_counts, self.obsm["design_matrix"]
+            self.layers["normed_counts"][:, self.varm["non_zero"]],
+            self.obsm["design_matrix"],
         )
 
-        V = (
-            nonzero_data.layers["_mu_LFC"]
-            + dispersions.values[None, :] * nonzero_data.layers["_mu_LFC"] ** 2
-        )
-        squared_pearson_res = (nonzero_data.X - nonzero_data.layers["_mu_LFC"]) ** 2 / V
-        diag_mul = (
-            nonzero_data.layers["_hat_diagonals"]
-            / (1 - nonzero_data.layers["_hat_diagonals"]) ** 2
-        )
+        # Calculate the squared pearson residuals for non-zero features
+        squared_pearson_res = self.X[:, self.varm["non_zero"]] - self.obsm["_mu_LFC"]
+        squared_pearson_res **= 2
+
+        # Calculate the overdispersion parameter tau
+        V = self.obsm["_mu_LFC"] ** 2
+        V *= dispersions[None, :]
+        V += self.obsm["_mu_LFC"]
+
+        # Calculate r^2 / (tau * num_vars)
+        squared_pearson_res /= V
+        squared_pearson_res /= num_vars
+
+        del V
+
+        # Calculate leverage modifier H / (1 - H)^2
+        diag_mul = 1 - self.obsm["_hat_diagonals"]
+        diag_mul **= 2
+        diag_mul = self.obsm["_hat_diagonals"] / diag_mul
+
+        # Multiply r^2 / (tau * num_vars) by H / (1 - H)^2 to get cook's distance
+        squared_pearson_res *= diag_mul
+
+        del diag_mul
 
         self.layers["cooks"] = np.full((self.n_obs, self.n_vars), np.nan)
-        self.layers["cooks"][:, self.varm["non_zero"]] = (
-            squared_pearson_res / num_vars * diag_mul
-        )
+        self.layers["cooks"][:, self.varm["non_zero"]] = squared_pearson_res
+
+        if not self.quiet:
+            print(f"... done in {time.time()-start:.2f} seconds.\n", file=sys.stderr)
 
     def refit(self) -> None:
         """Refit Cook outliers.
diff --git a/pydeseq2/utils.py b/pydeseq2/utils.py
index ca95d1c..95d93ac 100644
--- a/pydeseq2/utils.py
+++ b/pydeseq2/utils.py
@@ -736,14 +736,14 @@ def fit_alpha_mle(
         )
 
 
-def trimmed_mean(x, trim: float = 0.1, **kwargs) -> Union[float, np.ndarray]:
+def trimmed_mean(x: np.ndarray, trim: float = 0.1, **kwargs) -> Union[float, np.ndarray]:
     """Return trimmed mean.
 
     Compute the mean after trimming data of its smallest and largest quantiles.
 
     Parameters
     ----------
-    features : ndarray
+    x : ndarray
         Data whose mean to compute.
 
     trim : float
@@ -771,7 +771,7 @@ def trimmed_mean(x, trim: float = 0.1, **kwargs) -> Union[float, np.ndarray]:
         return s[ntrim : n - ntrim].mean()
 
 
-def trimmed_cell_variance(counts: pd.DataFrame, cells: pd.Series) -> pd.Series:
+def trimmed_cell_variance(counts: np.ndarray, cells: pd.Series) -> np.ndarray:
     """Return trimmed variance of counts according to condition.
 
     Compute the variance after trimming data of its smallest and largest elements,
@@ -780,7 +780,7 @@ def trimmed_cell_variance(counts: pd.DataFrame, cells: pd.Series) -> pd.Series:
 
     Parameters
     ----------
-    counts : pandas.DataFrame
+    counts : ndarray
         Sample-wise gene counts.
 
     cells : pandas.Series
@@ -788,7 +788,7 @@ def trimmed_cell_variance(counts: pd.DataFrame, cells: pd.Series) -> pd.Series:
 
     Returns
     -------
-    pandas.Series :
+    ndarray :
         Gene-wise trimmed variance estimate.
     """
     # how much to trim at different n
@@ -799,23 +799,27 @@ def trimmed_cell_variance(counts: pd.DataFrame, cells: pd.Series) -> pd.Series:
         return 2 if x >= 23.5 else 1 if x >= 3.5 else 0
 
     ns = cells.value_counts()
-    cell_means = pd.DataFrame(index=counts.columns)
+    sqerror = np.zeros_like(counts)
+
     for lvl in cells.unique():
-        cell_means[lvl] = trimmed_mean(
-            counts.loc[cells == lvl], trim=trimratio[trimfn(ns[lvl])], axis=0
+        cell_means = cast(
+            np.ndarray,
+            trimmed_mean(
+                counts[cells == lvl, :], trim=trimratio[trimfn(ns[lvl])], axis=0
+            ),
         )
-    qmat = cell_means[cells].T
-    qmat.index = cells.index
-    sqerror = (counts - qmat) ** 2
+        sqerror[cells == lvl, :] = counts[cells == lvl, :] - cell_means[None, :]
 
-    varEst = pd.DataFrame(index=counts.columns)
-    for lvl in cells.unique():
+    sqerror **= 2
+
+    varEst = np.zeros((len(ns), counts.shape[1]), dtype=float)
+    for i, lvl in enumerate(cells.unique()):
         scale = [2.04, 1.86, 1.51][trimfn(ns[lvl])]
-        varEst[lvl] = scale * trimmed_mean(
-            sqerror[cells == lvl], trim=trimratio[trimfn(ns[lvl])], axis=0
+        varEst[i, :] = scale * trimmed_mean(
+            sqerror[cells == lvl, :], trim=trimratio[trimfn(ns[lvl])], axis=0
         )
 
-    return varEst.max(axis=1)
+    return varEst.max(axis=0)
 
 
 def trimmed_variance(
@@ -1080,16 +1084,16 @@ def n_or_more_replicates(design_matrix: pd.DataFrame, min_replicates: int) -> pd
 
 
 def robust_method_of_moments_disp(
-    normed_counts: pd.DataFrame, design_matrix: pd.DataFrame
-) -> pd.Series:
+    normed_counts: np.ndarray, design_matrix: pd.DataFrame
+) -> np.ndarray:
     """Perform dispersion estimation using a method of trimmed moments.
 
     Used for outlier detection based on Cook's distance.
 
     Parameters
     ----------
-    normed_counts : pandas.DataFrame
-        DF of deseq2-normalized read counts. Rows: samples, columns: genes.
+    normed_counts : ndarray
+        Array of deseq2-normalized read counts. Rows: samples, columns: genes.
 
     design_matrix : pandas.DataFrame
         A DataFrame with experiment design information (to split cohorts).
@@ -1097,7 +1101,7 @@ def robust_method_of_moments_disp(
 
     Returns
     -------
-    pandas.Series
+    ndarray
         Trimmed method of moment dispersion estimates.
         Used for outlier detection based on Cook's distance.
     """
@@ -1107,7 +1111,7 @@ def robust_method_of_moments_disp(
         # 1 - group rows by unique combinations of design factors
         # 2 - keep only groups with 3 or more replicates
         # 3 - filter the counts matrix to only keep rows in those groups
-        filtered_counts = normed_counts.loc[three_or_more, :]
+        filtered_counts = normed_counts[three_or_more.values, :]
         filtered_design = design_matrix.loc[three_or_more, :]
         cell_id = pd.Series(
             filtered_design.groupby(
@@ -1117,15 +1121,14 @@ def robust_method_of_moments_disp(
         )
         v = trimmed_cell_variance(filtered_counts, cell_id)
     else:
-        v = pd.Series(
-            trimmed_variance(normed_counts.values), index=normed_counts.columns
-        )
+        v = trimmed_variance(normed_counts)
+
     m = normed_counts.mean(0)
     alpha = (v - m) / m**2
     # cannot use the typical min_disp = 1e-8 here or else all counts in the same
     # group as the outlier count will get an extreme Cook's distance
     minDisp = 0.04
-    alpha = cast(pd.Series, np.maximum(alpha, minDisp))
+    np.maximum(alpha, minDisp, out=alpha)
     return alpha
 
 


Test patch:
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 64f6717..18e5701 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -68,6 +68,35 @@ def test_size_factors_poscounts(counts_df, metadata):
     np.testing.assert_almost_equal(dds.obsm["size_factors"].ravel(), r_size_factors)
 
 
+def test_size_factors_control_genes(counts_df, metadata):
+    """Test that the size_factors calculation properly takes control_genes"""
+
+    dds = DeseqDataSet(counts=counts_df, metadata=metadata, design_factors="condition")
+
+    dds.fit_size_factors(control_genes=["gene4"])
+
+    np.testing.assert_almost_equal(
+        dds.obsm["size_factors"].ravel(),
+        counts_df["gene4"] / np.exp(np.log(counts_df["gene4"]).mean()),
+    )
+
+    dds.fit_size_factors(fit_type="poscounts", control_genes=[3])
+
+    np.testing.assert_almost_equal(
+        dds.obsm["size_factors"].ravel(),
+        counts_df["gene4"] / np.exp(np.log(counts_df["gene4"]).mean()),
+    )
+
+    dds.fit_size_factors(fit_type="poscounts")
+
+    np.testing.assert_raises(
+        AssertionError,
+        np.testing.assert_array_equal,
+        dds.obsm["size_factors"].ravel(),
+        counts_df["gene4"] / np.exp(np.log(counts_df["gene4"]).mean()),
+    )
+
+
 def test_deseq_independent_filtering_parametric_fit(counts_df, metadata, tol=0.02):
     """Test that the outputs of the DESeq2 function match those of the original R
     package, up to a tolerance in relative error.