Repository: owkin/PyDESeq2
Pull request #101 (2023-03-20 16:39:32)

ENH multi-level design support
#### Reference Issue or PRs

Related to #100 

#### What does your PR implement? Be specific.

This PR extends multi-factor analysis from bi-level to n-level designs. Continuous variables are left to future work.


Patch:
diff --git a/datasets/README.md b/datasets/README.md
index f4add70..bbff6fb 100644
--- a/datasets/README.md
+++ b/datasets/README.md
@@ -3,6 +3,11 @@
 This directory stores example data which can be retrieved using the `load_example_data` function from
 `pydeseq2.utils`. Only synthetic data is provided for now, but new datasets might be available in the future.
 
+The `tests_clinical.csv` and `test_counts.csv` files were generated using a custom modification of the 
+`makeExampleDESeqDataSet` method of [DESeq2](https://bioconductor.org/packages/release/bioc/html/DESeq2.html), 
+to handle several factors. 
+
+
 ## Folder organisation
 
 ```
diff --git a/docs/source/api/docstrings/pydeseq2.dds.DeseqDataSet.rst b/docs/source/api/docstrings/pydeseq2.dds.DeseqDataSet.rst
index db1de66..0dfd0bb 100644
--- a/docs/source/api/docstrings/pydeseq2.dds.DeseqDataSet.rst
+++ b/docs/source/api/docstrings/pydeseq2.dds.DeseqDataSet.rst
@@ -18,4 +18,5 @@
         ~DeseqDataSet.fit_genewise_dispersions
         ~DeseqDataSet.fit_size_factors
         ~DeseqDataSet.refit
+        ~DeseqDataSet.vst
 
diff --git a/examples/plot_minimal_pydeseq2_pipeline.py b/examples/plot_minimal_pydeseq2_pipeline.py
index fe77db6..9dfe9e9 100644
--- a/examples/plot_minimal_pydeseq2_pipeline.py
+++ b/examples/plot_minimal_pydeseq2_pipeline.py
@@ -260,7 +260,7 @@ if SAVE:
 # LFC shrinkage. This is implemented by the :meth:`lfc_shrink() <DeseqStats.lfc_shrink>`
 # method.
 
-stat_res.lfc_shrink()
+stat_res.lfc_shrink(coeff="condition_B_vs_A")
 
 if SAVE:
     with open(os.path.join(OUTPUT_PATH, "shrunk_stat_results.pkl"), "wb") as f:
@@ -378,4 +378,4 @@ stat_res_Y_vs_X.summary()
 # only shrink the LFCs of a :class:`DeseqStats` object based on its
 # ``contrast`` argument.
 
-stat_res_B_vs_A.lfc_shrink()
+stat_res_B_vs_A.lfc_shrink(coeff="condition_B_vs_A")
diff --git a/examples/plot_step_by_step.py b/examples/plot_step_by_step.py
index 6a5e7ef..e673fc5 100644
--- a/examples/plot_step_by_step.py
+++ b/examples/plot_step_by_step.py
@@ -238,7 +238,7 @@ if SAVE:
 # For visualization or post-processing purposes, it might be suitable to perform
 # LFC shrinkage. This is implemented by the `lfc_shrink` method.
 
-stat_res.lfc_shrink()
+stat_res.lfc_shrink(coeff="condition_B_vs_A")
 
 
 # %%
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index d9930c0..bc91613 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -35,7 +35,8 @@ warnings.simplefilter("ignore", DomainWarning)
 # Ignore AnnData's FutureWarning about implicit data conversion.
 warnings.simplefilter("ignore", FutureWarning)
 
-# TODO: support loading / saving DeseqDataSets from AnnData objects.
+
+# TODO: implement a quiet (non-verbose) mode ?
 
 
 class DeseqDataSet(ad.AnnData):
@@ -63,14 +64,13 @@ class DeseqDataSet(ad.AnnData):
         Must be indexed by sample barcodes.
 
     design_factors : str or list
-        Name of the columns of clinical to be used as design variables. If a list,
-        the last factor will be considered the variable of interest by default.
+        Name of the columns of clinical to be used as design variables.
         Only bi-level factors are supported. (default: ``'condition'``).
 
-    reference_level : str
-        The factor to use as a reference. Must be one of the values taken by the design.
-        If None, the reference will be chosen alphabetically (last in order).
-        (default: ``None``).
+    tested_level : list or None
+        An optional list of two strings of the form ``["factor", "test_level"]``
+        specifying the factor of interest and the (non-control) level we're testing, e.g.
+        ``["condition", "B"]``. (default: ``None``).
 
     min_mu : float
         Threshold for mean estimates. (default: ``0.5``).
@@ -161,7 +161,7 @@ class DeseqDataSet(ad.AnnData):
         counts: Optional[pd.DataFrame] = None,
         clinical: Optional[pd.DataFrame] = None,
         design_factors: Union[str, List[str]] = "condition",
-        reference_level: Optional[str] = None,
+        tested_level: Optional[List[str]] = None,
         min_mu: float = 0.5,
         min_disp: float = 1e-8,
         max_disp: float = 10.0,
@@ -202,7 +202,7 @@ class DeseqDataSet(ad.AnnData):
         self.obsm["design_matrix"] = build_design_matrix(
             clinical_df=self.obs,
             design_factors=self.design_factors,
-            ref=reference_level,
+            tested_level=tested_level,
             expanded=False,
             intercept=True,
         )
@@ -217,6 +217,50 @@ class DeseqDataSet(ad.AnnData):
         self.batch_size = batch_size
         self.joblib_verbosity = joblib_verbosity
 
+    def vst(self, use_design: bool = False) -> None:
+        """
+        Fit a variance stabilizing transformation, and apply it to normalized counts.
+
+        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``).
+        """
+
+        # Start by fitting median-of-ratio size factors, if not already present.
+        if "size_factors" not in self.obsm:
+            self.fit_size_factors()
+
+        if use_design:
+            # Check that the dispersion trend curve was fitted. If not, fit it.
+            # This will call previous functions in a cascade.
+            if "trend_coeffs" not in self.uns:
+                self.fit_dispersion_trend()
+        else:
+            # Reduce the design matrix to an intercept and reconstruct at the end
+            self.obsm["design_matrix_buffer"] = self.obsm["design_matrix"].copy()
+            self.obsm["design_matrix"] = pd.DataFrame(
+                1, index=self.obs_names, columns=[["intercept"]]
+            )
+            # Fit the trend curve with an intercept design
+            self.fit_genewise_dispersions()
+            self.fit_dispersion_trend()
+
+            # Restore the design matrix and free buffer
+            self.obsm["design_matrix"] = self.obsm["design_matrix_buffer"].copy()
+            del self.obsm["design_matrix_buffer"]
+
+        # Apply VST
+        a0, a1 = self.uns["trend_coeffs"]
+        cts = self.layers["normed_counts"]
+        self.layers["vst_counts"] = np.log2(
+            (1 + a1 + 2 * a0 * cts + 2 * np.sqrt(a0 * cts * (1 + a1 + a0 * cts)))
+            / (4 * a0)
+        )
+
     def deseq2(self) -> None:
         """Perform dispersion and log fold-change (LFC) estimation.
 
diff --git a/pydeseq2/ds.py b/pydeseq2/ds.py
index 96ee426..daef8d2 100644
--- a/pydeseq2/ds.py
+++ b/pydeseq2/ds.py
@@ -1,7 +1,6 @@
 import time
 from typing import List
 from typing import Optional
-from typing import cast
 
 # import anndata as ad
 import numpy as np
@@ -36,7 +35,7 @@ class DeseqStats:
 
     contrast : list or None
         A list of three strings, in the following format:
-        ``['variable_of_interest', 'tested_level', 'reference_level']``.
+        ``['variable_of_interest', 'tested_level', 'tested_level']``.
         Names must correspond to the clinical data passed to the DeseqDataSet.
         E.g., ``['condition', 'B', 'A']`` will measure the LFC of 'condition B' compared
         to 'condition A'. If None, the last variable from the design matrix is chosen
@@ -73,6 +72,9 @@ class DeseqStats:
     base_mean : pandas.Series
         Genewise means of normalized counts.
 
+    contrast_vector : ndarray
+        Vector encoding the contrast (variable being tested).
+
     contrast_idx : int
         Index of the LFC column corresponding to the variable being tested.
 
@@ -130,64 +132,23 @@ class DeseqStats:
 
         self.dds = dds
 
-        if contrast is not None:  # Test contrast if provided
-            assert len(contrast) == 3, "The contrast should contain three strings."
-            assert (
-                contrast[0] in self.dds.design_factors
-            ), "The contrast variable should be one of the design factors."
-            assert (
-                contrast[1] in self.dds.obs[contrast[0]].values
-                and contrast[2] in self.dds.obs[contrast[0]].values
-            ), "The contrast levels should correspond to design factors levels."
-            self.contrast = contrast
-        else:  # Build contrast if None
-            factor = self.dds.design_factors[-1]
-            levels = np.unique(self.dds.obs[factor]).astype(str)
-            if (
-                "_".join([factor, levels[0]])
-                == self.dds.obsm["design_matrix"].columns[-1]
-            ):
-                self.contrast = [factor, levels[0], levels[1]]
-            else:
-                self.contrast = [factor, levels[1], levels[0]]
-
         self.alpha = alpha
         self.cooks_filter = cooks_filter
         self.independent_filter = independent_filter
-        self.prior_LFC_var = prior_LFC_var
         self.base_mean = self.dds.varm["_normed_means"].copy()
+        self.prior_LFC_var = prior_LFC_var
+
+        # Check the validity of the contrast (if provided) or build it.
+        self._build_contrast(contrast)
 
         # Initialize the design matrix and LFCs. If the chosen reference level are the
         # same as in dds, keep them unchanged. Otherwise, change reference level.
         self.design_matrix = self.dds.obsm["design_matrix"].copy()
         self.LFC = self.dds.varm["LFC"].copy()
-        if self.contrast is None:
-            alternative_level = self.dds.obsm["design_matrix"].columns[-1]
-        else:
-            alternative_level = "_".join([self.contrast[0], self.contrast[1]])
-
-        if alternative_level in self.dds.obsm["design_matrix"].columns:
-            self.contrast_idx = self.dds.varm["LFC"].columns.get_loc(alternative_level)
-        else:  # The reference level is not the same as in dds: change it
-            reference_level = "_".join([self.contrast[0], self.contrast[2]])
-            self.contrast_idx = self.dds.varm["LFC"].columns.get_loc(reference_level)
-            # Change the design matrix reference level
-            self.design_matrix.rename(
-                columns={
-                    self.design_matrix.columns[self.contrast_idx]: alternative_level
-                },
-                inplace=True,
-            )
-            self.design_matrix.iloc[:, self.contrast_idx] = (
-                1 - self.design_matrix.iloc[:, self.contrast_idx]
-            )
-            # Rename and update LFC coefficients accordingly
-            self.LFC.rename(
-                columns={self.LFC.columns[self.contrast_idx]: alternative_level},
-                inplace=True,
-            )
-            self.LFC.iloc[:, 0] += self.LFC.iloc[:, self.contrast_idx]
-            self.LFC.iloc[:, self.contrast_idx] *= -1
+
+        # Build a contrast vector corresponding to the variable and levels of interest
+        self._build_contrast_vector()
+
         # Set a flag to indicate that LFCs are unshrunk
         self.shrunk_LFCs = False
         self.n_processes = get_num_processes(n_cpus)
@@ -229,9 +190,7 @@ class DeseqStats:
         # Store the results in a DataFrame, in log2 scale for LFCs.
         self.results_df = pd.DataFrame(index=self.dds.var_names)
         self.results_df["baseMean"] = self.base_mean
-        self.results_df["log2FoldChange"] = self.LFC.iloc[:, self.contrast_idx] / np.log(
-            2
-        )
+        self.results_df["log2FoldChange"] = self.LFC @ self.contrast_vector / np.log(2)
         self.results_df["lfcSE"] = self.SE / np.log(2)
         self.results_df["stat"] = self.statistics
         self.results_df["pvalue"] = self.p_values
@@ -289,7 +248,7 @@ class DeseqStats:
                     lfc=LFCs[i],
                     mu=mu[:, i],
                     ridge_factor=ridge_factor,
-                    idx=self.contrast_idx,
+                    contrast=self.contrast_vector,
                 )
                 for i in range(num_genes)
             )
@@ -298,11 +257,9 @@ class DeseqStats:
 
         pvals, stats, se = zip(*res)
 
-        # need casting because mypy recognize them as TimeStampSeries
-        # reason is unknown...
-        self.p_values = cast(pd.Series, pd.Series(pvals, index=self.dds.var_names))
-        self.statistics = cast(pd.Series, pd.Series(stats, index=self.dds.var_names))
-        self.SE = cast(pd.Series, pd.Series(se, index=self.dds.var_names))
+        self.p_values: pd.Series = pd.Series(pvals, index=self.dds.var_names)
+        self.statistics: pd.Series = pd.Series(stats, index=self.dds.var_names)
+        self.SE: pd.Series = pd.Series(se, index=self.dds.var_names)
 
         # Account for possible all_zeroes due to outlier refitting in DESeqDataSet
         if self.dds.refit_cooks and self.dds.varm["replaced"].sum() > 0:
@@ -310,23 +267,34 @@ class DeseqStats:
             self.statistics.loc[self.dds.new_all_zeroes_genes] = 0.0
             self.p_values.loc[self.dds.new_all_zeroes_genes] = 1.0
 
-    def lfc_shrink(self) -> None:
+    def lfc_shrink(self, coeff: str) -> None:
         """LFC shrinkage with an apeGLM prior :cite:p:`DeseqStats-zhu2019heavy`.
 
         Shrinks LFCs using a heavy-tailed Cauchy prior, leaving p-values unchanged.
 
-        Returns
-        -------
-        pandas.DataFrame or None
-            If pvalues were already computed, return the results DataFrame with MAP LFCs,
-            but unmodified stats and pvalues.
+        Parameters
+        ----------
+        coeff : str
+            The LFC coefficient to shrink.
+            If the desired coefficient is not available, it may be set from the
+            :class:`pydeseq2.dds.DeseqDataSet` argument ``tested_level``.
         """
+
+        if coeff not in self.LFC.columns:
+            raise KeyError(
+                f"The coeff argument ('{coeff}') should be one the LFC columns. "
+                f"If not available,  it can be set from DeseqDataSet's `tested_level` "
+                f"argument."
+            )
+
+        coeff_idx = self.LFC.columns.get_loc(coeff)
+
         size = 1.0 / self.dds.varm["dispersions"]
         offset = np.log(self.dds.obsm["size_factors"])
 
         # Set priors
         prior_no_shrink_scale = 15
-        prior_var = self._fit_prior_var()
+        prior_var = self._fit_prior_var(coeff_idx=coeff_idx)
         prior_scale = np.minimum(np.sqrt(prior_var), 1)
 
         design_matrix = self.design_matrix.values
@@ -347,7 +315,7 @@ class DeseqStats:
                     prior_no_shrink_scale=prior_no_shrink_scale,
                     prior_scale=prior_scale,
                     optimizer="L-BFGS-B",
-                    shrink_index=self.contrast_idx,
+                    shrink_index=coeff_idx,
                 )
                 for i in self.dds.non_zero_idx
             )
@@ -356,9 +324,9 @@ class DeseqStats:
 
         lfcs, inv_hessians, l_bfgs_b_converged_ = zip(*res)
 
-        self.LFC.iloc[:, self.contrast_idx].update(
+        self.LFC.iloc[:, coeff_idx].update(
             pd.Series(
-                np.array(lfcs)[:, self.contrast_idx],
+                np.array(lfcs)[:, coeff_idx],
                 index=self.dds.non_zero_genes,
             )
         )
@@ -367,7 +335,7 @@ class DeseqStats:
             pd.Series(
                 np.array(
                     [
-                        np.sqrt(np.abs(inv_hess[self.contrast_idx, self.contrast_idx]))
+                        np.sqrt(np.abs(inv_hess[coeff_idx, coeff_idx]))
                         for inv_hess in inv_hessians
                     ]
                 ),
@@ -385,9 +353,7 @@ class DeseqStats:
 
         # Replace in results dataframe, if it exists
         if hasattr(self, "results_df"):
-            self.results_df["log2FoldChange"] = self.LFC.iloc[
-                :, self.contrast_idx
-            ] / np.log(2)
+            self.results_df["log2FoldChange"] = self.LFC.iloc[:, coeff_idx] / np.log(2)
             self.results_df["lfcSE"] = self.SE / np.log(2)
 
             print(
@@ -448,7 +414,6 @@ class DeseqStats:
     def _p_value_adjustment(self) -> None:
         """Compute adjusted p-values using the Benjamini-Hochberg method.
 
-
         Does not correct the p-value trend.
         This method and the `_independent_filtering` are mutually exclusive.
         """
@@ -511,13 +476,18 @@ class DeseqStats:
 
         self.p_values[cooks_outlier] = np.nan
 
-    def _fit_prior_var(self, min_var: float = 1e-6, max_var: float = 400.0) -> float:
+    def _fit_prior_var(
+        self, coeff_idx: str, min_var: float = 1e-6, max_var: float = 400.0
+    ) -> float:
         """Estimate the prior variance of the apeGLM model.
 
         Returns shrinkage factors.
 
         Parameters
         ----------
+        coeff_idx : str
+            Index of the coefficient to shrink.
+
         min_var : float
             Lower bound for prior variance. (default: ``1e-6``).
 
@@ -530,8 +500,8 @@ class DeseqStats:
             Estimated prior variance.
         """
 
-        keep = ~self.LFC.iloc[:, self.contrast_idx].isna()
-        S = self.LFC[keep].iloc[:, self.contrast_idx] ** 2
+        keep = ~self.LFC.iloc[:, coeff_idx].isna()
+        S = self.LFC[keep].iloc[:, coeff_idx] ** 2
         D = self.SE[keep] ** 2
 
         def objective(a: float) -> float:
@@ -544,3 +514,87 @@ class DeseqStats:
             return min_var
         else:
             return root_scalar(objective, bracket=[min_var, max_var]).root
+
+    def _build_contrast(self, contrast: Optional[List[str]] = None) -> None:
+        """
+        Check the validity of the contrast (if provided). If not, build a default
+        contrast, corresponding to the last column of the design matrix.
+
+        A contrast should be a list of three strings, in the following format:
+        ``['variable_of_interest', 'tested_level', 'reference_level']``.
+        Names must correspond to the clinical data passed to the DeseqDataSet.
+        E.g., ``['condition', 'B', 'A']`` will measure the LFC of 'condition B'
+        compared to 'condition A'. If None, the last variable from the design matrix
+        is chosen as the variable of interest, and the reference level is picked
+        alphabetically.
+
+        Parameters
+        ----------
+        contrast : list or None
+            A list of three strings, in the following format:
+            ``['variable_of_interest', 'tested_level', 'reference_level']``.
+            (default: ``None``).
+        """
+
+        if contrast is not None:  # Test contrast if provided
+            if len(contrast) != 3:
+                raise ValueError("The contrast should contain three strings.")
+            if contrast[0] not in self.dds.design_factors:
+                raise KeyError(
+                    f"The contrast variable ('{contrast[0]}') should be one "
+                    f"of the design factors."
+                )
+            if contrast[1] not in self.dds.obs[contrast[0]].values:
+                raise KeyError(
+                    f"The tested level ('{contrast[1]}') should correspond to one of the"
+                    f" levels of '{contrast[0]}'"
+                )
+            if contrast[2] not in self.dds.obs[contrast[0]].values:
+                raise KeyError(
+                    f"The reference level ('{contrast[2]}') should correspond to one of"
+                    f" the levels of '{contrast[0]}'"
+                )
+            self.contrast = contrast
+        else:  # Build contrast if None
+            factor = self.dds.design_factors[-1]
+            factor_col = next(
+                col
+                for col in self.dds.obsm["design_matrix"].columns
+                if col.startswith(factor)
+            )
+            split_col = factor_col.split("_")
+            self.contrast = [split_col[0], split_col[1], split_col[-1]]
+
+    def _build_contrast_vector(self) -> None:
+        """
+        Build a vector corresponding to the desired contrast.
+
+        Allows to test any pair of levels without refitting LFCs.
+        """
+        factor = self.contrast[0]
+        alternative = self.contrast[1]
+        ref = self.contrast[2]
+        contrast_level = f"{factor}_{alternative}_vs_{ref}"
+
+        self.contrast_vector = np.zeros(self.LFC.shape[-1])
+        if contrast_level in self.design_matrix.columns:
+            self.contrast_idx = self.LFC.columns.get_loc(contrast_level)
+            self.contrast_vector[self.contrast_idx] = 1
+        elif f"{factor}_{ref}_vs_{alternative}" in self.design_matrix.columns:
+            # Reference and alternative are inverted
+            self.contrast_idx = self.LFC.columns.get_loc(
+                f"{factor}_{ref}_vs_{alternative}"
+            )
+            self.contrast_vector[self.contrast_idx] = -1
+        else:
+            # Need to change reference
+            # Get any column corresponding to the desired factor and extract old ref
+            old_ref = next(
+                col for col in self.LFC.columns if col.startswith(factor)
+            ).split("_")[-1]
+            new_alternative_idx = self.LFC.columns.get_loc(
+                f"{factor}_{alternative}_vs_{old_ref}"
+            )
+            new_ref_idx = self.LFC.columns.get_loc(f"{factor}_{ref}_vs_{old_ref}")
+            self.contrast_vector[new_alternative_idx] = 1
+            self.contrast_vector[new_ref_idx] = -1
diff --git a/pydeseq2/utils.py b/pydeseq2/utils.py
index 9e1ebb2..d774f6c 100644
--- a/pydeseq2/utils.py
+++ b/pydeseq2/utils.py
@@ -139,7 +139,7 @@ def test_valid_counts(counts: Union[pd.DataFrame, np.ndarray]) -> None:
 def build_design_matrix(
     clinical_df: pd.DataFrame,
     design_factors: Union[str, List[str]] = "condition",
-    ref: Optional[str] = None,
+    tested_level: Optional[List[str]] = None,
     expanded: bool = False,
     intercept: bool = True,
 ) -> pd.DataFrame:
@@ -159,9 +159,10 @@ def build_design_matrix(
         Name of the columns of clinical_df to be used as design_matrix variables.
         (default: ``"condition"``).
 
-    ref : str
-        The factor to use as a reference. Must be one of the values taken by the design.
-        If None, the reference will be chosen alphabetically (last in order).
+    tested_level : list or None
+        An optional list of two strings of the form ``["factor", "ref_level"]``
+        specifying the factor of interest and the desired reference level, e.g.
+        ``["condition", "A"]``.
         (default: ``None``).
 
     expanded : bool
@@ -183,33 +184,53 @@ def build_design_matrix(
     ):  # if there is a single factor, convert to singleton list
         design_factors = [design_factors]
 
-    design_matrix = pd.get_dummies(clinical_df[design_factors])
-
     for factor in design_factors:
-        # Check that each factor has exactly 2 levels
-        if len(np.unique(clinical_df[factor])) != 2:
+        # Check that each factor has at least 2 levels
+        if len(np.unique(clinical_df[factor])) < 2:
             raise ValueError(
-                f"Factors should take exactly two values, but {factor} "
-                f"takes values {np.unique(clinical_df[factor])}."
+                f"Factors should take at least two values, but {factor} "
+                f"takes the single value '{np.unique(clinical_df[factor])}'."
             )
-    factor = design_factors[-1]
-    if ref is None:
-        ref = "_".join([factor, np.sort(np.unique(clinical_df[factor]).astype(str))[0]])
-        ref_level = design_matrix.pop(ref)
-    else:  # Put reference level last
-        try:
-            ref_level = design_matrix.pop(f"{factor}_{ref}")
-        except KeyError as e:
-            print(
-                "The reference level must correspond to"
-                " one of the design factor values."
+
+    design_matrix = pd.get_dummies(clinical_df[design_factors], drop_first=not expanded)
+
+    if tested_level is not None:
+        if len(tested_level) != 2:
+            raise KeyError("The tested level should contain 2 strings.")
+        if tested_level[1] not in clinical_df[tested_level[0]].values:
+            raise KeyError(
+                f"The clinical data should contain a '{tested_level[0]}' column"
+                f" with a '{tested_level[1]}' level."
             )
-            raise e
-    design_matrix.insert(design_matrix.shape[-1], ref, ref_level)
-    if not expanded:  # drop duplicate factors
-        design_matrix.drop(
-            columns=[col for col in design_matrix.columns[::-2]], axis=1, inplace=True
-        )  # Drops every other column, starting from the last
+
+        # Check that the reference level is still in the matrix
+        test_level = "_".join(tested_level)
+        if test_level not in design_matrix.columns:
+            # Add the desired level and remove one
+            factor_cols = [
+                col for col in design_matrix.columns if col.startswith(tested_level[0])
+            ]
+            design_matrix[test_level] = 1 - design_matrix[factor_cols].sum(1)
+            design_matrix.drop(factor_cols[0], axis="columns", inplace=True)
+
+    if not expanded:
+        # Add reference level as column name suffix
+        for factor in design_factors:
+            if tested_level is not None and factor == tested_level[0]:
+                # The reference is the unique level that is no longer there
+                ref = next(
+                    level
+                    for level in clinical_df[factor]
+                    if f"{factor}_{level}" not in design_matrix.columns
+                )
+            else:
+                # The first level is the one that should have been dropped
+                ref = np.unique(clinical_df[factor])[0]
+            design_matrix.columns = [
+                f"{col}_vs_{ref}" if col.startswith(factor) else col
+                for col in design_matrix.columns
+            ]
+
     if intercept:
         design_matrix.insert(0, "intercept", 1)
     return design_matrix
@@ -782,7 +803,7 @@ def wald_test(
     lfc: np.ndarray,
     mu: float,
     ridge_factor: np.ndarray,
-    idx: int = -1,
+    contrast: np.ndarray,
 ) -> Tuple[float, float, float]:
     """Run Wald test for differential expression.
 
@@ -806,8 +827,8 @@ def wald_test(
     ridge_factor : ndarray
         Regularization factors.
 
-    idx : int
-        Index of design factor (in design matrix). (default: ``-1``).
+    contrast : ndarray
+        Vector encoding the contrast that is being tested.
 
     Returns
     -------
@@ -825,10 +846,10 @@ def wald_test(
     W = np.diag(mu / (1 + mu * disp))
     M = design_matrix.T @ W @ design_matrix
     H = np.linalg.inv(M + ridge_factor)
-    S = H @ M @ H
+    Hc = H @ contrast
     # Evaluate standard error and Wald statistic
-    wald_se: float = np.sqrt(S[idx, idx])
-    wald_statistic: float = lfc[idx] / wald_se
+    wald_se: float = np.sqrt(Hc.T @ M @ Hc)
+    wald_statistic: float = contrast @ lfc / wald_se
     wald_p_value = 2 * norm.sf(np.abs(wald_statistic))
     return wald_p_value, wald_statistic, wald_se
 


Test patch:
diff --git a/tests/data/README.md b/tests/data/README.md
index 683ce6e..0df3165 100644
--- a/tests/data/README.md
+++ b/tests/data/README.md
@@ -2,9 +2,13 @@
 
 This folder contains data for the pytest CI.
 
-The `tests_clinical.csv` and `test_counts.csv` files were generated using the `makeExampleDESeqDataSet` 
-of [DESeq2](https://bioconductor.org/packages/release/bioc/html/DESeq2.html) (with parameters `m = 100`, `n = 10`,
-and `betaSD = 1`).
+The files in `single_factor` and `multi_factor` contain the outputs of DESeq2 (v1.34.0) on the synthetic data provided
+in `/datasets/synthetic/`, respectively using `~condition` and `~condition + group` as design. More precisely:
 
-The remainder of the files (`r_test_*`) correspond to the output of DESeq2 (v1.34.0) with default settings on this dataset, 
-against which the output of PyDESeq2 is tested in the `test_pydeseq2.py` pytest when opening PRs.
\ No newline at end of file
+- `r_iterative_size_factors.csv` contains DESeq2's `estimateSizeFactorsIterate` output,
+- `r_lfc_shrink.csv` contains DESeq2 results after running `lfcShrink`,
+- `r_test_dispersions.csv` contains DESeq2 dispersions estimates (post-filtering and refitting),
+- `r_test_res.csv` contains DESeq2's `results` output,
+- `r_test_size_factors.csv` contains DESeq2's `estimateSizeFactors` output,
+- `r_vst.csv` contains DESeq2's `varianceStabilizingTransformation` output with `blind=TRUE` and `fitType="parametric"`,
+- `r_vst_with_design.csv` contains DESeq2's `varianceStabilizingTransformation` output with `blind=FALSE` and `fitType="parametric"`.
\ No newline at end of file
diff --git a/tests/data/single_factor/r_vst.csv b/tests/data/single_factor/r_vst.csv
new file mode 100644
index 0000000..0ae9c52
--- /dev/null
+++ b/tests/data/single_factor/r_vst.csv
@@ -0,0 +1,11 @@
+"","sample1","sample2","sample3","sample4","sample5","sample6","sample7","sample8","sample9","sample10","sample11","sample12","sample13","sample14","sample15","sample16","sample17","sample18","sample19","sample20","sample21","sample22","sample23","sample24","sample25","sample26","sample27","sample28","sample29","sample30","sample31","sample32","sample33","sample34","sample35","sample36","sample37","sample38","sample39","sample40","sample41","sample42","sample43","sample44","sample45","sample46","sample47","sample48","sample49","sample50","sample51","sample52","sample53","sample54","sample55","sample56","sample57","sample58","sample59","sample60","sample61","sample62","sample63","sample64","sample65","sample66","sample67","sample68","sample69","sample70","sample71","sample72","sample73","sample74","sample75","sample76","sample77","sample78","sample79","sample80","sample81","sample82","sample83","sample84","sample85","sample86","sample87","sample88","sample89","sample90","sample91","sample92","sample93","sample94","sample95","sample96","sample97","sample98","sample99","sample100"
+"gene1",5.04110709395604,4.24731263665037,4.6444922568408,4.27753997399955,4.31287502255098,3.90639993053475,3.90639993053475,4.39465863453275,4.89593407190211,5.24540858603705,4.53275039705426,4.71187486005075,4.50492580120502,5.76641221069808,4.85958051113913,6.09938739294276,4.22892932776322,5.84737010143047,4.55986077167108,4.23276158187784,4.46470073658036,4.97331145680698,4.40470224343279,5.31320854899517,4.24692118522698,5.58145635185857,3.90639993053475,5.00393489036174,4.29922835160006,4.23610380208107,4.81931561066275,3.90639993053475,4.96083900835411,4.65001762507949,4.47472329605445,4.41764779761476,4.52936856112822,3.90639993053475,4.42972950499431,5.17535916774395,4.32078601433704,5.04668467669856,4.46332205815783,5.54178670093241,4.25949137111742,3.90639993053475,5.11011238856864,4.41966482442235,5.42899706993301,4.58221974118338,5.18551140573441,4.87577561708095,5.29757294204838,5.12202550713102,4.65347312783471,4.5006965891644,4.80301105931374,5.67074257060676,5.22315462025469,4.52915573807794,4.72978446553888,4.29342764878885,4.91974057203888,5.41131530746111,4.76601832528728,5.15690559067124,5.37898522546302,4.87498559386838,4.415248584258,5.71190211352045,5.73920305412469,4.91583690310765,4.74843383607094,4.90467760914364,4.96666732989512,4.79450219731,4.60859993406858,5.04434552712468,4.97529186778471,5.48946010480508,4.52692538068361,4.45058215927641,4.90803492647296,5.07623446827891,5.7331944979534,4.75142424348637,5.20056444429186,4.97818480971189,5.16577392253431,4.81506951199612,3.90639993053475,4.95599266608811,4.83245842814059,5.72962264698006,4.39470401883235,4.76842167661968,4.33673625137178,5.22925893848332,5.44643110925359,5.88819548148005
+"gene2",5.38138373999374,5.98776596679974,4.6444922568408,5.05288025905266,5.21379131165658,4.88872900081318,5.31881430714792,5.35666696174447,5.74617156958584,4.96707802418041,5.42412734684098,5.46190688776824,4.92898655804911,5.29114573480476,6.07414828108701,5.33050738647249,5.52963457425083,5.2203939736305,5.15208824904277,5.24540858603705,4.86189153671866,5.49244122418272,5.51931620520483,5.45695610206686,5.29968522613836,5.51683196428453,5.57256044072795,5.60864452182457,5.3690686318354,5.8466154111864,5.46395288330041,5.62309643661323,5.26693940688655,4.94700898729297,4.76732128592982,5.17535916774395,5.0028957046876,5.15852946755355,5.44065085297831,5.62074738456303,5.73238714582962,5.42940032693502,5.58261247486227,5.21753230710055,4.68830349516006,5.39546472196244,5.60241880072018,5.22515015420892,5.82095453428012,5.55417113541989,5.8119638745356,5.54882325017073,6.06125519988914,5.5812207345249,5.54743745030931,6.16149178805503,5.66721140474223,5.75514570254157,5.49427961971276,5.76074295358258,5.29746227766118,5.78739975055133,5.35279685591621,4.92619150213677,5.03171944035226,5.4727831050269,5.49950344143982,5.17103867729767,5.56124621649659,5.10990911580924,4.86779405708399,5.55628820871717,5.96085962340469,6.17906210498546,5.76676694709224,5.30564808158969,5.02187889748249,4.86558254179463,6.01809470800165,5.77043076357162,5.46075756767801,5.72000259606073,5.41195940182634,5.52091563248034,4.83671938940712,6.93193911212323,6.09284030050012,6.05107687523118,6.40141353698257,5.22968537878753,5.86108512095158,6.04910584119304,5.36276275990613,5.16418245618518,5.52050807084074,5.50889265841308,5.4481314992325,6.24174949853372,5.27777236645892,5.25822936477786
+"gene3",4.57260559562058,4.38741470902142,5.10868602571642,4.4298439020043,4.88640224621299,4.40462818913784,5.45089063459605,5.02844654179347,5.214866617756,4.81630357905239,5.17886783527573,4.71187486005075,4.59590407754956,5.116413217557,5.23193392741788,4.93298499311661,4.95509680326643,4.60959480819522,3.90639993053475,4.63011174083191,4.23007091044943,4.35020624111892,4.6885917511608,4.26039228498293,4.49351035306986,4.26691225529936,5.50218714630825,5.00393489036174,4.99491824061484,3.90639993053475,4.32009514949799,3.90639993053475,4.80303148899908,4.55214048012604,4.82488996061365,4.5309414424061,4.52936856112822,5.20289566142941,4.20964513027446,5.17535916774395,5.23774663083511,5.04668467669856,4.54817390061422,4.88319672947171,4.76090838578896,3.90639993053475,5.16546469094941,4.32620988104311,4.89522365440025,4.38644050262622,4.62278879216389,3.90639993053475,6.00286509833062,4.32226125565993,4.28305340764572,3.90639993053475,4.42955867454486,3.90639993053475,4.41884355652748,4.62368178746474,4.27861350778072,4.37970734249569,4.71351910608834,4.58181142860657,4.76601832528728,4.62023985927691,4.56632339400683,4.25446965721268,4.415248584258,4.26307593190428,4.65647188680464,5.26005883494912,4.81383239213712,4.22775396903159,4.56466421747971,4.30856878165688,4.2229072215984,4.63706805716222,3.90639993053475,4.23840373908919,4.63841068386048,4.759455919178,4.70406308983809,4.26791456796197,4.37756441670915,4.50812944854385,4.97431946162618,4.60733920531614,4.61126642847382,4.31812107307723,5.12532043747068,4.61453241246802,5.30980833686914,3.90639993053475,4.25250301548186,4.4754616300581,4.64639206486436,4.51443428708131,4.55979505184117,4.70549249525506
+"gene4",7.06707589709277,6.32519488461959,7.71023820472443,7.00020423257605,7.57416533578186,5.8745194991298,7.68357273957781,6.30684703210763,6.96597797293252,7.2237499223187,7.38727629753508,5.93706039823475,7.13019861818193,6.44359363414318,7.66594359007695,7.09741997016699,7.49795595770808,7.59654366652137,7.61431545308689,6.28479706544742,6.79884719356467,6.84903562316212,7.07596501764677,6.52300555494802,6.47121550538279,7.7691456144451,7.55541237039493,7.52001834161209,7.28606425441911,7.11358164568079,7.28350723364332,6.46022351569673,6.84999833674695,6.78694696308595,8.14859815559456,7.05478009554724,7.65481276101848,7.05308627929658,6.96896437182704,7.05478009554724,7.58493277860203,6.98722177993478,7.34735512823581,7.59288597777907,7.31713189661219,6.39241089111043,6.97167762597171,7.15189277064297,6.71078874372737,7.16182460272112,7.29343703540422,6.36470042726026,7.00702607434902,6.34280925091369,7.19636418680319,5.96631636372838,7.15011309535814,5.9615034890017,6.82116871298937,6.82187202090134,6.64392763905999,5.25625654308582,7.30064159530108,6.99457270339434,6.95571546318873,6.55910883787392,6.8761210054242,6.47336381308007,6.72478368338905,6.76808244359844,7.33312440727815,5.82269605395026,7.04996912143411,6.69021793084148,6.92251489482268,5.85206489301835,7.09071395240256,6.21150675373108,6.97539934399698,6.24171391823007,7.15773539095839,7.01929460040401,7.02034434491261,7.19740833170314,7.27047651394953,7.03887769863569,7.16049676109225,6.53194502267638,7.21886377210665,7.08000560738104,7.24180472569671,6.74328071320216,7.42537566721996,6.88396851803583,6.50485640865049,6.32121302654305,6.94954684544664,6.97167762597171,6.59230519329161,6.54931685675434
+"gene5",5.27977775349027,5.01242678829521,5.52917416190478,6.14395377142904,5.46132267016367,5.61085552688703,5.16835414309355,5.35666696174447,5.06863432239993,5.71958753399194,4.62780146866586,5.71220824019663,5.84538767082605,5.57384436069832,5.85227382867697,5.69081326112896,5.23064696170947,4.80846731072459,5.85883077686226,5.94954286594267,4.91162718023219,5.38437497432934,5.77594610109243,5.4902008728484,6.02556910586182,4.62194527447819,5.26040016559635,6.69004008592371,5.32315603006249,5.25825776961192,5.01570670769428,5.7228635474245,5.46010281383784,5.73495634024146,5.68811927974949,6.23907273359131,5.48718185665909,5.98725038378144,5.44065085297831,5.3444631477771,5.39527000791195,5.37439418207758,5.55717339771854,5.43092149015496,4.99943179060371,5.50420210037335,5.60241880072018,5.71601863160511,5.16755915628061,4.96111642042481,5.40204327174255,5.5672452380761,5.58134994350081,5.49424516006405,5.98754618245661,6.14323259520034,5.21887308147853,6.1240170401695,6.06114233593862,6.09419160105966,5.80969696885977,5.78739975055133,5.88091039433995,5.899478633058,5.98769878809005,6.44893114653051,5.46073239205348,6.06561165616487,5.98909322571119,6.02753721226075,5.30365625949733,5.51841194563855,4.93072525155525,5.19123435693069,5.28744286789935,6.23034879451679,5.02187889748249,6.19141585667611,5.30856804066297,5.979980112896,5.61680431229272,6.34326942899687,5.9331890169793,5.55373990042258,6.49879240932551,5.48615881385995,6.63487958727271,5.57446344469266,5.37944959817645,6.07081913839363,6.29926691235672,6.00607877457362,5.71601863160511,6.84577769053359,5.71688228219104,4.97832091245982,6.15352521369212,5.28937987810541,6.36699490607727,5.5077145306243
+"gene6",3.90639993053475,4.96318506457132,4.54731715530875,4.99675646607325,4.47937263671711,4.5151139606457,5.53109198553411,4.39465863453275,4.67874694503383,3.90639993053475,4.91651835286154,3.90639993053475,4.39625000777316,4.98013181623019,4.85958051113913,3.90639993053475,4.95509680326643,4.31505298622751,5.12071622125526,4.36697144065255,5.00336522429312,5.32616120282182,4.8888710122369,3.90639993053475,5.145075058655,4.41493215529955,4.73419351646966,5.69071480493294,3.90639993053475,4.37166805744251,4.97976079221372,4.32798932013999,4.60530528071583,4.43509374769838,4.56125830532635,4.62572834783573,4.39043368111895,5.11212429580896,4.42972950499431,4.41764779761476,4.61933123048158,4.34697934429408,4.46332205815783,5.21753230710055,4.40451598041393,3.90639993053475,5.31562716684857,4.83245842814059,4.89522365440025,4.24661903077946,5.05753390446044,4.29265866955699,4.98525880264372,3.90639993053475,4.81646702181797,4.86597275735565,4.80301105931374,4.82879152257126,4.62739393484642,3.90639993053475,4.64480715925768,3.90639993053475,4.78833353577395,5.17070992705654,5.60504772420003,3.90639993053475,5.33575103628752,5.25267609283373,4.32258542707319,4.61444274475167,3.90639993053475,4.57480027762229,4.50597084586879,4.68567356515885,4.90085523822078,4.47334086467673,4.2229072215984,4.54083712727771,5.13189585967976,4.47895235448419,4.18561543703913,4.45058215927641,4.52736283465815,5.02441627929254,4.56983126722696,4.33341731043495,3.90639993053475,4.78799141211741,5.0395702658471,3.90639993053475,5.07154798309456,3.90639993053475,4.62853782275497,3.90639993053475,5.07592570524837,4.56210553053287,4.85504934898149,4.51443428708131,5.13116524639484,5.5077145306243
+"gene7",5.20626518213863,6.4307471888351,5.75210406715117,5.74808550473783,6.01939586675176,5.8745194991298,6.13178297125779,5.95194888534063,5.45762477496646,5.62410993865574,5.71434824907693,5.32734909791476,5.55537558435893,5.73025744292854,5.06385888528879,5.5826032690162,5.67761863116132,5.03696211651269,5.79031038657492,5.49291611964151,5.83614340584532,5.68244859820213,5.14478019249239,6.11673133697475,5.145075058655,5.61248937747078,5.46533478118324,5.52050807084074,5.49647181397846,5.25825776961192,5.05031768963774,5.58075364322135,5.90006701024256,6.05111486808571,5.6138317348913,5.26328055784143,5.8094330953228,5.24543303639875,5.80448912653929,5.49087181109023,5.23774663083511,7.14102284074606,5.26582165079316,5.50600150487825,5.98769621400992,5.93554371352455,5.86283058601078,5.55265299728728,6.40444973018289,5.43356280169462,5.24390145105428,5.73736446157673,6.11714238755379,5.24220541394816,5.3129473376966,5.41958134103275,5.58157497151541,5.15634558288039,5.76937316488221,5.76074295358258,6.0610987606876,5.83249716568378,5.27120888672809,5.64503349702179,5.23799835582669,5.49993552310687,6.28901440098437,5.21266063276808,5.73944730119365,5.01007270740896,6.5533819115425,6.50561549201741,5.7757501161297,6.52022614520828,5.63890672797129,5.75410652002562,6.18071466943569,5.70610201213676,5.26711315363607,6.17684628978418,5.90810242918654,5.388660147373,5.7284573023893,5.45234292741033,5.46840552849758,6.32399501173324,5.38571881845051,5.4993304827329,5.85147258005921,6.27612665515958,4.65506623431438,5.43348087169527,6.04950417355331,5.42502558647655,5.60864452182457,5.86893917164358,5.67609640966005,6.87376921628941,5.89696827981492,5.80353625768594
+"gene8",6.14467087636535,5.71799462942032,5.91550572881734,5.33586208024885,5.93204323763414,6.0096602308351,5.64195288696069,6.63942328316715,5.97963564651229,6.05941803041011,6.04025404334894,5.98754618245661,5.64157545093473,6.76385500596844,5.00074087291706,6.64588152903788,4.99965620724135,6.33069844247543,5.87539319546207,6.22421601005952,6.20828001589787,5.99192127349128,5.45386533306459,7.17563818562694,6.20843787452718,6.93396692822002,6.59492940071907,5.76760163633124,6.80576166346598,5.86759608299833,6.18658861547704,6.46022351569673,6.072941859461,5.39016933201054,6.52316970045278,6.1810753587624,5.85577162435136,6.45174340203264,5.62566635815065,6.3824574071253,6.28881493815904,6.19841921791823,6.07415946719002,6.20690608438626,6.05131589754557,6.07368311692547,5.97517328165559,6.10693522873252,5.56068360124212,6.13664172953092,5.74037735152393,6.14154075132359,5.40029452763489,6.29582332531445,5.8845003617069,5.54468677263023,5.21887308147853,5.75514570254157,5.94755515722157,5.41601092050764,6.23217091731421,5.83249716568378,6.13208416441675,5.84057019050038,5.98769878809005,6.49899924177197,6.26847398654057,6.45862625312591,5.44928948817043,6.28103655384944,5.96133370977168,5.69675901610658,5.25175578906085,6.20978390018059,6.01403709277561,5.91337259133489,6.48868023076156,6.43166054375869,5.9722655403051,5.65078289134718,4.85780705232998,5.81286109053848,5.58077108528882,6.1360997887631,6.35955658533746,5.92339714935639,5.63924002671109,6.11567299868842,5.27795806796843,5.89245220032049,5.7759079284252,5.79229606289696,5.7536406607985,6.04510940617195,5.76760163633124,5.92990496262424,6.38813863601087,5.97517328165559,5.48490734126876,7.24031978174728
+"gene9",6.05480399907477,6.22762410355796,6.53522149115282,5.8854574125321,5.99093119131383,6.0096602308351,6.10760366848201,6.05504057820223,6.70206198135364,5.76466666690829,4.71091569701823,6.00394431904936,6.40837401807306,6.13488469216531,5.69081326112896,5.85227382867697,5.92876090834259,6.53876657309679,5.46981752233177,6.22421601005952,6.04427474507967,6.2967841497099,5.8507881207722,5.92523578167103,6.15619465725473,5.61248937747078,5.79064468153218,5.42514345954339,6.03826397499398,6.06037427087577,5.91665603502672,5.68393471042058,6.86513217996736,6.8270209510656,5.56161639414299,5.84484205546352,6.11565629283893,5.8460878156094,6.12925335431688,5.84484205546352,5.29284421546781,5.87197535857028,5.83192750007993,6.28753144005393,5.87413038827588,6.95557940609427,5.83311948477393,5.82562124989778,6.11850213014837,5.8110767565011,6.57846188714706,5.30102845882608,5.44830354406908,6.08933910413573,5.17373134991574,6.24919405303723,6.03950880150654,6.18376519617616,5.59317333859299,5.83992874428376,5.33970877568687,5.32654494060728,4.78833353577395,6.22811267379654,6.74684354509594,5.90591170528687,5.57336538183939,6.33458294224889,5.86210658576285,6.06986040057169,5.96133370977168,6.06711264159453,5.84838285765499,5.32431472905752,6.01403709277561,6.13297359512698,5.81863260364644,5.89762641878417,5.9722655403051,5.36789263492867,5.3740098093694,5.87113782480949,5.51616826471844,5.78822566039613,5.34290195407164,5.66383247211807,5.76890399838076,5.98310801489499,5.74961185698385,5.72191437167912,5.31579123214392,6.09075620060001,5.82562124989778,6.65222021000827,6.41712344428707,6.11448649052405,5.90289850547098,5.54879916942742,6.78580981523118,6.21235355154699
+"gene10",4.48461312582745,4.91109371236248,4.43112355685642,4.99675646607325,4.80377838167823,5.31016612374788,4.85142717283359,5.07582160918152,4.67874694503383,4.91969148890903,4.53275039705426,4.90898495828709,4.74698593817716,5.69300708787934,4.46331870448127,5.23193392741788,4.85862539548854,5.2203939736305,4.90722362186415,4.81630357905239,4.62424112829588,5.84668232027055,5.27134724359215,5.23355519773888,4.96201763735775,4.77843461674337,4.88072772310509,5.80432268654145,5.45557251556794,4.37166805744251,3.90639993053475,4.8362803761711,6.17662668667575,4.88232511621218,4.76732128592982,4.8509110587817,5.40180868158091,5.15852946755355,4.85019207503742,4.97220168588376,3.90639993053475,5.19047380887536,4.54817390061422,5.21753230710055,5.0502037545777,5.20925145472013,5.26790999434977,4.83245842814059,4.97115078133684,5.01027037358841,4.49330166171214,5.229278998814,4.49256527164257,6.0610987606876,4.95174684922395,5.07149491885543,5.12229510658457,5.79527374920661,5.74171582273564,4.52915573807794,5.52981529905257,4.57282271760458,4.85662973399156,5.27352641517063,5.29846224726948,5.22797930177233,4.56632339400683,5.58859461021889,4.49296958353567,6.00584709830513,4.86779405708399,5.51841194563855,5.69860111608122,4.80272135813137,4.56466421747971,4.95092364901734,4.35312025741755,4.92924257605653,5.77046542632653,4.56611089380066,4.38852392356471,4.91015002087856,5.25591627041862,3.90639993053475,4.37756441670915,5.70441395575873,5.13079579294071,4.83972447110787,5.16577392253431,6.22770230970669,5.22490461338371,4.26312177185313,6.04950417355331,5.75915349171755,4.59374789830502,4.23636224946994,4.7578556011564,5.16546469094941,5.13116524639484,4.92967703958655
diff --git a/tests/data/single_factor/r_vst_with_design.csv b/tests/data/single_factor/r_vst_with_design.csv
new file mode 100644
index 0000000..f477970
--- /dev/null
+++ b/tests/data/single_factor/r_vst_with_design.csv
@@ -0,0 +1,11 @@
+"","sample1","sample2","sample3","sample4","sample5","sample6","sample7","sample8","sample9","sample10","sample11","sample12","sample13","sample14","sample15","sample16","sample17","sample18","sample19","sample20","sample21","sample22","sample23","sample24","sample25","sample26","sample27","sample28","sample29","sample30","sample31","sample32","sample33","sample34","sample35","sample36","sample37","sample38","sample39","sample40","sample41","sample42","sample43","sample44","sample45","sample46","sample47","sample48","sample49","sample50","sample51","sample52","sample53","sample54","sample55","sample56","sample57","sample58","sample59","sample60","sample61","sample62","sample63","sample64","sample65","sample66","sample67","sample68","sample69","sample70","sample71","sample72","sample73","sample74","sample75","sample76","sample77","sample78","sample79","sample80","sample81","sample82","sample83","sample84","sample85","sample86","sample87","sample88","sample89","sample90","sample91","sample92","sample93","sample94","sample95","sample96","sample97","sample98","sample99","sample100"
+"gene1",5.15188577077204,4.40174327214961,4.77639472593687,4.43022281894852,4.46352019254414,4.08073386831412,4.08073386831412,4.5406129339256,5.01424748866013,5.3460247592283,4.67087966030885,4.84007466884323,4.64462098742504,5.84374091926207,4.979818414006,6.16395378451766,4.38442486454377,5.92144112392647,4.69646998951852,4.38803500869893,4.6066696917302,5.08757811967048,4.55008314447337,5.41057427895843,4.40137448164242,5.66659793016854,4.08073386831412,5.11661919912923,4.45065976006065,4.39118356994307,4.94170173530474,4.08073386831412,5.07575330785696,4.7816149293192,4.61612464446728,4.56229054928712,4.66768783778319,4.08073386831412,4.57368432217848,5.27939877992168,4.47097583490813,5.15717886931126,4.60536914621491,5.62867035969357,4.41321732023935,4.08073386831412,5.21739861923921,4.56419266129158,5.52095953332127,4.71757969303984,5.28905082733315,4.99515439184409,5.39568267982756,5.22871486375103,4.78487972286996,4.64063029119078,4.92627193009595,5.75204855868008,5.32485129276651,4.66748697541779,4.85700725969574,4.4451935394628,5.0368016638169,5.50409061574604,4.89127399117978,5.26185784801724,5.47325828662498,4.99440621186768,4.56002805841379,5.79148008875772,5.81764890063073,5.03310290780184,4.87464254951838,5.02253033924907,5.08127876514478,4.91822066486707,4.74249119877235,5.15495898144975,5.08945586473665,5.57867700070401,4.66538198509024,4.59335186055529,5.02571099072279,5.18522787811702,5.81188854629463,4.87747066644894,5.3033647179355,5.0921989174048,5.27028703037226,4.93768317508674,4.08073386831412,5.07115910340064,4.95414145983085,5.80846449952818,4.54065572580809,4.89354732401987,4.48600898793645,5.33065856519251,5.53759650122584,5.96066141229176
+"gene2",5.47554517645182,6.05642328945557,4.77639472593687,5.16305889435507,5.31594458622321,5.00742268393106,5.41591411468145,5.45198258470373,5.82433024217546,5.08166813297949,5.51631322926792,5.55236840249985,5.04556299514444,5.38956228314752,6.13962321175383,5.42705390846331,5.61705652373982,5.3222251474079,5.2572794843596,5.3460247592283,4.98200666836623,5.58152411895311,5.60719691673865,5.54764241523781,5.39769425555525,5.60482335967922,5.6580906606296,5.69260557715343,5.4638040645559,5.9207163432048,5.55432159965885,5.7064344510006,5.36651655518967,5.06264353379207,4.89250645600796,5.27939877992168,5.11563352673828,5.26340123580948,5.53208003025634,5.70418645816005,5.81111458580963,5.52134430443788,5.66770363079147,5.3195030114567,4.81779380590215,5.4889724847405,5.68664921182534,5.32674966595318,5.89607758611067,5.64050844540654,5.88744740231782,5.63539620942578,6.1271979559339,5.66637259131873,5.63407154211932,6.22386341308896,5.7486668265357,5.8329356289008,5.58327994315195,5.8383035052146,5.39557729333203,5.8638744069977,5.44829397729518,5.04291434758202,5.14297779298356,5.56275201131668,5.58826940598248,5.27529156594596,5.64727241725273,5.21720554658181,4.98759588557767,5.64253234830366,6.03053128655786,6.2408232977897,5.8440811649321,5.40337314895631,5.13364121865156,4.98550171216854,6.08562182911605,5.84759542843369,5.55127123760922,5.79924348515849,5.50470502220029,5.60872513026141,4.9581748668164,6.97170423263467,6.15764144681995,6.11739072556396,6.45584513654346,5.33106427290183,5.93461410185425,6.11549173497186,5.45779293142341,5.2687742961278,5.60833571156008,5.59723838954522,5.53921938442692,6.30137034446396,5.37682912578172,5.35822608124627
+"gene3",4.70850222680366,4.53378295171736,5.21604384147846,4.57379220989414,5.00521885275466,4.55001331572984,5.5418528421068,5.1398723804706,5.31696739549808,4.93885109196513,5.28273442581356,4.84007466884323,4.73050144822003,5.22338353715035,5.33320354237821,5.04935214227435,5.07030987899852,4.74343079814791,4.08073386831412,4.7628096991036,4.38550027701368,4.49870556691844,4.81806624896613,4.41406613143638,4.63384963709741,4.42020914032473,5.59083286501822,5.11661919912923,5.10806729509865,4.08073386831412,4.47032472271342,4.08073386831412,4.92629126199397,4.68918198131135,4.94697765571108,4.66917233359138,4.66768783778319,5.30558173365349,4.3662592391926,5.27939877992168,5.3387340640047,5.15717886931126,4.68543768513651,5.00218278276894,4.88644068502629,4.08073386831412,5.26999309432674,4.47608769677565,5.01357453966483,4.53286443898555,4.7558925253508,4.08073386831412,6.07095800893559,4.47236619856901,4.4354179092876,4.08073386831412,4.57352321229553,4.08073386831412,4.56341817985584,4.75673601378887,4.43123435611639,4.52651635642495,4.84162908950453,4.71719415794666,4.89127399117978,4.75348494067161,4.70257111240641,4.40848609393085,4.56002805841379,4.41659459622587,4.78771306643583,5.35996733918172,4.93651238499745,4.38331763937936,4.7010047134256,4.45946194978131,4.37875190901151,4.76938098627883,4.08073386831412,4.39335026831382,4.77064934770743,4.88506689705643,4.83269000593633,4.4211535199372,4.52449604347986,4.64764404517331,5.08853386304479,4.74130052753168,4.74500956280777,4.46846424911592,5.23184503049142,4.7480941953545,5.40733557592229,4.08073386831412,4.40663324551183,4.61682119210904,4.77818957756358,4.65359370107013,4.69640794756134,4.83404121581877
+"gene4",7.10369450489444,6.38205703428671,7.73470252265655,7.03835178057671,7.60084500918134,5.94752029742203,7.70845746061675,6.36430695837461,7.00492927454742,7.2569945673589,7.4172958135097,6.00763842246031,7.16542333697261,6.49671634781243,7.69110987150111,7.13336260340202,7.52595469733247,7.62284687180669,7.64032313875761,6.34298193405643,6.8419359158003,6.89084421626128,7.11238446182112,6.57373316056836,6.52349493698854,7.79270421530661,7.58241128905712,7.547629264329,7.31804523323817,7.14916872083907,7.31553921406035,6.51283721941758,6.89178269446898,6.83034405074353,8.16702714841079,7.0916757609162,7.68015831407199,7.09002025377314,7.00784496384684,7.0916757609162,7.61143070480929,7.02567249555614,7.37813507702535,7.61925038349048,7.34849896137887,6.44712517740466,7.01049407188648,7.18664929578265,6.75620375012356,7.19636861577154,7.32527136327531,6.42029221396178,7.045015149893,6.39910216495875,7.23017803992004,6.03578142576633,7.18490782161418,6.03115075083919,6.86368403780847,6.86436938491895,6.69117855484581,5.35634843471408,7.33233319709589,7.03285150418963,6.99491059959387,6.60877707106943,6.91725231226239,6.52557811510314,6.76982206520655,6.8119721858379,7.36417956742442,5.8977494220606,7.08697368738767,6.73619131917466,6.96250790668533,5.92595004785291,7.12680502901762,6.27215287324314,7.01412795074183,6.30133596239875,7.19236676226075,7.05700011304077,7.0580256803367,7.23120032550025,7.3027695268334,7.07613439593972,7.19506911966758,6.58240861117832,7.25220934944514,7.11633483399279,7.27467875888306,6.78782534214855,7.45468521561099,6.92490535431832,6.55612350070276,6.37820448607664,6.98888917965873,7.01049407188648,6.64101545613558,6.5992706047494
+"gene5",5.37873835466932,5.12467428864133,5.61661654841462,6.20693931900714,5.5518106938252,5.69472106165657,5.27273967401167,5.45198258470373,5.17801269629665,5.79884567109593,4.76062739241133,5.79177345860004,5.91953727852402,5.65931841690087,5.92615071948271,5.77127343792939,5.33197911472018,4.93143514371856,5.93244865737312,6.01964444034624,5.02911435616976,5.47839731543459,5.85288601858095,5.57938445920529,6.09281983359542,4.75509577782212,5.36029221679324,6.73601832668824,5.42005017707771,5.35825311612684,5.1277857268054,5.80198561955582,5.5506462044885,5.81357755387668,5.76869263411257,6.29878377744674,5.57650124675927,6.05592704017915,5.53208003025634,5.44035181211308,5.48878679062451,5.46888113293865,5.64337860054478,5.52279575630576,5.11234808157338,5.59275760623387,5.68664921182534,5.79542518235963,5.27198398541006,5.07601629544213,5.49524660626003,5.65300822483732,5.66649616362811,5.58324703080232,6.05621174626756,6.20624348649061,5.32077839806141,6.18770608486137,6.12708919833697,6.15894424178251,5.88527158042356,5.8638744069977,5.95366089140955,5.97150541771936,6.05635862963815,6.50189005690649,5.55124720464732,6.13139605738663,6.05770079630103,6.09471530269849,5.40147612651964,5.60633293546671,5.0472106587303,5.29449241232188,5.38603643478096,6.29035454414213,5.13364121865156,6.2527506170807,5.40615423321998,6.04892982488693,5.70041320651756,6.39954753252802,6.00391529052593,5.64009619674695,6.55024079417995,5.57552425331387,6.68238363389965,5.65991042843523,5.47370104150324,6.13641462657797,6.3569752445334,6.07405200257831,5.79542518235963,6.88766838785099,5.79625290024495,5.09232797081691,6.21617513889777,5.38788082117759,6.42251361013997,5.59611292070092
+"gene6",4.08073386831412,5.07797740539937,4.6846289671625,5.10981069406019,4.62051091726214,4.65423510286738,5.61844926318484,4.5406129339256,4.80876193555065,4.08073386831412,5.03374857421791,4.08073386831412,4.54211341310052,5.09404510101469,4.979818414006,4.08073386831412,5.07030987899852,4.46557276246261,5.22747108750038,4.51450956533074,5.11607886664267,5.42291315285502,5.00755719477093,4.08073386831412,5.25061474118211,4.55972966431379,4.86117625827517,5.77117911608428,4.08073386831412,4.51893719811139,5.09369328649192,4.47776480834818,4.73937965443598,4.57874342616242,4.69778932123699,4.75866914054766,4.53662937067935,5.21930959853684,4.57368432217848,4.56229054928712,4.75262671232076,4.49566385912014,4.60536914621491,5.3195030114567,4.54990750972067,4.08073386831412,5.41287811054643,4.95414145983085,5.01357453966483,4.40108981920823,5.16747582836123,4.44446891376157,5.09890680791747,4.08073386831412,4.93900577483706,4.9858712192227,4.92627193009595,4.95067053726534,4.76024243658924,4.08073386831412,4.77669222912063,4.08073386831412,4.91238421119801,5.27497905412055,5.68916432473501,4.08073386831412,5.43205006036122,5.35294085147744,4.47267172037888,4.74800950544008,4.08073386831412,4.710574331632,4.64560711464007,4.81530818212932,5.01890925951096,4.61482046269474,4.37875190901151,4.67851238054287,5.23809206323592,4.6201144090924,4.34362537358982,4.59335186055529,4.66579484706365,5.13604853986285,4.70588290549571,4.48288076320401,4.08073386831412,4.91206052368548,5.15042739544018,4.08073386831412,5.18077868864679,4.08073386831412,4.76132295499545,4.08073386831412,5.18493474057936,4.69858914572352,4.97552813357349,4.65359370107013,5.23739791081462,5.59611292070092
+"gene7",5.30878631609857,6.48426582516173,5.83001883780196,5.82616543193521,6.08687481191557,5.94752029742203,6.1951972735632,6.02195889500517,5.54828070529726,5.70740437626593,5.79382432895452,5.42404487444465,5.64165987958821,5.80907301338948,5.17347951309013,5.66769482636493,5.75863417063627,5.14795245432401,5.86666712576157,5.58197767982505,5.91066025103481,5.76326054103095,5.2503345392559,6.18067895902902,5.25061474118211,5.69628437278148,5.55564085427318,5.60833571156008,5.58537373954347,5.35825311612684,5.16062678607308,5.66592587866231,5.97207094798289,6.11742733031517,5.69756880529461,5.36303382958083,5.88501831432708,5.34604802611539,5.88027328293783,5.58002523477268,5.3387340640047,7.17601326777081,5.365452584411,5.59447649844444,6.05635615207553,6.00617979003568,5.93629078782287,5.63905715973722,6.45878625658578,5.52531610152187,5.34459058991166,5.81588619897956,6.1810753997318,5.34297670053617,5.41032546944895,5.51197609685434,5.66671137551598,5.2613256055262,5.84658097893374,5.8383035052146,6.12704720854525,5.9071592177615,5.37058074851931,5.72743172122682,5.33897357909591,5.58868212075033,6.34706006114859,5.31486912410297,5.81788306984109,5.12244121081341,6.60321696744462,6.55685992678861,5.85269801203047,6.57103606731887,5.72156670235438,5.83193908404629,6.24241868610891,5.78592195414554,5.36668194294516,6.23868421040295,5.97979487447578,5.48248344437677,5.80734741882844,5.5432390275254,5.55857249937193,6.38089610250146,5.47967871069532,5.5881042004469,5.92538116623924,6.33459854652834,4.78638494156401,5.52523792223474,6.11587550311374,5.51717023284508,5.69260557715343,5.94215903206812,5.75717618753541,6.91495895692069,5.96909260154978,5.87935879855704
+"gene8",6.20763123104529,5.79731898414464,5.98691205343867,5.4321558672408,6.00281343089057,6.07750026388523,5.72448265212562,6.6868000736242,6.04859831580305,6.12542765252204,6.10696420434883,6.05621174626756,5.72412134222792,6.80785579531633,5.11358970006109,6.69307792969235,5.11256093139991,6.38738224811069,5.94835973860299,6.28442961001978,6.26903633018669,6.0604229242978,5.54469216692713,7.2098885310312,6.26918879425188,6.97368315952286,6.64356458931853,5.84488176522399,6.84867206868335,5.94086873847314,6.24808973331499,6.51283721941758,6.13846045867977,5.48392259010079,6.57389244806218,6.2427669015811,5.92951026526214,6.50461617122789,5.70889390535897,6.43748560422975,6.34686717640749,6.25951324558289,6.13963399305898,6.2677093662536,6.117621015345,6.13917488242438,6.04430396802265,6.17123178402497,5.64673451461618,6.19988460001859,5.8187747536763,6.20461112835194,5.49357872682357,6.35364478564138,5.95711052865714,5.63144229026949,5.32077839806141,5.8329356289008,6.01773244319177,5.50856993202156,6.29211501989043,5.9071592177615,6.19548782939717,5.91491101559664,6.05635862963815,6.5504414344541,6.3272001321939,6.51128867415851,5.54032461400927,6.33934577306899,6.03098740531738,5.77696978992833,5.35206500876176,6.27048883673336,6.08171461346018,5.98486127151444,6.54043209594718,6.48515095248439,6.04150587156159,5.73293599758437,4.97813920785721,5.88830858988085,5.66594255969477,6.19936176300371,6.41531248991205,5.9944994048468,5.72188574736605,6.17965824871627,5.37700592060904,5.9647522303576,5.85284939985105,5.86857243086382,5.83149234431811,6.11164155431396,5.84488176522399,6.00075717297615,6.44298748371563,6.04430396802265,5.57432913286479,7.27322415766453
+"gene9",6.12098177949224,6.28772212848399,6.58558859485673,5.95803020185283,6.05946991019158,6.07750026388523,6.17187636830922,6.12120972867116,6.74771317242837,5.84206670737048,4.83916791684755,6.07199702024223,6.46258785965499,6.19818951535888,5.77127343792939,5.92615071948271,5.9996570293055,6.58902945211368,5.5599205768767,6.28442961001978,6.11083746914096,6.35457402762931,5.92472379014958,5.99626732751264,6.21875122352933,5.69628437278148,5.86698788578929,5.5172826954951,6.1050472011873,6.12634908306324,5.98801797763061,5.76468407985136,6.90653713034541,6.86938702011608,5.64762633756077,5.91901330227162,6.17964213695873,5.92020966217466,6.19275704962974,5.91901330227162,5.39117962955902,5.94507597138285,5.90661225657696,6.34562602081728,5.94714644622877,6.99477778401224,5.90775673951507,5.9005576751122,6.18238683142044,5.88659591922247,6.62756976854303,5.39897347827466,5.53938358947464,6.15426605568309,5.27785128890345,6.3085646016973,6.10624631698122,6.24536377667197,5.67780481606995,5.914295056482,5.43582125010383,5.42327874219952,4.91238421119801,6.28819414451742,6.79129359431898,5.97768897662832,5.65886038760904,6.39114113691889,5.93559531079482,6.13549063835329,6.03098740531738,6.13284254921158,5.9224137566589,5.42115403457871,6.08171461346018,6.19634585578404,5.89384863478151,5.96972515998342,6.04150587156159,5.46268299171409,5.46851468106739,5.94427131816324,5.60418924612334,5.86466684410819,5.4388640761185,5.74543106017753,5.84613095946803,6.05194015142866,5.82762902481204,5.80107584996461,5.41303439232758,6.15563221067607,5.9005576751122,6.69924011626757,6.47106455393479,6.17851394549407,5.97479256863395,5.63537319072842,6.82923646688058,6.27297077817151
+"gene10",4.62545507291409,5.02860892999485,4.57499905700495,5.10981069406019,4.92699802342151,5.40767636050061,4.97209867571447,5.18483591274414,4.80876193555065,5.03675515610774,4.67087966030885,5.0266110537433,4.87327326114916,5.7733751797469,4.60536598261491,5.33320354237821,4.97891405493023,5.3222251474079,5.02494236630003,4.93885109196513,4.75726434875442,5.9207806002311,5.37071245701952,5.33474606131073,5.07687065834343,4.90301919962429,4.99984436773317,5.88011354666107,5.54632171596419,4.51893719811139,4.08073386831412,4.95775928993629,6.23847221556851,5.00135726343856,4.89250645600796,4.97161003306227,5.49502286138974,5.26340123580948,4.97092932371276,5.08652590020056,4.08073386831412,5.29376923136359,4.68543768513651,5.3195030114567,5.16051865335628,5.31162653653682,5.36744045274342,4.95414145983085,5.08552950768727,5.1226287141114,4.63365272959159,5.3306776501791,4.63295792262526,6.12704720854525,5.06713440529361,5.18072831283198,5.22897097661784,5.87142970571955,5.82005803308924,4.66748697541779,5.61722922728496,4.70870722037833,4.97702447833438,5.37278694886004,5.39652957773671,5.32944115626183,4.70257111240641,5.67342518130982,4.63333940300603,6.0738289490815,4.98759588557767,5.60633293546671,5.77873472820642,4.92599779644474,4.7010047134256,5.06635410087097,4.5014523998658,5.0458056070733,5.84762867726179,4.70237049352213,4.53482876134644,5.02771485141908,5.35602458321391,4.08073386831412,4.52449604347986,5.78430441462054,5.23704689734634,4.96101957976182,5.27028703037226,6.28779768453843,5.32651607741508,4.4166377857569,6.11587550311374,5.83677913121148,4.72846532289449,4.39142704400005,4.88355329088834,5.26999309432674,5.23739791081462,5.04621732214708
diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py
index 1a11ff8..d0cb8f3 100644
--- a/tests/test_edge_cases.py
+++ b/tests/test_edge_cases.py
@@ -115,7 +115,7 @@ def test_nan_factors():
         DeseqDataSet(counts=counts_df, clinical=clinical_df, design_factors="condition")
 
 
-def test_one_factors():
+def test_one_factor():
     """Test that a ValueError is thrown when the design factor takes only one value ."""
     counts_df = pd.DataFrame(
         {"gene1": [0, 1], "gene2": [4, 12]}, index=["sample1", "sample2"]
@@ -126,21 +126,6 @@ def test_one_factors():
         DeseqDataSet(counts=counts_df, clinical=clinical_df, design_factors="condition")
 
 
-def test_too_many_factors():
-    """Test that a ValueError is thrown when the design factor takes
-    more than two values."""
-    counts_df = pd.DataFrame(
-        {"gene1": [0, 1, 5], "gene2": [4, 12, 8]},
-        index=["sample1", "sample2", "sample3"],
-    )
-    clinical_df = pd.DataFrame(
-        {"condition": [0, 1, 2]}, index=["sample1", "sample2", "sample3"]
-    )
-
-    with pytest.raises(ValueError):
-        DeseqDataSet(counts=counts_df, clinical=clinical_df, design_factors="condition")
-
-
 def test_reference_level():
     """Test that a ValueError is thrown when the reference level is not one of the
     design factor values."""
@@ -154,10 +139,38 @@ def test_reference_level():
             counts=counts_df,
             clinical=clinical_df,
             design_factors="condition",
-            reference_level="control",
+            tested_level="control",
         )
 
 
+def test_lfc_shrinkage_coeff():
+    """Test that a KeyError is thrown when attempting to shrink an unexisting LFC
+    coefficient.
+    """
+
+    counts_df = load_example_data(
+        modality="raw_counts",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    clinical_df = load_example_data(
+        modality="clinical",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    dds = DeseqDataSet(
+        counts=counts_df, clinical=clinical_df, design_factors="condition"
+    )
+    dds.deseq2()
+
+    res = DeseqStats(dds)
+    res.summary()
+    with pytest.raises(KeyError):
+        res.lfc_shrink(coeff="this_coeff_does_not_exist")
+
+
 def test_indexes():
     """Test that a ValueError is thrown when the count matrix and the clinical data
     don't have the same index."""
@@ -174,6 +187,48 @@ def test_indexes():
         )
 
 
+def test_contrast():
+    """Test that KeyErrors/ValueErrors are thrown when invalid contrasts are passed to a
+    DeseqStats."""
+
+    counts_df = load_example_data(
+        modality="raw_counts",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    clinical_df = load_example_data(
+        modality="clinical",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    # Run analysis
+    dds = DeseqDataSet(
+        counts=counts_df,
+        clinical=clinical_df,
+        refit_cooks=False,
+        design_factors=["condition", "group"],
+    )
+    dds.deseq2()
+
+    # Too short
+    with pytest.raises(ValueError):
+        DeseqStats(dds, contrast=["condition", "B"])
+
+    # Unexisting factor
+    with pytest.raises(KeyError):
+        DeseqStats(dds, contrast=["batch", "Y", "X"])
+
+    # Unexisting factor reference level
+    with pytest.raises(KeyError):
+        DeseqStats(dds, contrast=["condition", "B", "C"])
+
+    # Unexisting tested level
+    with pytest.raises(KeyError):
+        DeseqStats(dds, contrast=["condition", "C", "B"])
+
+
 def test_cooks_not_refitted():
     """Test that an AttributeError is thrown when a `DeseqStats` object is initialized
     from a `DeseqDataSet` whose `refit_cooks` attribute is set to True, but whose
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index f7f92c0..f2b475c 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -152,7 +152,7 @@ def test_lfc_shrinkage(tol=0.02):
     res = DeseqStats(dds)
     res.summary()
     res.SE = r_res.lfcSE * np.log(2)
-    res.lfc_shrink()
+    res.lfc_shrink(coeff="condition_B_vs_A")
     shrunk_res = res.results_df
 
     # Check that the same LFC are found (up to tol)
@@ -163,8 +163,6 @@ def test_lfc_shrinkage(tol=0.02):
 
 
 # Multi-factor tests
-
-
 def test_multifactor_deseq(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.
@@ -255,7 +253,7 @@ def test_multifactor_lfc_shrinkage(tol=0.02):
     res = DeseqStats(dds)
     res.summary()
     res.SE = r_res.lfcSE * np.log(2)
-    res.lfc_shrink()
+    res.lfc_shrink(coeff="condition_B_vs_A")
     shrunk_res = res.results_df
 
     # Check that the same LFC found (up to tol)
@@ -367,3 +365,49 @@ def test_anndata_init(tol=0.02):
     ).max() < tol
     assert (abs(r_res.pvalue - res_df.pvalue) / r_res.pvalue).max() < tol
     assert (abs(r_res.padj - res_df.padj) / r_res.padj).max() < tol
+
+
+def test_vst(tol=0.02):
+    """
+    Test the output of VST compared with DESeq2.
+    """
+
+    # Load data
+    counts_df = load_example_data(
+        modality="raw_counts",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    clinical_df = load_example_data(
+        modality="clinical",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    # Load R data
+    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
+
+    r_vst = pd.read_csv(
+        os.path.join(test_path, "data/single_factor/r_vst.csv"), index_col=0
+    ).T
+
+    r_vst_with_design = pd.read_csv(
+        os.path.join(test_path, "data/single_factor/r_vst_with_design.csv"), index_col=0
+    ).T
+
+    # Test blind design
+    dds = DeseqDataSet(
+        counts=counts_df, clinical=clinical_df, design_factors=["condition"]
+    )
+    dds.vst(use_design=False)
+    assert (np.abs(r_vst - dds.layers["vst_counts"]) / r_vst).max().max() < tol
+
+    # Test full design
+    dds = DeseqDataSet(
+        counts=counts_df, clinical=clinical_df, design_factors=["condition"]
+    )
+    dds.vst(use_design=True)
+    assert (
+        np.abs(r_vst_with_design - dds.layers["vst_counts"]) / r_vst_with_design
+    ).max().max() < tol