Add support for user-supplied propensities in DML and DRLearner estimators

[Oren-T]

Motivation

If you have data from a randomized experiment with a known assignment design — say subjects randomized within blocks, where the true propensity is just the known treatment probability within each block — there's currently no way to tell the DML or DR learners about it. The estimators always fit model_t / model_propensity internally, and the cross-fitting machinery makes it effectively impossible to inject precomputed values from outside: models are cloned per fold and only ever see anonymous row slices, so even a custom "model" holding the known values can't align them to the rows it's asked to score.

This has come up repeatedly: #730 asks for exactly this for ForestDRLearner, #535 for A/B test data, and #646 for reusable first-stage models more generally. The workarounds suggested in those threads (a DummyClassifier for constant probabilities, or appending the propensity as the last column of W with a passthrough model) do work in narrow cases, but they're fragile and roundabout — and notably, the closure-over-an-array variant some users reach for is silently wrong under cross-fitting and bootstrap resampling, since the wrapper has no way to know which rows it's being asked about. In the block-randomized case, fitting a model to estimate something you know by design is strictly worse than just using the known values.

What this adds

A propensity keyword on fit (and score) of the DML family (when discrete_treatment=True) and the DRLearner family:

est = LinearDML(discrete_treatment=True)
est.fit(y, T, X=X, W=W, propensity=p)  # p = known per-unit assignment probability, e.g. the within-block rate

When provided, the first-stage treatment/propensity model is not fitted at all; the supplied values are used directly in residualization (DML: T_res = T_onehot - p[:, 1:]) or in the doubly robust correction (DRLearner, where they remain subject to min_propensity clipping and trimming_threshold trimming). Binary treatments accept a single column (the probability of the non-control treatment); multi-valued treatments take one column per category, ordered to match the fitted categories.

One point the docs are careful about: the supplied values must be the true by-design probabilities conditional on everything that influenced assignment — including design variables like randomization blocks even when those aren't part of X or W. That conditioning set is what identification actually requires (and it's what makes this feature useful in the first place: in a block design where the blocks aren't recoverable from X/W, no fitted propensity model can be correct, but the known per-unit rate is).

The implementation follows the existing pattern for per-row fit arguments (sample_var, freq_weight, Z): the array is validated in _OrthoLearner.fit and plumbed through _fit_nuisances into _crossfit, whose kwargs mechanism already slices per fold. As a result cross-fitting, mc_iters, cache_values/refit_final, groups, and bootstrap inference all compose with it without special handling. Estimators that don't model propensities are unaffected — the kwarg defaults to None everywhere and is rejected where unsupported. In score signatures the new parameter is placed last, so existing positional callers are unaffected (there's a test pinning this).

A few behavioral notes:

• Since no propensity model is fitted, score on new data also requires the propensity argument (a clear error message explains this).
• models_t / models_propensity raise an informative error after a bypassed fit instead of returning never-fitted models; nuisance_scores_t / nuisance_scores_propensity entries are None.
• Supplied values of exactly 0 or 1 trigger a warning (no overlap, or accidentally swapped columns).

Testing

New test module econml/tests/test_user_propensity.py (22 tests), including: a hidden-block RCT where the supplied propensities recover the ATE with covering confidence intervals while a fitted propensity model is badly biased; exact treatment-residual identities; proof the first-stage model is never fitted (via a classifier whose fit raises); multi-valued treatments with unit-varying propensities; category-ordering checks with a swapped-column negative control; clipping/trimming interaction via n_samples_trimmed_; composition with groups, freq_weight/sample_var, mc_iters, refit_final, and bootstrap inference; positional-argument compatibility of score; and input validation. The existing test_dml, test_drlearner, test_ortho_learner, test_refit, and test_inference suites all pass. The doc FAQ snippets added to dml.rst / dr.rst were executed against the testsetup data.

Natural follow-ups not included here to keep the change contained: accepting propensity in CausalForestDML.tune(), RScorer, and the policy learners.

Fixes #730. Also addresses the propensity side of #535 and #646.