Repository: owkin/PyDESeq2
Pull request #284 (2024-05-30 19:01:43)

Poscount implementation & diag(XXT) optimization
Thank you for implementing DESeq2 in python. I have been using it for single-cell data.

Included is an implementation of the poscount size_factor method from DESeq2 (R).

Included is an enhancement in `irls_solver` that replaces the full matrix multiplication to get `diag(X(XTX)-1XT)` with a more memory and runtime-efficient calculation of only the diagonal values. The result is identical.



asistradition: I am also adding an optimization for `wald_test` that replaces the diagonal weight matrix `W` with a broadcast vector multiplication

Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 7dd39f0..cc0264b 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,1 +1,1 @@
-* @BorisMuzellec @maikia @arthurPignetOwkin @mandreux-owkin
+* @BorisMuzellec @maikia @arthurPignetOwkin @mandreux-owkin @umarteauowkin
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index aafd3db..a385f1d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,7 +7,7 @@ repos:
       args: # arguments to configure black
        - --line-length=89
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.5
+    rev: v0.4.8
     hooks:
     - id: ruff
       args: [ --fix, --exit-non-zero-on-fix]
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index 44f06bc..47eee3b 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -498,7 +498,7 @@ class DeseqDataSet(ad.AnnData):
             self.refit()
 
     def fit_size_factors(
-        self, fit_type: Literal["ratio", "iterative"] = "ratio"
+        self, fit_type: Literal["ratio", "poscounts", "iterative"] = "ratio"
     ) -> None:
         """Fit sample-wise deseq2 normalization (size) factors.
 
@@ -506,6 +506,16 @@ class DeseqDataSet(ad.AnnData):
         unless each gene has at least one sample with zero read counts, in which case it
         switches to the ``iterative`` method.
 
+        Also available is the 'poscounts' method implemented in DESeq2 for the
+        single-cell or metagenomics use case where there may be few or no features which
+        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)
+
+        The "poscounts" method calculates the n-th root of the product of the non-zero
+        (positive) counts
+
         Parameters
         ----------
         fit_type : str
@@ -513,9 +523,35 @@ class DeseqDataSet(ad.AnnData):
         """
         if not self.quiet:
             print("Fitting size factors...", file=sys.stderr)
+
         start = time.time()
+
         if fit_type == "iterative":
             self._fit_iterate_size_factors()
+
+        elif fit_type == "poscounts":
+
+            # Calculate logcounts for x > 0 and take the mean for each gene
+            log_counts = np.zeros_like(self.X, dtype=float)
+            np.log(self.X, out=log_counts, where=self.X != 0)
+            logmeans = log_counts.mean(0)
+
+            # Determine which genes are usable (finite logmeans)
+            self.filtered_genes = (~np.isinf(logmeans)) & (logmeans > 0)
+
+            # Calculate size factor per sample
+            def sizeFactor(x):
+                _mask = np.logical_and(self.filtered_genes, x > 0)
+                return np.exp(np.median(np.log(x[_mask]) - logmeans[_mask]))
+
+            sf = np.apply_along_axis(sizeFactor, 1, self.X)
+            del log_counts
+
+            # Normalize size factors to a geometric mean of 1 to match DESeq
+            self.obsm["size_factors"] = sf / (np.exp(np.mean(np.log(sf))))
+            self.layers["normed_counts"] = self.X / self.obsm["size_factors"][:, None]
+            self.logmeans = logmeans
+
         # Test whether it is possible to use median-of-ratios.
         elif (self.X == 0).any(0).all():
             # There is at least a zero for each gene
@@ -526,14 +562,15 @@ class DeseqDataSet(ad.AnnData):
                 stacklevel=2,
             )
             self._fit_iterate_size_factors()
+
         else:
             self.logmeans, self.filtered_genes = deseq2_norm_fit(self.X)
             (
                 self.layers["normed_counts"],
                 self.obsm["size_factors"],
             ) = deseq2_norm_transform(self.X, self.logmeans, self.filtered_genes)
-        end = time.time()
 
+        end = time.time()
         self.varm["_normed_means"] = self.layers["normed_counts"].mean(0)
 
         if not self.quiet:
diff --git a/pydeseq2/ds.py b/pydeseq2/ds.py
index c127de8..eabbcac 100644
--- a/pydeseq2/ds.py
+++ b/pydeseq2/ds.py
@@ -260,16 +260,17 @@ class DeseqStats:
         self.results_df["pvalue"] = self.p_values
         self.results_df["padj"] = self.padj
 
-        if self.contrast[1] == self.contrast[2] == "":
-            # The factor is continuous
-            print(f"Log2 fold change & Wald test p-value: " f"{self.contrast[0]}")
-        else:
-            # The factor is categorical
-            print(
-                f"Log2 fold change & Wald test p-value: "
-                f"{self.contrast[0]} {self.contrast[1]} vs {self.contrast[2]}"
-            )
-        print(self.results_df)
+        if not self.quiet:
+            if self.contrast[1] == self.contrast[2] == "":
+                # The factor is continuous
+                print(f"Log2 fold change & Wald test p-value: " f"{self.contrast[0]}")
+            else:
+                # The factor is categorical
+                print(
+                    f"Log2 fold change & Wald test p-value: "
+                    f"{self.contrast[0]} {self.contrast[1]} vs {self.contrast[2]}"
+                )
+            print(self.results_df)
 
     def run_wald_test(self) -> None:
         """Perform a Wald test.
diff --git a/pydeseq2/utils.py b/pydeseq2/utils.py
index 9d3a8c8..ca95d1c 100644
--- a/pydeseq2/utils.py
+++ b/pydeseq2/utils.py
@@ -593,10 +593,17 @@ def irls_solver(
         dev_ratio = np.abs(dev - old_dev) / (np.abs(dev) + 0.1)
 
     # Compute H diagonal (useful for Cook distance outlier filtering)
+    # Calculate only the diagonal for X(XTWX)-1XT using einsum
+    # This is numerically equivalent to the more expensive calculation
+    # np.diag(X @ (X^T @ np.inv(X^T @ np.diag(W) @ X + lambda) @ X^T)
     W = mu / (1.0 + mu * disp)
+    H = np.einsum(
+        "ij,jk,ki->i", X, np.linalg.inv((X.T * W[None, :]) @ X + ridge_factor), X.T
+    )
+
     W_sq = np.sqrt(W)
-    XtWX = (X.T * W) @ X + ridge_factor
-    H = W_sq * np.diag(X @ np.linalg.inv(XtWX) @ X.T) * W_sq
+    H = W_sq * H * W_sq
+
     # Return an UNthresholded mu (as in the R code)
     # Previous quantities are estimated with a threshold though
     mu = size_factors * np.exp(X @ beta)
@@ -880,7 +887,7 @@ def wald_test(
     design_matrix: np.ndarray,
     disp: float,
     lfc: np.ndarray,
-    mu: float,
+    mu: np.ndarray,
     ridge_factor: np.ndarray,
     contrast: np.ndarray,
     lfc_null: float,
@@ -902,7 +909,7 @@ def wald_test(
     lfc : ndarray
         Log-fold change estimate (in natural log scale).
 
-    mu : float
+    mu : ndarray
         Mean estimation for the NB model.
 
     ridge_factor : ndarray
@@ -929,8 +936,8 @@ def wald_test(
         Standard error of the Wald statistic.
     """
     # Build covariance matrix estimator
-    W = np.diag(mu / (1 + mu * disp))
-    M = design_matrix.T @ W @ design_matrix
+    W = mu / (1 + mu * disp)
+    M = (design_matrix.T * W[None, :]) @ design_matrix
     H = np.linalg.inv(M + ridge_factor)
     Hc = H @ contrast
     # Evaluate standard error and Wald statistic


Test patch:
diff --git a/tests/data/single_factor/r_test_size_factors_poscount.csv b/tests/data/single_factor/r_test_size_factors_poscount.csv
new file mode 100644
index 0000000..80f6448
--- /dev/null
+++ b/tests/data/single_factor/r_test_size_factors_poscount.csv
@@ -0,0 +1,101 @@
+"","sizeFactor"
+"sample1",1.12868391752277
+"sample2",0.906262068755213
+"sample3",0.834189989589321
+"sample4",0.933676593163132
+"sample5",0.672473009851897
+"sample6",1.01583900659264
+"sample7",1.02446649272126
+"sample8",1.07461918414369
+"sample9",0.69904449508282
+"sample10",1.29713011611685
+"sample11",0.973512204866054
+"sample12",1.34701973596749
+"sample13",0.796678160455071
+"sample14",1.23690937424838
+"sample15",1.07698079281936
+"sample16",0.940667358408543
+"sample17",1.23799601765783
+"sample18",0.963902791269114
+"sample19",1.37188012877762
+"sample20",0.929700189747951
+"sample21",0.784930536577835
+"sample22",0.670017050733884
+"sample23",1.14163988798133
+"sample24",1.02920734976071
+"sample25",1.00425847253746
+"sample26",0.789508630118609
+"sample27",1.07821722355883
+"sample28",1.20097085979113
+"sample29",0.943587559085369
+"sample30",0.703273158410287
+"sample31",1.51930608774199
+"sample32",1.20599435311314
+"sample33",0.929078405857473
+"sample34",0.79696896493937
+"sample35",1.04815188517397
+"sample36",0.867822390357068
+"sample37",1.40386481015313
+"sample38",1.48620347284242
+"sample39",1.1183358418485
+"sample40",1.05309482244205
+"sample41",0.755615509447497
+"sample42",0.751412663608778
+"sample43",1.11695371705217
+"sample44",1.41540660853722
+"sample45",1.03210710820019
+"sample46",0.69845467790907
+"sample47",1.04158383403793
+"sample48",1.26462263410147
+"sample49",1.07523377909479
+"sample50",0.880974756102388
+"sample51",0.919812498538184
+"sample52",1.85377815812203
+"sample53",0.984261878847999
+"sample54",0.882177728366319
+"sample55",0.85443207246528
+"sample56",1.37188012877762
+"sample57",1.59049346900441
+"sample58",1.96039495192515
+"sample59",1.01828943882477
+"sample60",0.889668269994777
+"sample61",0.943479245051929
+"sample62",0.846512938142076
+"sample63",1.21403491251153
+"sample64",1.00546910954697
+"sample65",0.863505672614189
+"sample66",1.41085489690346
+"sample67",0.88784140085361
+"sample68",1.52473560041543
+"sample69",1.01416837381677
+"sample70",1.09320905234401
+"sample71",0.700950758177703
+"sample72",0.988814287334817
+"sample73",1.11188296123575
+"sample74",1.22294953706274
+"sample75",0.879964407876309
+"sample76",0.76086144586774
+"sample77",0.750284576941207
+"sample78",0.948466975576099
+"sample79",1.02446649272126
+"sample80",0.764401231719072
+"sample81",1.20410113160285
+"sample82",0.80451507819027
+"sample83",1.05617793338424
+"sample84",1.07821722355883
+"sample85",0.468747773198709
+"sample86",0.677071163896507
+"sample87",0.914586752518411
+"sample88",1.39278263659514
+"sample89",0.957481704325058
+"sample90",0.852479793557742
+"sample91",1.48620347284242
+"sample92",1.12868391752277
+"sample93",0.851428075390773
+"sample94",1.37418667233149
+"sample95",0.898820033940812
+"sample96",1.05156527864003
+"sample97",0.693157916621703
+"sample98",0.689319636513982
+"sample99",1.209435142853
+"sample100",0.747374543462677
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index fa750ae..64f6717 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -35,6 +35,39 @@ def metadata():
     )
 
 
+def test_size_factors_ratio(counts_df, metadata):
+    """Test that the size_factors calcuation by ratio matches R."""
+
+    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
+
+    r_size_factors = pd.read_csv(
+        os.path.join(test_path, "data/single_factor/r_test_size_factors.csv"),
+        index_col=0,
+    )["x"].values
+
+    dds = DeseqDataSet(counts=counts_df, metadata=metadata, design_factors="condition")
+    dds.fit_size_factors()
+
+    np.testing.assert_almost_equal(
+        dds.obsm["size_factors"].ravel(), r_size_factors.ravel()
+    )
+
+
+def test_size_factors_poscounts(counts_df, metadata):
+    """Test that the size_factors calcuation by poscount matches R."""
+    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
+
+    dds = DeseqDataSet(counts=counts_df, metadata=metadata, design_factors="condition")
+    dds.fit_size_factors("poscounts")
+
+    r_size_factors = pd.read_csv(
+        os.path.join(test_path, "data/single_factor/r_test_size_factors_poscount.csv"),
+        index_col=0,
+    )["sizeFactor"].values
+
+    np.testing.assert_almost_equal(dds.obsm["size_factors"].ravel(), r_size_factors)
+
+
 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.