Small fixes

src/rnanorm/methods/between_sample.py

@@ -1,6 +1,9 @@
 """Between sample normalizations."""
+from typing import Any, Optional
+
 import numpy as np
 from scipy.stats import gmean, rankdata, scoreatpercentile
+from sklearn import config_context
 from sklearn.base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
 from sklearn.utils.validation import check_is_fitted
 
@@ -70,7 +73,9 @@ def _get_norm_factors(self, X: Numeric2D) -> Numeric1D:
         :param X: Expression raw count matrix (n_samples, n_features)
         """
         X = remove_zero_genes(X)
-        lib_size = LibrarySize().fit_transform(X)
+        # Make sure that global set_config(transform_output="pandas")
+        # does not affect this method - we need numpy output here.
+        lib_size = LibrarySize().set_output(transform="default").fit_transform(X)
 
         # Compute upper quartile count for each sample.
         # No numpy method can be used as drop-in replacement for R's quantile.
@@ -97,7 +102,7 @@ def _reset(self) -> None:
         if hasattr(self, "geometric_mean_"):
             del self.geometric_mean_
 
-    def fit(self, X: Numeric2D) -> Self:
+    def fit(self, X: Numeric2D, y: Optional[Numeric1D] = None, **fit_params: Any) -> Self:
         """Fit.
 
         :param X: Expression raw count matrix (n_samples, n_features)
@@ -122,7 +127,8 @@ def transform(self, X: Numeric2D) -> Numeric2D:
 
         # Compute effective library sizes
         factors = self.get_norm_factors(X)
-        effective_lib_size = LibrarySize().fit_transform(X) * factors
+        lib_size = LibrarySize().set_output(transform="default").fit_transform(X)
+        effective_lib_size = lib_size * factors
 
         # Make CPM, but with effective library size
         return X / effective_lib_size[:, np.newaxis] * 1e6
@@ -241,8 +247,10 @@ def _get_norm_factors(self, X: Numeric2D) -> Numeric1D:
         """
         X = remove_zero_genes(X)
 
-        lib_size = LibrarySize().fit_transform(X)
-        lib_size_ref = LibrarySize().fit_transform(self.ref_[np.newaxis, :])
+        # ensure that output of transform will be a np.array
+        with config_context(transform_output="default"):
+            lib_size = LibrarySize().fit_transform(X)
+            lib_size_ref = LibrarySize().fit_transform(self.ref_[np.newaxis, :])
 
         # Values 0 cause a lot of troubles and warnings in log / division.
         # But computing with np.nan is OK, and is handled gracefully.
@@ -329,7 +337,7 @@ def _get_ref(self, X: Numeric2D) -> Numeric1D:
         ref_index = np.argmin(np.fabs(f75 - np.mean(f75)))
         return X[ref_index, :]
 
-    def fit(self, X: Numeric2D) -> Self:
+    def fit(self, X: Numeric2D, y: Optional[Numeric1D] = None, **fit_params: Any) -> Self:
         """Fit.
 
         :param X: Expression raw count matrix (n_samples, n_features)
@@ -354,7 +362,8 @@ def transform(self, X: Numeric2D) -> Numeric2D:
         """
         # Compute effective library sizes
         factors = self.get_norm_factors(X)
-        effective_lib_size = LibrarySize().fit_transform(X) * factors
+        lib_size = LibrarySize().set_output(transform="default").fit_transform(X)
+        effective_lib_size = lib_size * factors
 
         # Method ``check_is_fitted`` is not called here, since it is
         # called in self.get_norm_factors

tests/test_ctf.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from sklearn import config_context
 
 from rnanorm import CTF
 
@@ -44,3 +45,11 @@ def test_ctf(exp, expected_factors, expected_ctf):
         expected_ctf.loc[["Sample_2"]],
         rtol=1e-3,
     )
+
+
+def test_global_set_output(exp):
+    """Ensure that global config does not break things."""
+    with config_context(transform_output="pandas"):
+        CTF().fit_transform(exp)
+
+    CTF().set_output(transform="pandas").fit_transform(exp)

tests/test_cuf.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from sklearn import config_context
 
 from rnanorm import CUF
 
@@ -44,3 +45,11 @@ def test_cuf(exp, expected_factors, expected_cuf):
         expected_cuf.loc[["Sample_2"]],
         rtol=1e-3,
     )
+
+
+def test_global_set_output(exp):
+    """Ensure that global config does not break things."""
+    with config_context(transform_output="pandas"):
+        CUF().fit_transform(exp)
+
+    CUF().set_output(transform="pandas").fit_transform(exp)

tests/test_sklearn_compat.py

@@ -0,0 +1,37 @@
+import pandas as pd
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+
+from rnanorm import CPM, CTF, CUF, FPKM, TMM, TPM, UQ
+from rnanorm.datasets import load_toy_data
+
+
+def test_grid_search():
+    """Test compatibility of all methods with sklearn machinery."""
+    ds = load_toy_data()
+    X = ds.exp
+    y = pd.Series([0, 0, 1, 1], index=X.index)
+    pipeline = Pipeline(
+        steps=[
+            ("normalization", CPM()),
+            ("scaler", StandardScaler()),
+            ("classifier", LogisticRegression()),
+        ]
+    )
+    params = {
+        "normalization": [
+            CPM(),
+            FPKM(gtf=ds.gtf_path),
+            TPM(gtf=ds.gtf_path),
+            UQ(),
+            CUF(),
+            TMM(),
+            CTF(),
+        ],
+    }
+    search = GridSearchCV(pipeline, params, cv=2, refit=False)
+    search.fit(X, y)
+    results = pd.DataFrame(search.cv_results_)
+    assert results.shape[0] == 7

tests/test_tmm.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from sklearn import config_context
 
 from rnanorm import TMM
 from rnanorm.datasets import load_gtex
@@ -70,3 +71,11 @@ def test_tmm_rnanorm_edger():
         rnanorm_factors,
         decimal=14,
     )
+
+
+def test_global_set_output(exp):
+    """Ensure that global config does not break things."""
+    with config_context(transform_output="pandas"):
+        TMM().fit_transform(exp)
+
+    TMM().set_output(transform="pandas").fit_transform(exp)

tests/test_uq.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from sklearn import config_context
 
 from rnanorm import UQ
 from rnanorm.datasets import load_gtex
@@ -71,3 +72,11 @@ def test_uq_rnanorm_edger():
         rnanorm_factors,
         decimal=14,
     )
+
+
+def test_global_set_output(exp):
+    """Ensure that global config does not break things."""
+    with config_context(transform_output="pandas"):
+        UQ().fit_transform(exp)
+
+    UQ().set_output(transform="pandas").fit_transform(exp)