metacausal.aggregation.fit_nuisance

metacausal.aggregation.fit_nuisance(X, T, Y, fold_spec, propensity_model=None, outcome_model=None, propensity_trim=0.01, random_state=None, outcome_type=None)[source]

Fit nuisance models via cross-fitting and return out-of-fold predictions.

For each fold j in fold_spec:

  1. Clone and fit propensity model on the training indices.

  2. Clone and fit outcome model on treated training units -> mu1 on test indices.

  3. Clone and fit outcome model on control training units -> mu0 on test indices.

  4. Clip propensity to [propensity_trim, 1 - propensity_trim].

For CrossFitSplit, all n positions are filled. For TrainAvgSplit, only fold_spec.test_indices[0] positions are filled; the rest remain np.nan.

Parameters:

  • X (array of shape (n, p))

  • T (array of shape (n,)) – Binary treatment assignment (0 or 1).

  • Y (array of shape (n,)) – Observed outcome.

  • fold_spec (FoldSpec) – Output of CrossFitSplit.split() or TrainAvgSplit.split().

  • propensity_model (sklearn classifier or None) – Model for P(T=1|X). Must support predict_proba. Default: HistGradientBoostingClassifier.

  • outcome_model (sklearn regressor / classifier or None) – Model for E[Y|X, T]. Cloned separately for treated and control fits. For continuous Y, must be a regressor; mu_hat is read from predict(). For binary Y, must be a classifier with predict_proba; mu_hat is read from predict_proba(X)[:, 1]. None selects an outcome-type-appropriate default (HistGradientBoostingRegressor or HistGradientBoostingClassifier).

  • propensity_trim (float) – Clip propensity scores to [trim, 1-trim] to enforce overlap.

  • outcome_type ("continuous", "binary", or None) – How to interpret Y. None (default) auto-detects from the value set of Y via metacausal.infer_outcome_type().

Returns:

NuisanceEstimates with out-of-fold predictions. e_hat is already trimmed.

Return type:

NuisanceEstimates