causalis.shared.rct_design.enrichment

Outcome-model enrichment with separate training and effect-evaluation users.

Module Contents

Classes

EnrichmentDiagnostics	Counts, assignment checks, and balance for a frozen enrichment rule.
PrognosticEnrichment	Learn a binary-outcome selection rule on train; evaluate on held-out users.

API

class causalis.shared.rct_design.enrichment.EnrichmentDiagnostics

Counts, assignment checks, and balance for a frozen enrichment rule.

counts reports users and conversions separately for each arm and sample. SRM uses the caller’s planned allocation, never an observed allocation. srm_selected is None when no evaluation users meet the threshold. balance includes all original covariates and the predicted outcome score, even covariates that are constant within the selected sample. Diagnostics do not establish that covariates were recorded before treatment or that the evaluation outcomes have never been used for model/rule selection.

counts: pandas.DataFrame

None

srm_full: causalis.shared.srm.SRMResult

None

srm_evaluation: causalis.shared.srm.SRMResult

None

srm_selected: causalis.shared.srm.SRMResult | None

None

balance: pandas.DataFrame

None

threshold: float

None

selected_fraction: float

None

warnings: tuple[str, ...]

None

class causalis.shared.rct_design.enrichment.PrognosticEnrichment(estimator: Any, *, training_source: str = 'treatment', train_fraction: float = 0.3, top_fraction: float | None = 0.2, threshold: float | None = None, split_seed: int = 42)

Learn a binary-outcome selection rule on train; evaluate on held-out users.

A deterministic hash of user_id and split_seed splits the entire experiment independently of treatment labels and outcomes. The estimator is cloned and fitted only on the specified arm of the training sample. One threshold is applied to both evaluation arms. No training user is returned by :meth:select.

Parameters

estimator : sklearn-compatible classifier Must implement fit, predict_proba, and fitted classes_. Pipelines are supported. The original estimator is never fitted. Set the estimator’s own random seed for reproducible model fitting. training_source : {“treatment”, “control”}, default “treatment” Treatment training learns :math:P(Y(1)=1\mid X), while control training learns :math:P(Y(0)=1\mid X). Neither is an uplift prediction. train_fraction : float, default 0.3 Probability of assignment to train; actual sample sizes may differ. top_fraction : float or None, default 0.2 Set the threshold to the 1 - top_fraction quantile of predictions over BOTH training arms. This is a training reference, not a separate per-arm ranking. Threshold ties are included, so the selected share may exceed this fraction. The evaluation share may also differ under drift. Set to None when supplying an explicit threshold. threshold : float or None, default None Fixed probability cutoff in [0, 1], as an alternative to top_fraction. split_seed : int, default 42 Controls the independent analysis split, not the experimental assignment or the estimator’s randomness. Do not search seeds using evaluation results.

Notes

Features must be measured before experimental assignment. This cannot be verified from column names. All preprocessing must be fitted within the estimator/pipeline on train; do not pass preprocessing fitted on evaluation. Evaluation outcomes must not be used to tune models, cutoffs, sample splits, or balance/SRM. Split independence assumes independently randomized users; cluster/block randomization needs a design-aware split and inference.

The target effect is conditional on the learned selected population, not the original audience. SRM and covariate balance hold in expectation under ordinary randomization, not necessarily in every realized sample. Selecting on a prognostic score does not guarantee improved power.

Examples

from sklearn.linear_model import LogisticRegression rule = PrognosticEnrichment(LogisticRegression(max_iter=1000)) rule.fit(data, features=[“prior_activity”, “prior_spend”]) selected = rule.select() diagnostics = rule.diagnostics(target_allocation={0: 0.5, 1: 0.5}) effect = DiffInMeans().fit(selected).estimate(method=”conversion_ztest”)

Initialization

fit(data: causalis.data_contracts.CausalData | causalis.data_contracts.RctCausalData, *, features: Sequence[str], outcome: str | None = None) → causalis.shared.rct_design.enrichment.PrognosticEnrichment

Freeze a model and threshold without using evaluation outcomes.

features must explicitly select declared pre-treatment confounders. A unique, non-missing user_id is required. For RctCausalData with multiple outcomes, select one outcome explicitly. Only a single binary treatment column (0=control, 1=treatment) is supported.

get_partition() → pandas.DataFrame

Return user IDs, train/evaluation labels, scores, and selection flags.

The ID column is always named user_id, regardless of the input name. The returned copy preserves input row order and index. Assignment and selection do not depend on input row order. selected is always False for training users. Mutating the returned frame cannot alter the rule.

get_selected_df() → pandas.DataFrame

Return selected evaluation rows, including when no inference is possible.

select() → causalis.data_contracts.CausalData

Return selected evaluation users for existing RCT estimators.

Constant confounders in the segment are omitted to satisfy CausalData; diagnostic balance still includes them. Both arms and a non-constant outcome must be present. With no events, use diagnostics/get_selected_df to inspect the segment rather than silently changing the cutoff.

diagnostics(*, target_allocation: dict[int, float], alpha: float = 0.001) → causalis.shared.rct_design.enrichment.EnrichmentDiagnostics

Check the frozen rule using the original planned binary allocation.

This does not adjust the cutoff or repair a failing check. Empty or single-arm selections remain inspectable. Few successes/non-successes are flagged because ordinary conversion z-tests use an approximation.