From f0c49898cfc5fd4ab46d564dfa2faa96530b0523 Mon Sep 17 00:00:00 2001 From: Adrian Hayler Date: Fri, 11 Sep 2026 17:09:32 +0200 Subject: [PATCH 1/3] Add automatic hurdle regression --- src/tabpfn_extensions/hurdle/__init__.py | 3 + src/tabpfn_extensions/hurdle/hurdle.py | 196 ++++++++++++++++++++++ tests/test_hurdle.py | 197 +++++++++++++++++++++++ 3 files changed, 396 insertions(+) create mode 100644 src/tabpfn_extensions/hurdle/__init__.py create mode 100644 src/tabpfn_extensions/hurdle/hurdle.py create mode 100644 tests/test_hurdle.py diff --git a/src/tabpfn_extensions/hurdle/__init__.py b/src/tabpfn_extensions/hurdle/__init__.py new file mode 100644 index 00000000..70a62997 --- /dev/null +++ b/src/tabpfn_extensions/hurdle/__init__.py @@ -0,0 +1,3 @@ +from tabpfn_extensions.hurdle.hurdle import AutoHurdleRegressor + +__all__ = ["AutoHurdleRegressor"] diff --git a/src/tabpfn_extensions/hurdle/hurdle.py b/src/tabpfn_extensions/hurdle/hurdle.py new file mode 100644 index 00000000..b92fc6ea --- /dev/null +++ b/src/tabpfn_extensions/hurdle/hurdle.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +from typing import Any, Literal + +import numpy as np +from sklearn.base import BaseEstimator, RegressorMixin, clone +from sklearn.utils import _safe_indexing +from sklearn.utils.validation import check_is_fitted + +from tabpfn_extensions.misc.sklearn_compat import validate_data +from tabpfn_extensions.utils import TabPFNClassifier, TabPFNRegressor + +_QUANTILE_GRID = ( + 0.01, + 0.05, + 0.1, + 0.15, + 0.2, + 0.25, + 0.3, + 0.35, + 0.4, + 0.45, + 0.5, + 0.55, + 0.6, + 0.65, + 0.7, + 0.75, + 0.8, + 0.85, + 0.9, + 0.95, + 0.99, +) + + +class AutoHurdleRegressor(RegressorMixin, BaseEstimator): + """Two-stage regression for non-negative targets with a point mass at zero. + + Args: + classifier: Cloneable classifier with predict_proba. Defaults to TabPFN. + regressor: Cloneable regressor with TabPFN's output_type and quantiles + prediction arguments. Defaults to TabPFN. + hurdle: Whether to split zero and positive targets. "auto" enables the + split for non-negative targets whose zero rate exceeds zero_threshold. + Otherwise, all rows go to the regressor. + zero_threshold: Training zero-rate threshold for automatic selection. + quantile_grid: Strictly increasing positive-stage quantile levels between + 0 and 1, with at least two entries. None uses the 21-point grid from + 0.01 to 0.99. Required levels outside the grid clamp to its endpoints. + + predict defaults to the mixture median, suitable for absolute error. With + positive probability p, it is zero for p <= 0.5 and otherwise the positive + distribution's (p - 0.5) / p quantile. Quantiles use linear interpolation on + the supplied grid, clamping levels outside that grid. Negative + positive-stage predictions are clipped to zero. Mean predictions multiply + the clipped positive-stage mean by p. + + Fitted attributes include hurdle_, zero_rate_, classifier_, and regressor_. + classifier_ is None when the hurdle is inactive. Both estimators are None + for an all-zero target with the hurdle enabled. + """ + + def __init__( + self, + classifier: Any = None, + regressor: Any = None, + *, + hurdle: Literal["auto"] | bool = "auto", + zero_threshold: float = 0.5, + quantile_grid: list[float] | tuple[float, ...] | np.ndarray | None = None, + ) -> None: + self.classifier = classifier + self.regressor = regressor + self.hurdle = hurdle + self.zero_threshold = zero_threshold + self.quantile_grid = quantile_grid + + def fit(self, X: Any, y: Any) -> AutoHurdleRegressor: + """Fit cloned estimators, preserving DataFrame columns and dtypes.""" + self.__dict__.pop("is_fitted_", None) + if self.hurdle not in ("auto", True, False): + raise ValueError("hurdle must be 'auto', True, or False.") + if not 0 <= self.zero_threshold <= 1: + raise ValueError("zero_threshold must be between 0 and 1.") + grid = np.array( + _QUANTILE_GRID if self.quantile_grid is None else self.quantile_grid, + dtype=float, + copy=True, + ) + if ( + grid.ndim != 1 + or grid.size < 2 + or not np.isfinite(grid).all() + or np.any((grid <= 0) | (grid >= 1)) + or np.any(np.diff(grid) <= 0) + ): + raise ValueError( + "quantile_grid must contain at least two finite, strictly increasing " + "levels strictly between 0 and 1." + ) + self.quantile_grid_ = grid + _, y = validate_data( + self, X, y, dtype=None, ensure_all_finite=False, y_numeric=True + ) + y = np.asarray(y, dtype=float) + if not np.isfinite(y).all(): + raise ValueError("Targets must be finite.") + self.zero_rate_ = float(np.mean(y == 0)) + self.hurdle_ = ( + bool(np.all(y >= 0) and self.zero_rate_ > self.zero_threshold) + if self.hurdle == "auto" + else bool(self.hurdle) + ) + if self.hurdle_ and np.any(y < 0): + raise ValueError("Hurdle modelling requires non-negative targets.") + positive = y > 0 + if self.hurdle_ and positive.all(): + raise ValueError("Hurdle modelling requires zero and positive targets.") + + classifier = None + regressor = None + if not self.hurdle_ or positive.any(): + regressor = clone( + self.regressor if self.regressor is not None else TabPFNRegressor() + ) + if self.hurdle_: + classifier = clone( + self.classifier + if self.classifier is not None + else TabPFNClassifier() + ) + classifier.fit(X, positive.astype(int)) + regressor.fit(_safe_indexing(X, positive), y[positive]) + else: + regressor.fit(X, y) + self.classifier_ = classifier + self.regressor_ = regressor + self.is_fitted_ = True + return self + + def predict( + self, + X: Any, + *, + output_type: Literal["mean", "median", "quantiles"] = "median", + quantiles: list[float] | None = None, + ) -> np.ndarray | list[np.ndarray]: + """Predict a mean, median, or list of quantile arrays in target units.""" + check_is_fitted(self, "is_fitted_") + validate_data(self, X, reset=False, dtype=None, ensure_all_finite=False) + if output_type not in ("mean", "median", "quantiles"): + raise ValueError("output_type must be 'mean', 'median', or 'quantiles'.") + if quantiles is None: + quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + if not quantiles or any(not 0 < q < 1 for q in quantiles): + raise ValueError("quantiles must be nonempty and strictly between 0 and 1.") + if not self.hurdle_: + return self.regressor_.predict( + X, output_type=output_type, quantiles=quantiles + ) + if self.regressor_ is None: + zeros = np.zeros(len(X)) + return ( + [zeros.copy() for _ in quantiles] + if output_type == "quantiles" + else zeros + ) + + positive_class = np.flatnonzero(self.classifier_.classes_ == 1).item() + p = np.clip( + np.asarray( + self.classifier_.predict_proba(X)[:, positive_class], dtype=float + ), + 0.0, + 1.0, + ) + if output_type == "mean": + return p * np.maximum(self.regressor_.predict(X, output_type="mean"), 0.0) + grid = np.asarray( + self.regressor_.predict( + X, output_type="quantiles", quantiles=self.quantile_grid_.tolist() + ), + dtype=float, + ) + levels = quantiles if output_type == "quantiles" else [0.5] + predictions = [] + for q in levels: + prediction = np.zeros(len(p)) + active = p > 1 - q + for row in np.flatnonzero(active): + level = (p[row] - (1 - q)) / p[row] + prediction[row] = np.interp(level, self.quantile_grid_, grid[:, row]) + predictions.append(np.maximum(prediction, 0.0)) + return predictions if output_type == "quantiles" else predictions[0] diff --git a/tests/test_hurdle.py b/tests/test_hurdle.py new file mode 100644 index 00000000..fad53bcb --- /dev/null +++ b/tests/test_hurdle.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import pytest +from pytest_mock import MockerFixture +from sklearn.base import clone +from sklearn.dummy import DummyClassifier, DummyRegressor +from sklearn.exceptions import NotFittedError + +from tabpfn_extensions.hurdle import AutoHurdleRegressor + + +@pytest.fixture +def model() -> AutoHurdleRegressor: + return AutoHurdleRegressor(DummyClassifier(), DummyRegressor()) + + +def test_fit_clones_and_preserves_frame( + model: AutoHurdleRegressor, mocker: MockerFixture +) -> None: + fit = mocker.spy(DummyRegressor, "fit") + X = pd.DataFrame({"value": [1, 2, 3, 4], "kind": ["a", "b", "a", "b"]}) + X.index = [7, 4, 9, 2] + y = pd.Series([0, 0, 0, 4], index=X.index) + model.fit(X, y) + assert model.hurdle_ + assert model.zero_rate_ == 0.75 + np.testing.assert_array_equal(model.classifier_.class_prior_, [0.75, 0.25]) + assert model.regressor_.constant_.item() == 4 + pd.testing.assert_frame_equal(fit.call_args.args[1], X.iloc[[3]]) + assert not hasattr(model.regressor, "constant_") + assert not hasattr(model.classifier, "classes_") + assert model.n_features_in_ == 2 + np.testing.assert_array_equal(model.feature_names_in_, X.columns) + assert clone(model).get_params()["zero_threshold"] == 0.5 + with pytest.raises(ValueError, match="feature names"): + model.predict(X[["kind", "value"]]) + + +@pytest.mark.parametrize( + ("y", "hurdle", "threshold", "expected"), + [ + ([0, 0, 1, 2], "auto", 0.5, False), + ([0, 0, 0, 2], "auto", 0.5, True), + ([0, 0, 0, -2], "auto", 0.5, False), + ([1, 2, 3, 4], "auto", 0.5, False), + ([0, 1, 2, 3], "auto", 0.1, True), + ([0, 0, 0, 2], False, 0.5, False), + ([0, 1, 2, 3], True, 0.5, True), + ], +) +def test_auto_gate( + model: AutoHurdleRegressor, + y: list[float], + hurdle: Any, + threshold: float, + expected: bool, +) -> None: + model.set_params(hurdle=hurdle, zero_threshold=threshold).fit(np.ones((4, 2)), y) + assert model.hurdle_ == expected + if not expected: + assert model.regressor_.constant_.item() == np.mean(y) + + +def test_mixture_predictions( + model: AutoHurdleRegressor, monkeypatch: pytest.MonkeyPatch +) -> None: + X = np.ones((8, 2)) + model.fit(X, [0, 0, 0, 0, 0, 1, 2, 3]) + p = np.array([0, 0.1, 0.5, 0.5001, 0.8, 1, -0.1, 1.1]) + # Reverse class order to check that probabilities follow the positive label. + model.classifier_.classes_ = np.array([1, 0]) + monkeypatch.setattr( + model.classifier_, "predict_proba", lambda _X: np.column_stack([p, 1 - p]) + ) + + def uniform_predict( + X: Any, *, output_type: str, quantiles: list[float] | None = None + ) -> Any: + if output_type == "mean": + return np.full(len(X), 15.0) + return [np.full(len(X), 10 + 10 * q) for q in quantiles] + + monkeypatch.setattr(model.regressor_, "predict", uniform_predict) + np.testing.assert_allclose(model.predict(X), [0, 0, 0, 10.1, 13.75, 15, 0, 15]) + np.testing.assert_allclose( + model.predict(X, output_type="mean"), np.clip(p, 0, 1) * 15 + ) + lower, median, upper = model.predict( + X, output_type="quantiles", quantiles=[0.1, 0.5, 0.9] + ) + np.testing.assert_allclose(median, model.predict(X)) + assert np.all(lower <= median) + assert np.all(median <= upper) + assert upper[4] == pytest.approx(18.75) + + monkeypatch.setattr( + model.regressor_, "predict", lambda X, **_kw: -np.ones((21, len(X))) + ) + np.testing.assert_array_equal(model.predict(X), 0) + + +def test_prediction_validation(model: AutoHurdleRegressor) -> None: + X = np.ones((4, 2)) + with pytest.raises(NotFittedError): + model.predict(X) + model.fit(X, np.zeros(4)) + with pytest.raises(ValueError, match="output_type"): + model.predict(X, output_type="mode") + with pytest.raises(ValueError, match="features"): + model.predict(np.ones((4, 3))) + + +def test_refit_and_all_zero(model: AutoHurdleRegressor) -> None: + X = np.ones((4, 2)) + model.fit(X, [0, 0, 0, 1]) + model.fit(X, [0, 0, 0, 0]) + assert model.classifier_ is None + assert model.regressor_ is None + np.testing.assert_array_equal(model.predict(X), np.zeros(4)) + np.testing.assert_array_equal(model.predict(X, output_type="mean"), np.zeros(4)) + np.testing.assert_array_equal( + model.predict(X, output_type="quantiles", quantiles=[0.1, 0.9]), + np.zeros((2, 4)), + ) + model.fit(X, [1, 2, 3, 4]) + assert not model.hurdle_ + assert model.classifier_ is None + with pytest.raises(ValueError): + model.fit(X, [0, 0, np.nan, 1]) + with pytest.raises(NotFittedError): + model.predict(X) + + +@pytest.mark.parametrize("y", [[0, -1, 2], [1, 2, 3], [0, np.inf, 1]]) +def test_invalid_forced_targets(model: AutoHurdleRegressor, y: list[float]) -> None: + with pytest.raises(ValueError): + model.set_params(hurdle=True).fit(np.ones((3, 2)), y) + + +@pytest.mark.parametrize("params", [{"hurdle": "yes"}, {"zero_threshold": -1}]) +def test_invalid_parameters(model: AutoHurdleRegressor, params: dict[str, Any]) -> None: + with pytest.raises(ValueError): + model.set_params(**params).fit(np.ones((4, 2)), [0, 0, 0, 1]) + + +@pytest.mark.parametrize("quantiles", [[], [0], [1], [np.nan], [-0.1]]) +def test_invalid_quantiles(model: AutoHurdleRegressor, quantiles: list[float]) -> None: + X = np.ones((4, 2)) + model.fit(X, [0, 0, 0, 0]) + with pytest.raises(ValueError, match="quantiles"): + model.predict(X, output_type="quantiles", quantiles=quantiles) + + +def test_custom_grid(model: AutoHurdleRegressor, mocker: MockerFixture) -> None: + grid = [0.2, 0.4, 0.6] + model.set_params(quantile_grid=grid).fit(np.ones((4, 2)), [0, 0, 0, 1]) + assert clone(model).quantile_grid == grid + grid[0] = 0.1 + mocker.patch.object( + model.classifier_, "predict_proba", return_value=np.array([[0.2, 0.8]]) + ) + predict = mocker.patch.object( + model.regressor_, + "predict", + return_value=[np.array([2]), np.array([4]), np.array([6])], + ) + X = np.ones((1, 2)) + np.testing.assert_allclose(model.predict(X), [3.75]) + predict.assert_called_once_with( + X, output_type="quantiles", quantiles=[0.2, 0.4, 0.6] + ) + np.testing.assert_allclose( + model.predict(X, output_type="quantiles", quantiles=[0.25, 0.9]), [[2], [6]] + ) + + +@pytest.mark.parametrize( + "grid", + [ + [], + [0.5], + [[0.1, 0.9]], + [0, 0.5], + [0.5, 1], + [0.5, np.nan], + [0.5, np.inf], + [0.5, 0.5], + [0.9, 0.1], + ], +) +def test_invalid_grid(model: AutoHurdleRegressor, grid: Any) -> None: + with pytest.raises(ValueError, match="quantile_grid"): + model.set_params(quantile_grid=grid).fit(np.ones((4, 2)), [0, 0, 0, 1]) From 6a26136c964a30d498847f03bf0dfac21bbb42a1 Mon Sep 17 00:00:00 2001 From: Adrian Hayler Date: Fri, 11 Sep 2026 17:10:13 +0200 Subject: [PATCH 2/3] Add hurdle regression example and integration test --- examples/hurdle/facebook_comments.py | 83 ++++++++++++++++++++++++++++ src/tabpfn_extensions/__init__.py | 2 + tests/test_hurdle_integration.py | 25 +++++++++ 3 files changed, 110 insertions(+) create mode 100644 examples/hurdle/facebook_comments.py create mode 100644 tests/test_hurdle_integration.py diff --git a/examples/hurdle/facebook_comments.py b/examples/hurdle/facebook_comments.py new file mode 100644 index 00000000..755c8249 --- /dev/null +++ b/examples/hurdle/facebook_comments.py @@ -0,0 +1,83 @@ +"""Compare mean and median predictions on UCI Facebook Comment Volume. + +Data: Singh, K. (2015), https://doi.org/10.24432/C5Q886 (CC BY 4.0). +The target is the number of comments in the next H hours. Training variant 1 +has 55.1% zero targets. We sample the published train and test sets separately. +Root mean squared error evaluates means. For an absolute-error objective, +the example also compares median predictions using mean absolute error. +Both comparisons include plain TabPFN and a constant training-target baseline. +Results depend on the sample and model version. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from urllib.request import urlopen +from zipfile import ZipFile + +import numpy as np +import pandas as pd +from sklearn.metrics import mean_absolute_error, root_mean_squared_error + +from tabpfn_extensions import AutoHurdleRegressor, TabPFNClassifier, TabPFNRegressor + + +def main(archive: Path, n_train: int, n_test: int, seed: int) -> None: + """Download the public data if needed, then fit and compare both models.""" + if not archive.exists(): + archive.parent.mkdir(parents=True, exist_ok=True) + with urlopen( + "https://archive.ics.uci.edu/static/public/363/" + "facebook%2Bcomment%2Bvolume%2Bdataset.zip", + timeout=60, + ) as response: + archive.write_bytes(response.read()) + with ZipFile(archive) as dataset: + train = pd.read_csv( + dataset.open("Dataset/Training/Features_Variant_1.csv"), header=None + ).sample(n=n_train, random_state=seed) + test = pd.read_csv( + dataset.open("Dataset/Testing/Features_TestSet.csv"), header=None + ).sample(n=n_test, random_state=seed) + X_train, y_train = train.iloc[:, :-1], train.iloc[:, -1] + X_test, y_test = test.iloc[:, :-1], test.iloc[:, -1] + print(f"Rows: {len(train)} train, {len(test)} test") + print( + f"Zero targets: {(y_train == 0).mean():.1%} train, {(y_test == 0).mean():.1%} test" + ) + baseline = TabPFNRegressor(n_estimators=1, random_state=seed) + hurdle = AutoHurdleRegressor( + classifier=TabPFNClassifier(n_estimators=1, random_state=seed), + regressor=baseline, + ).fit(X_train, y_train) + baseline.fit(X_train, y_train) + print(f"Hurdle enabled: {hurdle.hurdle_}") + mean_predictions = { + "Constant mean": np.full(len(test), np.mean(y_train)), + "TabPFN mean": baseline.predict(X_test, output_type="mean"), + "Hurdle mean": hurdle.predict(X_test, output_type="mean"), + } + for label, prediction in mean_predictions.items(): + print(f"{label:20s} RMSE: {root_mean_squared_error(y_test, prediction):.4f}") + median_predictions = { + "Constant median": np.full(len(test), np.median(y_train)), + "TabPFN median": baseline.predict(X_test, output_type="median"), + "Hurdle median": hurdle.predict(X_test, output_type="median"), + } + for label, prediction in median_predictions.items(): + print(f"{label:20s} MAE: {mean_absolute_error(y_test, prediction):.4f}") + + +if __name__ == "__main__": + fast = os.environ.get("FAST_TEST_MODE") == "1" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--archive", type=Path, default=Path("downloads/facebook_comments.zip") + ) + parser.add_argument("--n-train", type=int, default=100 if fast else 1000) + parser.add_argument("--n-test", type=int, default=50 if fast else 1000) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + main(args.archive, args.n_train, args.n_test, args.seed) diff --git a/src/tabpfn_extensions/__init__.py b/src/tabpfn_extensions/__init__.py index db60d761..8144ad56 100644 --- a/src/tabpfn_extensions/__init__.py +++ b/src/tabpfn_extensions/__init__.py @@ -9,6 +9,7 @@ from tabpfn_common_utils.telemetry.interactive import opt_in from .embedding import TabPFNEmbedding +from .hurdle import AutoHurdleRegressor from .many_class import ManyClassClassifier from .unsupervised import TabPFNUnsupervisedModel @@ -16,6 +17,7 @@ from .utils import TabPFNClassifier, TabPFNRegressor, is_tabpfn __all__ = [ + "AutoHurdleRegressor", "TabPFNClassifier", "TabPFNRegressor", "is_tabpfn", diff --git a/tests/test_hurdle_integration.py b/tests/test_hurdle_integration.py new file mode 100644 index 00000000..52c4704f --- /dev/null +++ b/tests/test_hurdle_integration.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +from tabpfn_extensions import AutoHurdleRegressor + + +@pytest.mark.local_compatible +@pytest.mark.client_compatible +def test_tabpfn_integration(tabpfn_classifier: Any, tabpfn_regressor: Any) -> None: + rng = np.random.default_rng(7) + X = rng.normal(size=(24, 3)) + y = np.where(X[:, 0] > 0.5, 1 + np.abs(X[:, 1]), 0) + model = AutoHurdleRegressor(tabpfn_classifier, tabpfn_regressor).fit(X, y) + prediction = model.predict(X[:4]) + assert prediction.shape == (4,) + assert np.isfinite(prediction).all() + assert (prediction >= 0).all() + model.set_params(hurdle=False).fit(X, y) + np.testing.assert_allclose( + model.predict(X[:4]), model.regressor_.predict(X[:4], output_type="median") + ) From 0bd5ae1a603ff0f3de7217444a81a81f936332a1 Mon Sep 17 00:00:00 2001 From: Adrian Hayler Date: Fri, 11 Sep 2026 17:24:22 +0200 Subject: [PATCH 3/3] Add hurdle regression changelog entry --- changelog/395.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog/395.added.md diff --git a/changelog/395.added.md b/changelog/395.added.md new file mode 100644 index 00000000..1d373da4 --- /dev/null +++ b/changelog/395.added.md @@ -0,0 +1 @@ +Add AutoHurdleRegressor for zero-inflated regression with configurable estimators and quantile grids, mean, median, and quantile predictions, and a Facebook Comment Volume example.