Regression API¶

StagecoachRegressor¶

class stagecoachml.regression.StagecoachRegressor(stage1_estimator, stage2_estimator, early_features=None, late_features=None, residual=True, use_stage1_pred_as_feature=True, inner_cv=None, random_state=None)[source]¶

Bases: StagecoachBase, RegressorMixin

Two-stage regressor for staggered feature arrival.

This estimator handles scenarios where features arrive in batches at different times. It trains a stage1 model on early features and a stage2 model that can use late features plus (optionally) the stage1 prediction.

Parameters:

stage1_estimator (BaseEstimator) – Sklearn regressor for the early features.
stage2_estimator (BaseEstimator) – Sklearn regressor for the late features (and optionally the stage1 prediction).
early_features (list[str] | None) – Column names for the early features. If None, the first half of the columns is used.
late_features (list[str] | None) – Column names for the late features. If None, the second half of the columns is used.
residual (bool) – If True, stage2 predicts y - stage1_pred; if False, it predicts y directly.
use_stage1_pred_as_feature (bool) – If True, the stage1 prediction is included as an input to stage2.
inner_cv (int | None) – Number of folds for cross-fitting the stage1 predictions during training. Helps avoid overfitting when the stage1 prediction is used as a stage2 feature.
random_state (int | None) – Random state for reproducibility.

stage1_estimator_¶

Fitted stage1 estimator.

Type:: Any

stage2_estimator_¶

Fitted stage2 estimator.

Type:: Any

__init__(stage1_estimator, stage2_estimator, early_features=None, late_features=None, residual=True, use_stage1_pred_as_feature=True, inner_cv=None, random_state=None)[source]¶

fit(X, y, sample_weight=None)[source]¶

Fit the two-stage regressor.

Parameters:

X (ndarray | DataFrame) – Training data of shape (n_samples, n_features).
y (ndarray | Series) – Target values of shape (n_samples,).
sample_weight (ndarray | None) – Per-sample weights of shape (n_samples,).

Returns:

The fitted estimator.

Return type:

StagecoachRegressor

predict_stage1(X)[source]¶

Predict using only early features (stage1).

Parameters:: X (ndarray | DataFrame) – Input data of shape (n_samples, n_features).
Returns:: Stage1 predictions of shape (n_samples,).
Return type:: ndarray

predict(X)[source]¶

Predict using both stages (full prediction).

Parameters:: X (ndarray | DataFrame) – Input data of shape (n_samples, n_features).
Returns:: Final predictions of shape (n_samples,).
Return type:: ndarray

set_fit_request(*, sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters¶

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in fit.

Returns¶

selfobject: The updated object.

set_score_request(*, sample_weight='$UNCHANGED$')¶

Configure whether metadata should be requested to be passed to the score method.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters¶

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns¶

selfobject: The updated object.

Usage Examples¶

Basic Usage¶

from stagecoachml import StagecoachRegressor
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor

# Load data
diabetes = load_diabetes(as_frame=True)
X = diabetes.frame.drop(columns=["target"])
y = diabetes.frame["target"]

# Split features
features = list(X.columns)
mid = len(features) // 2
early_features = features[:mid]
late_features = features[mid:]

# Create model
model = StagecoachRegressor(
    stage1_estimator=LinearRegression(),
    stage2_estimator=RandomForestRegressor(),
    early_features=early_features,
    late_features=late_features,
    residual=True,
    use_stage1_pred_as_feature=True,
)

# Train and predict
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
model.fit(X_train, y_train)

# Get stage-1 predictions (early features only)
stage1_pred = model.predict_stage1(X_test)

# Get final predictions (all features)
final_pred = model.predict(X_test)