Speed up MathFeatures reductions with NumPy

benchmarks/test_creation.py

@@ -18,8 +18,8 @@
     "func", [["sum", "mean"], ["sum", "mean", "std", "min", "max"]]
 )
 def test_math_features_transform(benchmark, df_tiny, func):
-    # MathFeatures aggregates row-wise, which is orders of magnitude slower per
-    # row than the vectorised transformers, hence the smallest dataframe.
+    # Keep the original dataset size so results remain comparable with the
+    # pre-NumPy baseline recorded for issue #576.
     creator = MathFeatures(variables=NUM_VARS, func=func)
     creator.fit(df_tiny)
     benchmark(creator.transform, df_tiny)

docs/user_guide/creation/MathFeatures.rst

@@ -6,9 +6,10 @@ MathFeatures
 ============
 
 :class:`MathFeatures()` applies basic functions to groups of features, returning one or
-more additional variables as a result. It uses `pandas.agg()` to create the features,
-so in essence, you can pass any function that is accepted by this method. One exception
-is that :class:`MathFeatures()` does not accept dictionaries for the parameter `func`.
+more additional variables as a result. It uses vectorized NumPy operations for common
+reductions and falls back to `pandas.agg()` for other functions. Thus, you can pass any
+function that is accepted by ``pandas.agg``. One exception is that
+:class:`MathFeatures()` does not accept dictionaries for the parameter `func`.
 
 The functions can be passed as strings, numpy methods, i.e., np.mean, or any function
 that you create, as long as it returns a scalar from a vector.
@@ -237,4 +238,4 @@ For tutorials about this and other feature engineering methods check out these r
 
 Both our book and courses are suitable for beginners and more advanced data scientists
 alike. By purchasing them you are supporting `Sole <https://linkedin.com/in/soledad-galli>`_,
-the main developer of feature-engine.
+the main developer of feature-engine.

feature_engine/creation/math_features.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -40,6 +41,33 @@
     np.prod: "prod",
 }
 
+# The kwargs preserve pandas' string-reduction defaults. In particular, pandas
+# uses one degree of freedom for ``std`` and ``var`` while NumPy uses zero.
+# NumPy callables have their direct pandas >= 3 semantics instead (ddof=0).
+_NUMPY_REDUCERS = {
+    "sum": (np.nansum, {}),
+    "mean": (np.nanmean, {}),
+    "std": (np.nanstd, {"ddof": 1}),
+    "var": (np.nanvar, {"ddof": 1}),
+    "min": (np.nanmin, {}),
+    "max": (np.nanmax, {}),
+    "prod": (np.nanprod, {}),
+    "median": (np.nanmedian, {}),
+    np.sum: (np.nansum, {}),
+    np.mean: (np.nanmean, {}),
+    np.std: (np.nanstd, {"ddof": 0}),
+    np.var: (np.nanvar, {"ddof": 0}),
+    np.min: (np.nanmin, {}),
+    np.max: (np.nanmax, {}),
+    np.prod: (np.nanprod, {}),
+    np.median: (np.median, {}),
+}
+
+
+def _get_numpy_reducer(func):
+    """Return the NumPy reducer for a supported aggregation."""
+    return _NUMPY_REDUCERS.get(func)
+
 
 @Substitution(
     missing_values=_missing_values_docstring,
@@ -54,8 +82,8 @@
 class MathFeatures(BaseCreation):
     """
     MathFeatures() applies functions across multiple features returning one or more
-    additional features as a result. It uses `pandas.agg()` to create the features,
-    setting `axis=1`.
+    additional features as a result. Common reductions use vectorized NumPy
+    operations. Other functions fall back to `pandas.agg()` with `axis=1`.
 
     For supported aggregation functions, see `pandas documentation
     <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.agg.html>`_.
@@ -234,10 +262,32 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
             else:
                 func = _FUNC_TO_STRING_ALIAS.get(func, func)
 
+        variables = X[self.variables]
+        functions = func if isinstance(func, list) else [func]
+        reducers = [_get_numpy_reducer(fun) for fun in functions]
+        values = variables.to_numpy()
+
+        # Nullable extension dtypes produce object arrays. Keep those, custom
+        # callables, and less common pandas aggregations on the exact legacy path.
+        if reducers and values.dtype.kind in "biuf" and all(reducers):
+            results = []
+            for reducer, kwargs in reducers:
+                # pandas' named reductions do not warn for empty/all-missing rows.
+                # NumPy returns the same values but emits RuntimeWarning for some
+                # reducers, so silence only those warnings on this equivalent path.
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", RuntimeWarning)
+                    result = reducer(values, axis=1, **kwargs)
+                results.append(pd.Series(result, index=X.index))
+
+            result = results[0] if len(results) == 1 else pd.concat(results, axis=1)
+        else:
+            result = variables.agg(func, axis=1)
+
         if len(new_variable_names) == 1:
-            X[new_variable_names[0]] = X[self.variables].agg(func, axis=1)
+            X[new_variable_names[0]] = result
         else:
-            X[new_variable_names] = X[self.variables].agg(func, axis=1)
+            X[new_variable_names] = result
 
         if self.drop_original:
             X.drop(columns=self.variables, inplace=True)

tests/test_creation/test_math_features.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import pandas as pd
 import pytest
@@ -288,6 +290,73 @@ def test_no_error_when_null_values_in_variable(df_vartypes):
     pd.testing.assert_frame_equal(X, ref)
 
 
+def test_standard_aggregations_match_pandas_with_missing_values():
+    X = pd.DataFrame(
+        {
+            "a": [1.0, np.nan, np.nan, 4.0],
+            "b": [3.0, 4.0, np.nan, 6.0],
+            "c": [5.0, 8.0, np.nan, np.nan],
+        }
+    )
+    functions = ["sum", "mean", "std", "var", "min", "max", "prod", "median"]
+    names = [f"result_{function}" for function in functions]
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)
+        expected = X.agg(functions, axis=1)
+    expected.columns = names
+
+    transformer = MathFeatures(
+        variables=list(X.columns),
+        func=functions,
+        new_variables_names=names,
+        missing_values="ignore",
+    )
+    result = transformer.fit_transform(X)
+
+    pd.testing.assert_frame_equal(result[names], expected)
+
+
+def test_nullable_dtypes_use_backwards_compatible_aggregation():
+    X = pd.DataFrame(
+        {
+            "a": pd.Series([1, pd.NA, 3], dtype="Int64"),
+            "b": pd.Series([2, 4, pd.NA], dtype="Int64"),
+        }
+    )
+    functions = ["sum", "mean"]
+    names = ["row_sum", "row_mean"]
+    expected = X.agg(functions, axis=1)
+    expected.columns = names
+
+    transformer = MathFeatures(
+        variables=list(X.columns),
+        func=functions,
+        new_variables_names=names,
+        missing_values="ignore",
+    )
+    result = transformer.fit_transform(X)
+
+    pd.testing.assert_frame_equal(result[names], expected)
+
+
+def test_custom_function_uses_pandas_aggregation_fallback(df_vartypes):
+    def peak_to_peak(row):
+        return row.max() - row.min()
+
+    expected = df_vartypes[["Age", "Marks"]].agg(peak_to_peak, axis=1)
+    transformer = MathFeatures(
+        variables=["Age", "Marks"],
+        func=peak_to_peak,
+        new_variables_names=["age_marks_range"],
+    )
+
+    result = transformer.fit_transform(df_vartypes)
+
+    pd.testing.assert_series_equal(
+        result["age_marks_range"], expected, check_names=False
+    )
+
+
 def test_drop_original_variables(df_vartypes):
     transformer = MathFeatures(
         variables=["Age", "Marks"],