ENH DeseqStats n_cpus (#379)

* n_cpu argument for DeseqStats

* Match n_cpus logic in DeseqStats to the logic in DeseqDataSet

* Add n_cpus to test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: asistradition

Diff for: pydeseq2/dds.py

@@ -321,16 +321,17 @@ def __init__(
         self.filtered_genes = None
 
         if inference:
-            if hasattr(inference, "n_cpus"):
-                if n_cpus:
+            if n_cpus:
+                if hasattr(inference, "n_cpus"):
                     inference.n_cpus = n_cpus
-            else:
-                warnings.warn(
-                    "The provided inference object does not have an n_cpus "
-                    "attribute, cannot override `n_cpus`.",
-                    UserWarning,
-                    stacklevel=2,
-                )
+                else:
+                    warnings.warn(
+                        "The provided inference object does not have an n_cpus "
+                        "attribute, cannot override `n_cpus`.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+
         # Initialize the inference object.
         self.inference = inference or DefaultInference(n_cpus=n_cpus)
 

Diff for: pydeseq2/ds.py

@@ -1,5 +1,6 @@
 import sys
 import time
+import warnings
 from typing import Literal
 
 # import anndata as ad
@@ -141,6 +142,7 @@ def __init__(
         ) = None,
         inference: Inference | None = None,
         quiet: bool = False,
+        n_cpus: int | None = None,
     ) -> None:
         assert (
             "LFC" in dds.varm
@@ -188,8 +190,20 @@ def __init__(
         self.shrunk_LFCs = False
         self.quiet = quiet
 
+        if inference:
+            if n_cpus:
+                if hasattr(inference, "n_cpus"):
+                    inference.n_cpus = n_cpus
+                else:
+                    warnings.warn(
+                        "The provided inference object does not have an n_cpus "
+                        "attribute, cannot override `n_cpus`.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+
         # Initialize the inference object.
-        self.inference = inference or DefaultInference()
+        self.inference = inference or DefaultInference(n_cpus=n_cpus)
 
         # If the `refit_cooks` attribute of the dds object is True, check that outliers
         # were actually refitted.

Diff for: tests/test_pydeseq2.py

@@ -194,9 +194,7 @@ def test_alt_hypothesis(alt_hypothesis, counts_df, metadata, tol=0.02):
     )
 
     dds = DeseqDataSet(
-        counts=counts_df,
-        metadata=metadata,
-        design="~condition",
+        counts=counts_df, metadata=metadata, design="~condition", n_cpus=2
     )
     dds.deseq2()
 
@@ -205,6 +203,7 @@ def test_alt_hypothesis(alt_hypothesis, counts_df, metadata, tol=0.02):
         contrast=["condition", "B", "A"],
         lfc_null=-0.5 if alt_hypothesis == "less" else 0.5,
         alt_hypothesis=alt_hypothesis,
+        n_cpus=2,
     )
     ds.summary()