Repository: owkin/PyDESeq2
Pull request #185 (2023-10-17 14:16:44)

Split vst functions into fit & transform
#### What does your PR implement? 
https://docs.google.com/presentation/d/1IyFTz4KdavaljiFbnCQitfE4N-9tdKZgWFjo8WOseKc/edit#slide=id.g29024929474_0_905

Additionally mutualized the synthetic data fixtures between the different tests in test_pydeseq2.py, let me know if this is not desirable.


hippolyteheuberger: TODO to finish the PR:
* Fix docstrings not written for deseq2_norm_fit() and deseq2_norm_transform()
* (optionally: improve test coverage of these two functions)
* Ensure test are passing
hippolyteheuberger: @SimonGrouard I have added the different tests :)

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1958739..5d36fbd 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -15,5 +15,4 @@ repos:
     rev: v1.7.0
     hooks:
     - id: mypy
-      exclude: ^(tests/|docs/source/conf.py)
-
+      exclude: ^(tests/|docs/source/conf.py)
\ No newline at end of file
diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py
index 11a0ab6..03bf68d 100644
--- a/pydeseq2/dds.py
+++ b/pydeseq2/dds.py
@@ -17,7 +17,8 @@ from scipy.stats import trim_mean  # type: ignore
 
 from pydeseq2.default_inference import DefaultInference
 from pydeseq2.inference import Inference
-from pydeseq2.preprocessing import deseq2_norm
+from pydeseq2.preprocessing import deseq2_norm_fit
+from pydeseq2.preprocessing import deseq2_norm_transform
 from pydeseq2.utils import build_design_matrix
 from pydeseq2.utils import dispersion_trend
 from pydeseq2.utils import make_scatter
@@ -145,6 +146,19 @@ 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()``.
+
+    filtered_genes: numpy.ndarray
+        Genes whose log means are different from -∞, computed in
+        preprocessing.deseq2_norm_fit().
+
     References
     ----------
     .. bibliography::
@@ -252,6 +266,8 @@ class DeseqDataSet(ad.AnnData):
         self.min_replicates = min_replicates
         self.beta_tol = beta_tol
         self.quiet = quiet
+        self.logmeans = None
+        self.filtered_genes = None
 
         # Initialize the inference object.
         self.inference = inference or DefaultInference()
@@ -261,8 +277,30 @@ class DeseqDataSet(ad.AnnData):
         use_design: bool = False,
         fit_type: Literal["parametric", "mean"] = "parametric",
     ) -> 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``).
+        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"``).
         """
-        Fit a variance stabilizing transformation, and apply it to normalized counts.
+        self.vst_fit(use_design=use_design, fit_type=fit_type)
+        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.
 
         Results are stored in ``dds.layers["vst_counts"]``.
 
@@ -277,9 +315,12 @@ class DeseqDataSet(ad.AnnData):
             gamma-family GLM. mean - use the mean of gene-wise dispersion estimates.
             (default: ``"parametric"``).
         """
-        # Start by fitting median-of-ratio size factors, if not already present.
-        if "size_factors" not in self.obsm:
-            self.fit_size_factors()
+        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 use_design:
             # Check that the dispersion trend curve was fitted. If not, fit it.
@@ -294,33 +335,77 @@ class DeseqDataSet(ad.AnnData):
             )
             # Fit the trend curve with an intercept design
             self.fit_genewise_dispersions()
-            if fit_type == "parametric":
+            if self.fit_type == "parametric":
                 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
-        if fit_type == "parametric":
+    def vst_transform(self, counts: Optional[np.ndarray] = None) -> np.ndarray:
+        """Apply the variance stabilizing transformation.
+
+        Uses the results from the ``vst_fit`` method.
+
+        Parameters
+        ----------
+        counts : numpy.ndarray
+            Counts to transform. If ``None``, use the counts from the current dataset.
+            (default: ``None``).
+
+        Returns
+        -------
+        numpy.ndarray
+            Variance stabilized counts.
+        """
+        if "size_factors" not in self.obsm:
+            raise RuntimeError(
+                "The vst_fit method should be called prior to vst_transform."
+            )
+
+        if counts is None:
+            # the transformed counts will be the current ones
+            normed_counts = self.layers["normed_counts"]
+        else:
+            if self.logmeans is None:
+                # the size factors were still computed iteratively
+                warnings.warn(
+                    "The size factors were fitted iteratively. They will "
+                    "be re-computed with the counts to be transformed. In a train/test "
+                    "setting with a downstream task, this would result in a leak of "
+                    "data from test to train set.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                logmeans, filtered_genes = deseq2_norm_fit(counts)
+            else:
+                logmeans, filtered_genes = self.logmeans, self.filtered_genes
+
+            normed_counts, _ = deseq2_norm_transform(counts, logmeans, filtered_genes)
+
+        if self.fit_type == "parametric":
             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)))
+            return np.log2(
+                (
+                    1
+                    + a1
+                    + 2 * a0 * normed_counts
+                    + 2 * np.sqrt(a0 * normed_counts * (1 + a1 + a0 * normed_counts))
+                )
                 / (4 * a0)
             )
-        elif fit_type == "mean":
+        elif self.fit_type == "mean":
             gene_dispersions = self.varm["genewise_dispersions"]
             use_for_mean = gene_dispersions > 10 * self.min_disp
             mean_disp = trim_mean(gene_dispersions[use_for_mean], proportiontocut=0.001)
-            self.layers["vst_counts"] = (
-                2 * np.arcsinh(np.sqrt(mean_disp * self.layers["normed_counts"]))
+            return (
+                2 * np.arcsinh(np.sqrt(mean_disp * normed_counts))
                 - np.log(mean_disp)
                 - np.log(4)
             ) / np.log(2)
         else:
             raise NotImplementedError(
-                f"Found fit_type '{fit_type}'. Expected 'parametric' or 'mean'."
+                f"Found fit_type '{self.fit_type}'. Expected 'parametric' or 'mean'."
             )
 
     def deseq2(self) -> None:
@@ -379,7 +464,11 @@ class DeseqDataSet(ad.AnnData):
             )
             self._fit_iterate_size_factors()
         else:
-            self.layers["normed_counts"], self.obsm["size_factors"] = deseq2_norm(self.X)
+            self.logmeans, self.filtered_genes = deseq2_norm_fit(self.X)
+            (
+                self.layers["normed_counts"],
+                self.obsm["size_factors"],
+            ) = deseq2_norm_transform(self.X, self.logmeans, self.filtered_genes)
         end = time.time()
 
         if not self.quiet:
diff --git a/pydeseq2/preprocessing.py b/pydeseq2/preprocessing.py
index 8b14c24..7704526 100644
--- a/pydeseq2/preprocessing.py
+++ b/pydeseq2/preprocessing.py
@@ -8,8 +8,7 @@ import pandas as pd
 def deseq2_norm(
     counts: Union[pd.DataFrame, np.ndarray]
 ) -> Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.DataFrame, np.ndarray]]:
-    """
-    Return normalized counts and size_factors.
+    """Return normalized counts and size_factors.
 
     Uses the median of ratios method.
 
@@ -27,12 +26,73 @@ def deseq2_norm(
     size_factors : pandas.DataFrame or ndarray
         DESeq2 normalization factors.
     """
+    logmeans, filtered_genes = deseq2_norm_fit(counts)
+    deseq2_counts, size_factors = deseq2_norm_transform(counts, logmeans, filtered_genes)
+    return deseq2_counts, size_factors
+
+
+def deseq2_norm_fit(
+    counts: Union[pd.DataFrame, np.ndarray]
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Return ``logmeans`` and ``filtered_genes``, needed in the median of ratios method.
+
+    ``Logmeans`` and ``filtered_genes`` can then be used to normalize external datasets.
+
+    Parameters
+    ----------
+    counts : pandas.DataFrame or ndarray
+            Raw counts. One column per gene, one row per sample.
+
+    Returns
+    -------
+    logmeans : ndarray
+        Gene-wise mean log counts.
+
+    filtered_genes : ndarray
+        Genes whose log means are different from -∞.
+    """
     # Compute gene-wise mean log counts
     with np.errstate(divide="ignore"):  # ignore division by zero warnings
         log_counts = np.log(counts)
     logmeans = log_counts.mean(0)
     # Filter out genes with -∞ log means
     filtered_genes = ~np.isinf(logmeans)
+
+    return logmeans, filtered_genes
+
+
+def deseq2_norm_transform(
+    counts: Union[pd.DataFrame, np.ndarray],
+    logmeans: np.ndarray,
+    filtered_genes: np.ndarray,
+) -> Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.DataFrame, np.ndarray]]:
+    """Return normalized counts and size factors from the median of ratios method.
+
+    Can be applied on external dataset, using the ``logmeans`` and ``filtered_genes``
+    previously computed in the ``fit`` function.
+
+    Parameters
+    ----------
+    counts : pandas.DataFrame or ndarray
+            Raw counts. One column per gene, one row per sample.
+
+    logmeans : ndarray
+        Gene-wise mean log counts.
+
+    filtered_genes : ndarray
+        Genes whose log means are different from -∞.
+
+    Returns
+    -------
+    deseq2_counts : pandas.DataFrame or ndarray
+        DESeq2 normalized counts.
+        One column per gene, rows are indexed by sample barcodes.
+
+    size_factors : pandas.DataFrame or ndarray
+        DESeq2 normalization factors.
+    """
+    with np.errstate(divide="ignore"):  # ignore division by zero warnings
+        log_counts = np.log(counts)
     # Subtract filtered log means from log counts
     if isinstance(log_counts, pd.DataFrame):
         log_ratios = log_counts.loc[:, filtered_genes] - logmeans[filtered_genes]


Test patch:
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 884eb5c..c3b2032 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -10,30 +10,38 @@ import tests
 from pydeseq2.dds import DeseqDataSet
 from pydeseq2.ds import DeseqStats
 from pydeseq2.preprocessing import deseq2_norm
+from pydeseq2.preprocessing import deseq2_norm_fit
+from pydeseq2.preprocessing import deseq2_norm_transform
 from pydeseq2.utils import load_example_data
 
 # Single-factor tests
 
 
-def test_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.
-    """
-
-    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
-
-    counts_df = load_example_data(
+@pytest.fixture
+def counts_df():
+    return load_example_data(
         modality="raw_counts",
         dataset="synthetic",
         debug=False,
     )
 
-    metadata = load_example_data(
+
+@pytest.fixture
+def metadata():
+    return load_example_data(
         modality="metadata",
         dataset="synthetic",
         debug=False,
     )
 
+
+def test_deseq(counts_df, metadata, tol=0.02):
+    """Test that the outputs of the DESeq2 function match those of the original R
+    package, up to a tolerance in relative error.
+    """
+
+    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
+
     r_res = pd.read_csv(
         os.path.join(test_path, "data/single_factor/r_test_res.csv"), index_col=0
     )
@@ -58,25 +66,13 @@ def test_deseq(tol=0.02):
 
 
 @pytest.mark.parametrize("alt_hypothesis", ["lessAbs", "greaterAbs", "less", "greater"])
-def test_alt_hypothesis(alt_hypothesis, tol=0.02):
+def test_alt_hypothesis(alt_hypothesis, counts_df, metadata, tol=0.02):
     """Test that the outputs of the DESeq2 function match those of the original R
     package, up to a tolerance in relative error, with the alternative hypothesis test.
     """
 
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     r_res = pd.read_csv(
         os.path.join(test_path, f"data/single_factor/r_test_res_{alt_hypothesis}.csv"),
         index_col=0,
@@ -120,7 +116,7 @@ def test_alt_hypothesis(alt_hypothesis, tol=0.02):
     # ).max() < tol
 
 
-def test_deseq_no_refit_cooks(tol=0.02):
+def test_deseq_no_refit_cooks(counts_df, metadata, tol=0.02):
     """Test that the outputs of the DESeq2 function *without cooks refit*
     match those of the original R package, up to a tolerance in relative error.
     Note: this is just to check that the workflow runs bug-free, as we expect no outliers
@@ -129,18 +125,6 @@ def test_deseq_no_refit_cooks(tol=0.02):
 
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     r_res = pd.read_csv(
         os.path.join(test_path, "data/single_factor/r_test_res.csv"), index_col=0
     )
@@ -169,7 +153,7 @@ def test_deseq_no_refit_cooks(tol=0.02):
     assert (abs(r_res.padj - res_df.padj) / r_res.padj).max() < tol
 
 
-def test_lfc_shrinkage(tol=0.02):
+def test_lfc_shrinkage(counts_df, metadata, tol=0.02):
     """Test that the outputs of the lfc_shrink function match those of the original
     R package (starting from the same inputs), up to a tolerance in relative error.
     """
@@ -183,18 +167,6 @@ def test_lfc_shrinkage(tol=0.02):
         index_col=0,
     )
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     r_size_factors = pd.read_csv(
         os.path.join(test_path, "data/single_factor/r_test_size_factors.csv"),
         index_col=0,
@@ -223,24 +195,13 @@ def test_lfc_shrinkage(tol=0.02):
     ).max() < tol
 
 
-def test_iterative_size_factors(tol=0.02):
+def test_iterative_size_factors(counts_df, metadata, tol=0.02):
     """Test that the outputs of the iterative size factor method match those of the
     original R package (starting from the same inputs), up to a tolerance in relative
     error.
     """
 
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
 
     r_size_factors = pd.read_csv(
         os.path.join(test_path, "data/single_factor/r_iterative_size_factors.csv"),
@@ -257,25 +218,13 @@ def test_iterative_size_factors(tol=0.02):
 
 
 # Multi-factor tests
-def test_multifactor_deseq(tol=0.02):
+def test_multifactor_deseq(counts_df, metadata, tol=0.02):
     """Test that the outputs of the DESeq2 function match those of the original R
     package, up to a tolerance in relative error.
     """
 
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     r_res = pd.read_csv(
         os.path.join(test_path, "data/multi_factor/r_test_res.csv"), index_col=0
     )
@@ -301,7 +250,7 @@ def test_multifactor_deseq(tol=0.02):
     assert (abs(r_res.padj - res_df.padj) / r_res.padj).max() < tol
 
 
-def test_multifactor_lfc_shrinkage(tol=0.02):
+def test_multifactor_lfc_shrinkage(counts_df, metadata, tol=0.02):
     """Test that the outputs of the lfc_shrink function match those of the original
     R package (starting from the same inputs), up to a tolerance in relative error.
     """
@@ -315,18 +264,6 @@ def test_multifactor_lfc_shrinkage(tol=0.02):
         index_col=0,
     )
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     r_size_factors = pd.read_csv(
         os.path.join(test_path, "data/multi_factor/r_test_size_factors.csv"),
         index_col=0,
@@ -458,24 +395,12 @@ def test_continuous_lfc_shrinkage(tol=0.02):
     ).max() < tol
 
 
-def test_contrast():
+def test_contrast(counts_df, metadata):
     """
     Check that the contrasts ['condition', 'B', 'A'] and ['condition', 'A', 'B'] give
     coherent results (change of sign in LFCs and Wald stats, same (adjusted p-values).
     """
 
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     dds = DeseqDataSet(
         counts=counts_df, metadata=metadata, design_factors=["group", "condition"]
     )
@@ -505,26 +430,13 @@ def test_contrast():
     )
 
 
-def test_anndata_init(tol=0.02):
+def test_anndata_init(counts_df, metadata, tol=0.02):
     """
     Test initializing dds with an AnnData object that already has filled in fields,
     including with the same names as those used by pydeseq2.
     """
     np.random.seed(42)
 
-    # Load data
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     # Make an anndata object
     adata = ad.AnnData(X=counts_df.astype(int), obs=metadata)
 
@@ -562,24 +474,10 @@ def test_anndata_init(tol=0.02):
     assert (abs(r_res.padj - res_df.padj) / r_res.padj).max() < tol
 
 
-def test_vst(tol=0.02):
+def test_vst(counts_df, metadata, 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,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     # Load R data
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
 
@@ -604,24 +502,10 @@ def test_vst(tol=0.02):
     ).max().max() < tol
 
 
-def test_mean_vst(tol=0.02):
+def test_mean_vst(counts_df, metadata, tol=0.02):
     """
     Test the output of VST with ``fitType="mean"`` compared with DESeq2.
     """
-
-    # Load data
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     # Load R data
     test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
 
@@ -635,22 +519,10 @@ def test_mean_vst(tol=0.02):
     assert (np.abs(r_vst - dds.layers["vst_counts"]) / r_vst).max().max() < tol
 
 
-def test_ref_level():
+def test_ref_level(counts_df, metadata):
     """Test that DeseqDataSet columns are created according to the passed reference
     level, if any.
     """
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     dds = DeseqDataSet(
         counts=counts_df,
         metadata=metadata,
@@ -669,22 +541,10 @@ def test_ref_level():
     ).all()
 
 
-def test_deseq2_norm():
+def test_deseq2_norm(counts_df, metadata):
     """Test that deseq2_norm() called on a pandas dataframe outputs the same results as
     DeseqDataSet.fit_size_factors()
     """
-    counts_df = load_example_data(
-        modality="raw_counts",
-        dataset="synthetic",
-        debug=False,
-    )
-
-    metadata = load_example_data(
-        modality="metadata",
-        dataset="synthetic",
-        debug=False,
-    )
-
     # Fit size factors from DeseqDataSet
     dds = DeseqDataSet(counts=counts_df, metadata=metadata)
     dds.fit_size_factors()
@@ -698,3 +558,67 @@ def test_deseq2_norm():
         s2,
         decimal=8,
     )
+
+
+@pytest.fixture
+def train_counts(counts_df):
+    return counts_df[25:75]
+
+
+@pytest.fixture
+def test_counts(counts_df):
+    return counts_df[0:25]
+
+
+@pytest.fixture
+def train_metadata(metadata):
+    return metadata[25:75]
+
+
+@pytest.fixture
+def train_dds(train_counts, train_metadata):
+    return DeseqDataSet(
+        counts=train_counts, metadata=train_metadata, design_factors="condition"
+    )
+
+
+def test_deseq2_norm_fit(train_counts):
+    # patient from 25 to 75
+    logmeans, filtered_genes = deseq2_norm_fit(train_counts)
+
+    # 10 genes
+    assert logmeans.shape == (10,)
+    assert filtered_genes.shape == (10,)
+
+
+def test_deseq2_norm_transform(train_counts, test_counts):
+    # First fit with some indices
+    logmeans, filtered_genes = deseq2_norm_fit(train_counts)
+
+    normed_counts, size_factors = deseq2_norm_transform(
+        test_counts, logmeans, filtered_genes
+    )
+    assert isinstance(normed_counts, pd.DataFrame)
+    # 25 samples, 10 genes
+    assert normed_counts.shape == (25, 10)
+    assert size_factors.shape == (25,)
+
+
+def test_vst_fit(train_dds):
+    # patient from 25 to 75
+    train_dds.vst_fit()
+
+    # the correct attributes are fit
+    assert "trend_coeffs" in train_dds.uns
+    assert "normed_counts" in train_dds.layers
+    assert "size_factors" in train_dds.obsm
+
+
+def test_vst_transform(train_dds, test_counts):
+    # First fit with some indices
+    train_dds.vst_fit()
+
+    result = train_dds.vst_transform(test_counts.to_numpy())
+    assert isinstance(result, np.ndarray)
+    # 25 samples, 10 genes
+    assert result.shape == (25, 10)