Repository: owkin/PyDESeq2 (python)

Issue:
[BUG] The vst() function does not work when the algorithm switches to a mean-based dispersion trend.
**Describe the bug**
Calling the vst() leads to the fit_type variable being hard-coded to the chosen fit_type. Consequently when the vst_transform is called it breaks with:
`
File ".venv/lib/python3.10/site-packages/pydeseq2/dds.py", line 411, in vst_transform
    a0, a1 = self.uns["trend_coeffs"]
KeyError: 'trend_coeffs'
`

**Expected behavior**
When switching to mean based, update fit_type to 'mean' in line 918.
`
self.fit_type = 'mean'
`

Pull request description:
fix : enable vst be used "blindly" and to fit its own dispersion
<!--
Thanks for contributing a pull request! Please ensure you have taken a look at
the contribution guidelines: https://pydeseq2.readthedocs.io/en/latest/usage/contributing.html
Make sure that you have added unit tests if applicable.
-->

#### Reference Issue or PRs
<!--
If this PR is related to an existing issue or PR please reference it, eg:
Fixes #1234.
You can use github keywords as described:
https://github.com/blog/1506-closing-issues-via-pull-requests
-->
Fix #263

#### What does your PR implement? Be specific.
Proposed solution :

- A `self.fit_type` is defined at class initialisation and will by default set the fit_type for DEA and VST.
- If needed, a fit type can be passed to deseq2() and vst() in order to launch the two analysis with a separate fit_type. It will set self.fit_type to the user one.
- self.vst_fit() and self.vst_transform() will call internally the self.fit_type.
- vst_transform() should always be called after vst_fit() or deseq2(). Raise an exception if needed trend_coef have not been computed by deseq2() or vst_fit()
- Add unit tests


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 69d6dea..aafd3db 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.4
+    rev: v0.4.5
     hooks:
     - id: ruff
       args: [ --fix, --exit-non-zero-on-fix]
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index 281657b..44f06bc 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -71,9 +71,12 @@ class DeseqDataSet(ad.AnnData):
         specifying the factor of interest and the reference (control) level against which
         we're testing, e.g. ``["condition", "A"]``. (default: ``None``).
 
-    trend_fit_type : str
-        Either "parametric" or "mean" for the type of fitting of the dispersions
-        trend curve. (default: ``"parametric"``).
+    fit_type: str
+        Either ``"parametric"`` or ``"mean"`` for the type of fitting of dispersions to
+        the mean intensity. ``"parametric"``: fit a dispersion-mean relation via a
+        robust gamma-family GLM. ``"mean"``: use the mean of gene-wise dispersion
+        estimates. Will set the fit type for the DEA and the vst transformation. If
+        needed, it can be set separately for each method.(default: ``"parametric"``).
 
     min_mu : float
         Threshold for mean estimates. (default: ``0.5``).
@@ -157,12 +160,6 @@ class DeseqDataSet(ad.AnnData):
     quiet : bool
         Suppress deseq2 status updates during fit.
 
-    fit_type: str
-        Either "parametric" or "mean" for the type of fitting of dispersions to the
-        mean intensity. "parametric": fit a dispersion-mean relation via a robust
-        gamma-family GLM. "mean": use the mean of gene-wise dispersion estimates.
-        (default: ``"parametric"``).
-
     logmeans: numpy.ndarray
         Gene-wise mean log counts, computed in ``preprocessing.deseq2_norm_fit()``.
 
@@ -174,6 +171,7 @@ class DeseqDataSet(ad.AnnData):
     ----------
     .. bibliography::
         :keyprefix: DeseqDataSet-
+
     """
 
     def __init__(
@@ -185,7 +183,7 @@ class DeseqDataSet(ad.AnnData):
         design_factors: Union[str, List[str]] = "condition",
         continuous_factors: Optional[List[str]] = None,
         ref_level: Optional[List[str]] = None,
-        trend_fit_type: Literal["parametric", "mean"] = "parametric",
+        fit_type: Literal["parametric", "mean"] = "parametric",
         min_mu: float = 0.5,
         min_disp: float = 1e-8,
         max_disp: float = 10.0,
@@ -219,6 +217,8 @@ class DeseqDataSet(ad.AnnData):
                 "Either adata or both counts and metadata arguments must be provided."
             )
 
+        self.fit_type = fit_type
+
         # Convert design_factors to list if a single string was provided.
         self.design_factors = (
             [design_factors] if isinstance(design_factors, str) else design_factors
@@ -271,7 +271,6 @@ class DeseqDataSet(ad.AnnData):
         # Check that the design matrix has full rank
         self._check_full_rank_design()
 
-        self.trend_fit_type = trend_fit_type
         self.min_mu = min_mu
         self.min_disp = min_disp
         self.max_disp = np.maximum(max_disp, self.n_obs)
@@ -299,7 +298,7 @@ class DeseqDataSet(ad.AnnData):
     def vst(
         self,
         use_design: bool = False,
-        fit_type: Literal["parametric", "mean"] = "parametric",
+        fit_type: Optional[Literal["parametric", "mean"]] = None,
     ) -> None:
         """Fit a variance stabilizing transformation, and apply it to normalized counts.
 
@@ -310,47 +309,63 @@ class DeseqDataSet(ad.AnnData):
         use_design : bool
             Whether to use the full design matrix to fit dispersions and the trend curve.
             If False, only an intercept is used. (default: ``False``).
+
         fit_type: str
-            Either "parametric" or "mean" for the type of fitting of dispersions to the
-            mean intensity. "parametric": fit a dispersion-mean relation via a robust
-            gamma-family GLM. "mean": use the mean of gene-wise dispersion estimates.
-            (default: ``"parametric"``).
+            * ``None``: fit_type provided at initialization to fit
+              the dispersions trend curve.
+            * ``"parametric"``: fit a dispersion-mean relation via a robust
+              gamma-family GLM.
+            * ``"mean"``: use the mean of gene-wise dispersion estimates.
+
+            (default: ``None``).
         """
-        self.vst_fit(use_design=use_design, fit_type=fit_type)
+        if fit_type is not None:
+            self.vst_fit_type = fit_type
+        else:
+            self.vst_fit_type = self.fit_type
+
+        print(f"Fit type used for VST : {self.vst_fit_type}")
+
+        self.vst_fit(use_design=use_design)
         self.layers["vst_counts"] = self.vst_transform()
 
     def vst_fit(
         self,
         use_design: bool = False,
-        fit_type: Literal["parametric", "mean"] = "parametric",
     ) -> None:
         """Fit a variance stabilizing transformation.
 
+        This method should be called before `vst_transform`.
+
         Results are stored in ``dds.layers["vst_counts"]``.
 
         Parameters
         ----------
         use_design : bool
             Whether to use the full design matrix to fit dispersions and the trend curve.
-            If False, only an intercept is used. (default: ``False``).
-        fit_type : str
-            Either "parametric" or "mean" for the type of fitting of dispersions to the
-            mean intensity. parametric - fit a dispersion-mean relation via a robust
-            gamma-family GLM. mean - use the mean of gene-wise dispersion estimates.
-            (default: ``"parametric"``).
+            If False, only an intercept is used.
+            Only useful if ``fit_type = "parametric"`.
+            (default: ``False``).
         """
-        self.fit_type = fit_type  # to re-use inside vst_transform
-
         # Start by fitting median-of-ratio size factors if not already present,
         # or if they were computed iteratively
         if "size_factors" not in self.obsm or self.logmeans is None:
             self.fit_size_factors()  # by default, fit_type != "iterative"
 
+        if not hasattr(self, "vst_fit_type"):
+            self.vst_fit_type = self.fit_type
+
         if use_design:
-            # Check that the dispersion trend curve was fitted. If not, fit it.
-            # This will call previous functions in a cascade.
-            if "disp_function" not in self.uns:
-                self.fit_dispersion_trend()
+            if self.vst_fit_type == "parametric":
+                self._fit_parametric_dispersion_trend(vst=True)
+            else:
+                warnings.warn(
+                    "use_design=True is only useful when fit_type='parametric'. ",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                self.fit_genewise_dispersions(vst=True)
+
         else:
             # Reduce the design matrix to an intercept and reconstruct at the end
             self.obsm["design_matrix_buffer"] = self.obsm["design_matrix"].copy()
@@ -358,9 +373,9 @@ class DeseqDataSet(ad.AnnData):
                 1, index=self.obs_names, columns=[["intercept"]]
             )
             # Fit the trend curve with an intercept design
-            self.fit_genewise_dispersions()
-            if self.fit_type == "parametric":
-                self.fit_dispersion_trend()
+            self.fit_genewise_dispersions(vst=True)
+            if self.vst_fit_type == "parametric":
+                self._fit_parametric_dispersion_trend(vst=True)
 
             # Restore the design matrix and free buffer
             self.obsm["design_matrix"] = self.obsm["design_matrix_buffer"].copy()
@@ -381,6 +396,11 @@ class DeseqDataSet(ad.AnnData):
         -------
         numpy.ndarray
             Variance stabilized counts.
+
+        Raises
+        ------
+        RuntimeError
+            If the size factors were not fitted before calling this method.
         """
         if "size_factors" not in self.obsm:
             raise RuntimeError(
@@ -407,8 +427,11 @@ class DeseqDataSet(ad.AnnData):
 
             normed_counts, _ = deseq2_norm_transform(counts, logmeans, filtered_genes)
 
-        if self.fit_type == "parametric":
-            a0, a1 = self.uns["trend_coeffs"]
+        if self.vst_fit_type == "parametric":
+            if "vst_trend_coeffs" not in self.uns:
+                raise RuntimeError("Fit the dispersion curve prior to applying VST.")
+
+            a0, a1 = self.uns["vst_trend_coeffs"]
             return np.log2(
                 (
                     1
@@ -418,8 +441,8 @@ class DeseqDataSet(ad.AnnData):
                 )
                 / (4 * a0)
             )
-        elif self.fit_type == "mean":
-            gene_dispersions = self.varm["genewise_dispersions"]
+        elif self.vst_fit_type == "mean":
+            gene_dispersions = self.varm["vst_genewise_dispersions"]
             use_for_mean = gene_dispersions > 10 * self.min_disp
             mean_disp = trim_mean(gene_dispersions[use_for_mean], proportiontocut=0.001)
             return (
@@ -429,14 +452,30 @@ class DeseqDataSet(ad.AnnData):
             ) / np.log(2)
         else:
             raise NotImplementedError(
-                f"Found fit_type '{self.fit_type}'. Expected 'parametric' or 'mean'."
+                f"Found fit_type '{self.vst_fit_type}'. Expected 'parametric' or 'mean'."
             )
 
-    def deseq2(self) -> None:
+    def deseq2(self, fit_type: Optional[Literal["parametric", "mean"]] = None) -> None:
         """Perform dispersion and log fold-change (LFC) estimation.
 
         Wrapper for the first part of the PyDESeq2 pipeline.
+
+
+        Parameters
+        ----------
+        fit_type : str
+            Either None, ``"parametric"`` or ``"mean"`` for the type of fitting of
+            dispersions to the mean intensity.``"parametric"``: fit a dispersion-mean
+            relation via a robust gamma-family GLM. ``"mean"``: use the mean of
+            gene-wise dispersion estimates.
+
+            If None, the fit_type provided at class initialization is used.
+            (default: ``None``).
         """
+        if fit_type is not None:
+            self.fit_type = fit_type
+            print(f"Using {self.fit_type} fit type.")
+
         # Compute DESeq2 normalization factors using the Median-of-ratios method
         self.fit_size_factors()
         # Fit an independent negative binomial model per gene
@@ -495,13 +534,21 @@ class DeseqDataSet(ad.AnnData):
             ) = deseq2_norm_transform(self.X, self.logmeans, self.filtered_genes)
         end = time.time()
 
+        self.varm["_normed_means"] = self.layers["normed_counts"].mean(0)
+
         if not self.quiet:
             print(f"... done in {end - start:.2f} seconds.\n", file=sys.stderr)
 
-    def fit_genewise_dispersions(self) -> None:
+    def fit_genewise_dispersions(self, vst=False) -> None:
         """Fit gene-wise dispersion estimates.
 
         Fits a negative binomial per gene, independently.
+
+        Parameters
+        ----------
+        vst : bool
+            Whether the dispersion estimates are being fitted as part of the VST
+            pipeline. (default: ``False``).
         """
         # Check that size factors are available. If not, compute them.
         if "size_factors" not in self.obsm:
@@ -545,8 +592,11 @@ class DeseqDataSet(ad.AnnData):
                 beta_tol=self.beta_tol,
             )
 
-        self.layers["_mu_hat"] = np.full((self.n_obs, self.n_vars), np.nan)
-        self.layers["_mu_hat"][:, self.varm["non_zero"]] = mu_hat_
+        mu_param_name = "_vst_mu_hat" if vst else "_mu_hat"
+        disp_param_name = "vst_genewise_dispersions" if vst else "genewise_dispersions"
+
+        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_
 
         if not self.quiet:
             print("Fitting dispersions...", file=sys.stderr)
@@ -554,7 +604,7 @@ class DeseqDataSet(ad.AnnData):
         dispersions_, l_bfgs_b_converged_ = self.inference.alpha_mle(
             counts=self.X[:, self.non_zero_idx],
             design_matrix=design_matrix,
-            mu=self.layers["_mu_hat"][:, self.non_zero_idx],
+            mu=self.layers[mu_param_name][:, self.non_zero_idx],
             alpha_hat=self.varm["_MoM_dispersions"][self.non_zero_idx],
             min_disp=self.min_disp,
             max_disp=self.max_disp,
@@ -564,36 +614,42 @@ class DeseqDataSet(ad.AnnData):
         if not self.quiet:
             print(f"... done in {end - start:.2f} seconds.\n", file=sys.stderr)
 
-        self.varm["genewise_dispersions"] = np.full(self.n_vars, np.nan)
-        self.varm["genewise_dispersions"][self.varm["non_zero"]] = np.clip(
+        self.varm[disp_param_name] = np.full(self.n_vars, np.nan)
+        self.varm[disp_param_name][self.varm["non_zero"]] = np.clip(
             dispersions_, self.min_disp, self.max_disp
         )
 
         self.varm["_genewise_converged"] = np.full(self.n_vars, np.nan)
         self.varm["_genewise_converged"][self.varm["non_zero"]] = l_bfgs_b_converged_
 
-    def fit_dispersion_trend(self) -> None:
+    def fit_dispersion_trend(self, vst: bool = False) -> None:
         """Fit the dispersion trend curve.
 
-        The type of the trend curve is determined by ``self.trend_fit_type``.
+        Parameters
+        ----------
+        vst : bool
+            Whether the dispersion trend curve is being fitted as part of the VST
+            pipeline. (default: ``False``).
         """
+        disp_param_name = "vst_genewise_dispersions" if vst else "genewise_dispersions"
+        fit_type = self.vst_fit_type if vst else self.fit_type
+
         # Check that genewise dispersions are available. If not, compute them.
-        if "genewise_dispersions" not in self.varm:
-            self.fit_genewise_dispersions()
+        if disp_param_name not in self.varm:
+            self.fit_genewise_dispersions(vst)
 
         if not self.quiet:
             print("Fitting dispersion trend curve...", file=sys.stderr)
         start = time.time()
-        self.varm["_normed_means"] = self.layers["normed_counts"].mean(0)
 
-        if self.trend_fit_type == "parametric":
-            self._fit_parametric_dispersion_trend()
-        elif self.trend_fit_type == "mean":
-            self._fit_mean_dispersion_trend()
+        if fit_type == "parametric":
+            self._fit_parametric_dispersion_trend(vst)
+        elif fit_type == "mean":
+            self._fit_mean_dispersion_trend(vst)
         else:
             raise NotImplementedError(
-                f"Expected 'parametric' or 'meand' trend curve fit "
-                f"types, received {self.trend_fit_type}"
+                f"Expected 'parametric' or 'mean' trend curve fit "
+                f"types, received {fit_type}"
             )
         end = time.time()
 
@@ -874,14 +930,25 @@ class DeseqDataSet(ad.AnnData):
             **kwargs,
         )
 
-    def _fit_parametric_dispersion_trend(self):
+    def _fit_parametric_dispersion_trend(self, vst: bool = False):
         r"""Fit the dispersion curve according to a parametric model.
 
         :math:`f(\mu) = \alpha_1/\mu + a_0`.
+
+        Parameters
+        ----------
+        vst : bool
+            Whether the dispersion trend curve is being fitted as part of the VST
+            pipeline. (default: ``False``).
         """
+        disp_param_name = "vst_genewise_dispersions" if vst else "genewise_dispersions"
+
+        if disp_param_name not in self.varm:
+            self.fit_genewise_dispersions(vst)
+
         # Exclude all-zero counts
         targets = pd.Series(
-            self[:, self.non_zero_genes].varm["genewise_dispersions"].copy(),
+            self[:, self.non_zero_genes].varm[disp_param_name].copy(),
             index=self.non_zero_genes,
         )
         covariates = pd.Series(
@@ -916,13 +983,11 @@ class DeseqDataSet(ad.AnnData):
                     stacklevel=2,
                 )
 
-                self._fit_mean_dispersion_trend()
+                self._fit_mean_dispersion_trend(vst)
                 return
 
             # Filter out genes that are too far away from the curve before refitting
-            pred_ratios = (
-                self[:, covariates.index].varm["genewise_dispersions"] / predictions
-            )
+            pred_ratios = self[:, covariates.index].varm[disp_param_name] / predictions
 
             targets.drop(
                 targets[(pred_ratios < 1e-4) | (pred_ratios >= 15)].index,
@@ -932,20 +997,31 @@ class DeseqDataSet(ad.AnnData):
                 covariates[(pred_ratios < 1e-4) | (pred_ratios >= 15)].index,
                 inplace=True,
             )
-        self.uns["trend_coeffs"] = pd.Series(coeffs, index=["a0", "a1"])
 
-        self.varm["fitted_dispersions"] = np.full(self.n_vars, np.nan)
-        self.uns["disp_function_type"] = "parametric"
-        self.varm["fitted_dispersions"][self.varm["non_zero"]] = self.disp_function(
-            self.varm["_normed_means"][self.varm["non_zero"]]
-        )
+        if vst:
+            self.uns["vst_trend_coeffs"] = pd.Series(coeffs, index=["a0", "a1"])
+        else:
+            self.uns["trend_coeffs"] = pd.Series(coeffs, index=["a0", "a1"])
+
+            self.varm["fitted_dispersions"] = np.full(self.n_vars, np.nan)
+            self.uns["disp_function_type"] = "parametric"
+            self.varm["fitted_dispersions"][self.varm["non_zero"]] = self.disp_function(
+                self.varm["_normed_means"][self.varm["non_zero"]]
+            )
+
+    def _fit_mean_dispersion_trend(self, vst: bool = False):
+        """Use the mean of dispersions as trend curve.
+
+        Parameters
+        ----------
+        vst : bool
+            Whether the dispersion trend curve is being fitted as part of the VST
+            pipeline. (default: ``False``).
+        """
+        disp_param_name = "vst_genewise_dispersions" if vst else "genewise_dispersions"
 
-    def _fit_mean_dispersion_trend(self):
-        """Use the mean of dispersions as trend curve."""
         self.uns["mean_disp"] = trim_mean(
-            self.varm["genewise_dispersions"][
-                self.varm["genewise_dispersions"] > 10 * self.min_disp
-            ],
+            self.varm[disp_param_name][self.varm[disp_param_name] > 10 * self.min_disp],
             proportiontocut=0.001,
         )
 


Test patch:
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 7f44a0b..fa750ae 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -51,7 +51,7 @@ def test_deseq_independent_filtering_parametric_fit(counts_df, metadata, tol=0.0
         counts=counts_df,
         metadata=metadata,
         design_factors="condition",
-        trend_fit_type="parametric",
+        fit_type="parametric",
     )
     dds.deseq2()
 
@@ -79,7 +79,7 @@ def test_deseq_independent_filtering_mean_fit(counts_df, metadata, tol=0.02):
         counts=counts_df,
         metadata=metadata,
         design_factors="condition",
-        trend_fit_type="mean",
+        fit_type="mean",
     )
     dds.deseq2()
 
@@ -112,7 +112,7 @@ def test_deseq_without_independent_filtering_parametric_fit(
         counts=counts_df,
         metadata=metadata,
         design_factors="condition",
-        trend_fit_type="parametric",
+        fit_type="parametric",
     )
     dds.deseq2()
 
@@ -744,7 +744,7 @@ def test_vst_fit(train_dds):
     train_dds.vst_fit()
 
     # the correct attributes are fit
-    assert "trend_coeffs" in train_dds.uns
+    assert "vst_trend_coeffs" in train_dds.uns
     assert "normed_counts" in train_dds.layers
     assert "size_factors" in train_dds.obsm
 
@@ -759,6 +759,49 @@ def test_vst_transform(train_dds, test_counts):
     assert result.shape == (25, 10)
 
 
+@pytest.mark.parametrize(
+    ("dea_fit_type", "vst_fit_type"),
+    [
+        ("mean", "parametric"),
+        ("parametric", "mean"),
+        ("parametric", "parametric"),
+        ("mean", "mean"),
+    ],
+)
+def test_vst_blind(train_counts, train_metadata, dea_fit_type, vst_fit_type):
+    """Test vst with combinatory dea dea_fit_type and fit_type"""
+    train_dds = DeseqDataSet(
+        counts=train_counts,
+        metadata=train_metadata,
+        design_factors="condition",
+        fit_type=dea_fit_type,
+    )
+    train_dds.deseq2()
+    if dea_fit_type == "parametric":
+        assert "trend_coeffs" in train_dds.uns
+    else:
+        assert "mean_disp" in train_dds.uns
+    assert "normed_counts" in train_dds.layers
+    assert "size_factors" in train_dds.obsm
+    assert train_dds.fit_type == dea_fit_type
+
+    train_dds.vst(use_design=False, fit_type=vst_fit_type)
+    # Check that the dea fit type hasn't changed
+    assert train_dds.fit_type == dea_fit_type
+
+
+def test_vst_transform_no_fit(train_counts, train_metadata, test_counts):
+    """Test vst_transform without calling vst_fit()"""
+    train_dds = DeseqDataSet(
+        counts=train_counts,
+        metadata=train_metadata,
+        design_factors="condition",
+        fit_type="parametric",
+    )
+    with pytest.raises(RuntimeError):
+        train_dds.vst_transform(test_counts.to_numpy())
+
+
 def assert_res_almost_equal(py_res, r_res, tol=0.02):
     # check that the same p-values are NaN
     assert (py_res.pvalue.isna() == r_res.pvalue.isna()).all()