Repository: owkin/PyDESeq2

Issue:
[BUG] deseq2_norm throws KeyError when called on a pandas DataFrame
**Describe the bug**
We are trying to use deseq2_norm to normalized your example data set and our own rna_seq experiment, in both cases the function does not work.

**To Reproduce**
python version = 3.11.3
code : 
counts_df = load_example_data(
    modality="raw_counts",
    dataset="synthetic",
    debug=False,
)

deseq2_norm(counts_df)

**Expected behavior**
We were expecting to get a dataframe with the normalized values of the gens in each sample.

**Screenshots**
![image](https://user-images.githubusercontent.com/130982623/235124582-72cd6a37-df72-43e0-8ce4-e7baf53f6679.png)
![image](https://user-images.githubusercontent.com/130982623/235124836-3c2f6e54-45e2-41ee-a170-9243dc0d45b2.png)


**Desktop (please complete the following information):**
 - OS: Linux
 - Version 22.04

**Additional context**
Add any other context about the problem here.


Discussion:

Patch:
diff --git a/pydeseq2/preprocessing.py b/pydeseq2/preprocessing.py
index 397f28c..2fa95ca 100644
--- a/pydeseq2/preprocessing.py
+++ b/pydeseq2/preprocessing.py
@@ -1,10 +1,13 @@
 from typing import Tuple
+from typing import Union
 
 import numpy as np
 import pandas as pd
 
 
-def deseq2_norm(counts: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
+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.
 
@@ -12,16 +15,16 @@ def deseq2_norm(counts: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
 
     Parameters
     ----------
-    counts : ndarray
+    counts : pandas.DataFrame or ndarray
             Raw counts. One column per gene, one row per sample.
 
     Returns
     -------
-    deseq2_counts : pandas.DataFrame
+    deseq2_counts : pandas.DataFrame or ndarray
         DESeq2 normalized counts.
         One column per gene, rows are indexed by sample barcodes.
 
-    size_factors : pandas.DataFrame
+    size_factors : pandas.DataFrame or ndarray
         DESeq2 normalization factors.
     """
 
@@ -32,7 +35,10 @@ def deseq2_norm(counts: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
     # Filter out genes with -∞ log means
     filtered_genes = ~np.isinf(logmeans)
     # Subtract filtered log means from log counts
-    log_ratios = log_counts[:, filtered_genes] - logmeans[filtered_genes]
+    if isinstance(log_counts, pd.DataFrame):
+        log_ratios = log_counts.loc[:, filtered_genes] - logmeans[filtered_genes]
+    else:
+        log_ratios = log_counts[:, filtered_genes] - logmeans[filtered_genes]
     # Compute sample-wise median of log ratios
     log_medians = np.median(log_ratios, axis=1)
     # Return raw counts divided by size factors (exponential of log ratios)


Test patch:
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 02cda60..ba33114 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -8,6 +8,7 @@ import pandas as pd
 import tests
 from pydeseq2.dds import DeseqDataSet
 from pydeseq2.ds import DeseqStats
+from pydeseq2.preprocessing import deseq2_norm
 from pydeseq2.utils import load_example_data
 
 # Single-factor tests
@@ -480,3 +481,34 @@ def test_ref_level():
             lambda x: 1 if x == "X" else 0 if x == "Y" else np.NaN
         )
     ).all()
+
+
+def test_deseq2_norm():
+    """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,
+    )
+
+    clinical_df = load_example_data(
+        modality="clinical",
+        dataset="synthetic",
+        debug=False,
+    )
+
+    # Fit size factors from DeseqDataSet
+    dds = DeseqDataSet(counts=counts_df, clinical=clinical_df)
+    dds.fit_size_factors()
+    s1 = dds.obsm["size_factors"]
+
+    # Fit size factors from counts directly
+    s2 = deseq2_norm(counts_df)[1]
+
+    np.testing.assert_almost_equal(
+        s1,
+        s2,
+        decimal=8,
+    )