DimReducer: add prefix param + tests (v0.3.4)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "centimators"
-version = "0.3.3"
+version = "0.3.4"
 description = "essential data transformers and model estimators for ML and data science competitions"
 readme = "README.md"
 authors = [

src/centimators/feature_transformers/dimreduction.py

@@ -1,6 +1,7 @@
 """Dimensionality reduction transformers for feature compression."""
 
 import narwhals as nw
+import numpy as np
 from narwhals.typing import FrameT
 from sklearn.decomposition import PCA
 from sklearn.manifold import TSNE
@@ -9,88 +10,67 @@
 
 
 class DimReducer(_BaseFeatureTransformer):
-    """
-    DimReducer applies dimensionality reduction to features using PCA, t-SNE, or UMAP.
+    """Dimensionality reduction using PCA, t-SNE, or UMAP.
 
-    This transformer reduces the dimensionality of input features by projecting them
-    into a lower-dimensional space using one of three methods: Principal Component
-    Analysis (PCA), t-distributed Stochastic Neighbor Embedding (t-SNE), or Uniform
-    Manifold Approximation and Projection (UMAP).
+    Reduces ``feature_names`` columns into ``n_components`` output columns
+    named ``{prefix}_{i}``.
 
     Args:
-        method (str): The dimensionality reduction method to use. Options are:
-            - 'pca': Principal Component Analysis (linear, preserves global structure)
-            - 'tsne': t-SNE (non-linear, preserves local structure, visualization)
-            - 'umap': UMAP (non-linear, preserves local + global structure)
-            Default: 'pca'
-        n_components (int): Number of dimensions in the reduced space. Default: 2
-        feature_names (list[str] | None): Names of columns to reduce. If None,
-            all columns are used.
-        **reducer_kwargs: Additional keyword arguments passed to the underlying
-            reducer (sklearn.decomposition.PCA, sklearn.manifold.TSNE, or umap.UMAP).
+        method: ``"pca"``, ``"umap"``, or ``"tsne"``.
+        n_components: Number of dimensions in the reduced space.
+        feature_names: Columns to reduce.  If None, all columns are used.
+        prefix: Output column name prefix.  Default ``"dim"`` produces
+            ``dim_0, dim_1, ...``.
+        **reducer_kwargs: Forwarded to the underlying reducer.
+            Common: ``random_state=42``.
+
+    Notes:
+        To use GPU-accelerated cuML backends, swap the import before
+        constructing DimReducer::
+
+            import umap
+            from cuml.manifold import UMAP as cuUMAP
+            umap.UMAP = cuUMAP  # drop-in replacement
 
     Examples:
         >>> import polars as pl
-        >>> from centimators.feature_transformers import DimReducer
-        >>> df = pl.DataFrame({
-        ...     'feature1': [1.0, 2.0, 3.0, 4.0],
-        ...     'feature2': [4.0, 5.0, 6.0, 7.0],
-        ...     'feature3': [7.0, 8.0, 9.0, 10.0],
-        ... })
-        >>>
-        >>> # PCA reduction
-        >>> reducer = DimReducer(method='pca', n_components=2)
-        >>> reduced = reducer.fit_transform(df)
-        >>> print(reduced.columns)  # ['dim_0', 'dim_1']
-        >>>
-        >>> # t-SNE for visualization
-        >>> reducer = DimReducer(method='tsne', n_components=2, random_state=42)
+        >>> df = pl.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]})
+        >>> reducer = DimReducer(method="pca", n_components=2)
         >>> reduced = reducer.fit_transform(df)
-        >>>
-        >>> # UMAP (requires umap-learn)
-        >>> reducer = DimReducer(method='umap', n_components=2, random_state=42)
-        >>> reduced = reducer.fit_transform(df)
-
-    Notes:
-        - PCA is deterministic and fast, suitable for preprocessing
-        - t-SNE is stochastic and slower, primarily for visualization (does not support
-          separate transform - uses fit_transform internally)
-        - UMAP balances speed and quality, good for both preprocessing and visualization
-        - UMAP requires the umap-learn package: `uv add 'centimators[all]'`
-        - All methods work with any narwhals-compatible backend (pandas, polars, etc.)
+        >>> reduced.columns
+        ['dim_0', 'dim_1']
+
+        >>> reducer = DimReducer(
+        ...     method="pca", n_components=2, prefix="emb_thesis"
+        ... )
+        >>> reducer.fit_transform(df).columns
+        ['emb_thesis_0', 'emb_thesis_1']
     """
 
+    _VALID_METHODS = ("pca", "tsne", "umap")
+
     def __init__(
         self,
         method: str = "pca",
         n_components: int = 2,
         feature_names: list[str] | None = None,
+        prefix: str = "dim",
         **reducer_kwargs,
     ):
         super().__init__(feature_names=feature_names)
-
-        valid_methods = ["pca", "tsne", "umap"]
-        if method not in valid_methods:
-            raise ValueError(f"method must be one of {valid_methods}, got '{method}'")
-
+        if method not in self._VALID_METHODS:
+            raise ValueError(
+                f"method must be one of {self._VALID_METHODS}, got {method!r}"
+            )
         self.method = method
         self.n_components = n_components
+        self.prefix = prefix
         self.reducer_kwargs = reducer_kwargs
         self._reducer = None
 
     def fit(self, X: FrameT, y=None):
-        """Fit the dimensionality reduction model.
-
-        Args:
-            X (FrameT): Input data frame.
-            y: Ignored. Kept for compatibility.
-
-        Returns:
-            DimReducer: The fitted transformer.
-        """
         super().fit(X, y)
 
-        # Initialize the appropriate reducer
         if self.method == "pca":
             self._reducer = PCA(n_components=self.n_components, **self.reducer_kwargs)
         elif self.method == "tsne":
@@ -100,65 +80,41 @@ def fit(self, X: FrameT, y=None):
                 import umap
             except ImportError as e:
                 raise ImportError(
-                    "DimReducer with method='umap' requires umap-learn. Install with:\n"
-                    "  uv add 'centimators[all]'\n"
-                    "or:\n"
-                    "  pip install 'centimators[all]'"
+                    "DimReducer with method='umap' requires umap-learn. "
+                    "Install with: uv pip install 'centimators[all]'"
                 ) from e
             self._reducer = umap.UMAP(
                 n_components=self.n_components, **self.reducer_kwargs
             )
 
-        # Fit the reducer on the selected features
         X_native = nw.from_native(X)
-        X_subset = X_native.select(self.feature_names)
-        X_numpy = X_subset.to_numpy()
+        X_numpy = np.asarray(
+            X_native.select(self.feature_names).to_numpy(), dtype=np.float32
+        )
 
-        # For t-SNE, we skip fit since it doesn't support separate fit/transform
         if self.method != "tsne":
             self._reducer.fit(X_numpy)
 
         return self
 
     @nw.narwhalify(allow_series=True)
     def transform(self, X: FrameT, y=None) -> FrameT:
-        """Transform features by reducing their dimensionality.
-
-        Args:
-            X (FrameT): Input data frame.
-            y: Ignored. Kept for compatibility.
-
-        Returns:
-            FrameT: Transformed data frame with reduced dimensionality.
-                Columns are named 'dim_0', 'dim_1', ..., 'dim_{n_components-1}'.
-        """
         if self._reducer is None:
             raise ValueError("Transformer not fitted. Call fit() first.")
 
-        # Extract features and convert to numpy
-        X_subset = X.select(self.feature_names)
-        X_numpy = X_subset.to_numpy()
+        X_numpy = np.asarray(X.select(self.feature_names).to_numpy(), dtype=np.float32)
 
-        # Apply dimensionality reduction
-        # Note: t-SNE doesn't support transform(), so we use fit_transform
         if self.method == "tsne":
             X_reduced = self._reducer.fit_transform(X_numpy)
         else:
             X_reduced = self._reducer.transform(X_numpy)
 
-        # Create output column names
-        output_cols = {f"dim_{i}": X_reduced[:, i] for i in range(self.n_components)}
+        X_reduced = np.asarray(X_reduced)
 
-        # Return as narwhals DataFrame with the same backend as input
+        output_cols = {
+            f"{self.prefix}_{i}": X_reduced[:, i] for i in range(self.n_components)
+        }
         return nw.from_dict(output_cols, backend=nw.get_native_namespace(X))
 
     def get_feature_names_out(self, input_features=None) -> list[str]:
-        """Return the output feature names.
-
-        Args:
-            input_features (list[str], optional): Ignored. Kept for compatibility.
-
-        Returns:
-            list[str]: List of output feature names: ['dim_0', 'dim_1', ...].
-        """
-        return [f"dim_{i}" for i in range(self.n_components)]
+        return [f"{self.prefix}_{i}" for i in range(self.n_components)]

tests/test_feature_transformers.py

@@ -10,6 +10,7 @@
     MovingAverageTransformer,
     LogReturnTransformer,
     GroupStatsTransformer,
+    DimReducer,
 )
 
 # EmbeddingTransformer import with optional dependency check
@@ -247,3 +248,66 @@ def simple_embedder(texts):
 
     feature_names = transformer.get_feature_names_out()
     assert feature_names == ["text_embed_0", "text_embed_1", "text_embed_2"]
+
+
+def _make_reduction_frame():
+    rng = np.random.default_rng(0)
+    return pl.DataFrame({f"f{i}": rng.standard_normal(20) for i in range(5)})
+
+
+def test_dimreducer_default_prefix_backward_compat():
+    """Default prefix must remain ``dim`` (locks 0.3.x column names)."""
+    df = _make_reduction_frame()
+    reducer = DimReducer(method="pca", n_components=2)
+    reduced = reducer.fit_transform(df)
+    assert reduced.columns == ["dim_0", "dim_1"]
+
+
+def test_dimreducer_custom_prefix():
+    df = _make_reduction_frame()
+    reducer = DimReducer(method="pca", n_components=3, prefix="emb_thesis")
+    reduced = reducer.fit_transform(df)
+    assert reduced.columns == ["emb_thesis_0", "emb_thesis_1", "emb_thesis_2"]
+
+
+def test_dimreducer_get_feature_names_out_matches_transform():
+    df = _make_reduction_frame()
+    for prefix in ("dim", "emb_thesis"):
+        reducer = DimReducer(method="pca", n_components=2, prefix=prefix)
+        reduced = reducer.fit_transform(df)
+        assert reducer.get_feature_names_out() == reduced.columns
+
+
+def test_dimreducer_invalid_method_raises_at_construction():
+    """Invalid method must fail eagerly in __init__, not at fit."""
+    with pytest.raises(ValueError, match="method must be one of"):
+        DimReducer(method="not_a_method")
+
+
+def test_dimreducer_pandas_backend():
+    """Transformer is backend-agnostic (narwhals)."""
+    pd = pytest.importorskip("pandas")
+    df = pd.DataFrame(
+        {f"f{i}": np.random.default_rng(1).standard_normal(20) for i in range(4)}
+    )
+    reducer = DimReducer(method="pca", n_components=2, prefix="p")
+    reduced = reducer.fit_transform(df)
+    assert list(reduced.columns) == ["p_0", "p_1"]
+    assert len(reduced) == 20
+
+
+def test_dimreducer_tsne():
+    df = _make_reduction_frame()
+    reducer = DimReducer(method="tsne", n_components=2, perplexity=5.0, random_state=0)
+    reduced = reducer.fit_transform(df)
+    assert reduced.columns == ["dim_0", "dim_1"]
+    assert len(reduced) == 20
+
+
+def test_dimreducer_umap():
+    pytest.importorskip("umap")
+    df = _make_reduction_frame()
+    reducer = DimReducer(method="umap", n_components=2, random_state=0)
+    reduced = reducer.fit_transform(df)
+    assert reduced.columns == ["dim_0", "dim_1"]
+    assert len(reduced) == 20