DoubleML / DoubleML/doubleml-for-py

[Feature Request]:

Open
#403 3 comments 1 reaction 0 assignees View on GitHub
enhancement new feature
Dominant language
Python
Stars
786
Forks
128
Avg merge
12h 2m
Merged PRs (30d)
1

Description

### Describe the feature you want to propose or implement @SvenKlaassen @PhilippBach

I would like to contribute a new model class for **sample selection models estimated with a control function (Heckman-type) correction**, where the target parameter is the **sample selection bias coefficient** itself.

**Model.** An outcome equation and a participation (selection) equation,

```
Y_i = g_0(X_i) + eps_i (Y_i observed only if D_i = 1)
D_i = 1{ Z_i'beta_0 + v_i > 0 }, Z_i = (X_i, U_i)
```

with jointly normal `(eps_i, v_i)`. `X_i` are the covariates of the outcome equation and `U_i` are additional variables that shift participation but are excluded from the outcome equation. Under joint normality, `E[eps_i | X_i, Z_i, D_i = 1] = theta_0 * h_i` with the inverse Mills ratio `h_i = phi(Z_i'beta_0) / Phi(Z_i'beta_0)`, so the observed outcome equation becomes

```
Y_i = g_0(X_i) + theta_0 * h_i + u_i, E[u_i | X_i] = 0
```

and the parameter of interest is `theta_0 = **sigma_{eps,v}`.**

**Why this is useful.** `theta_0` determines whether outcomes observed in a self-selected subsample can be used to say anything about the units that do not report. If `theta_0 != 0`, imputing outcomes for non-reporters from a model fitted on reporters is systematically biased. The class provides both a sqrt(n)-consistent estimate of the bias and a test of its absence. Typical applications are voluntary corporate disclosure (our motivation is carbon emissions disclosure), survey non-response, and wage equations with participation decisions — settings where selection is on *unobservables* and the researcher has covariates that shift participation but not the outcome.

**Relation to the existing `DoubleMLSSM`.** These are different estimands with different identifying assumptions, not competing implementations of the same thing. `DoubleMLSSM` targets an ATE under outcome attrition using inverse probability weighting, under MAR or under nonignorable nonresponse with an instrument. The proposed class targets the selection bias coefficient in a *structural* outcome equation via a control function, and identifies coefficients of the outcome equation under selection on unobservables. In practice they answer different questions and would typically be used by different literatures.

### Propose a possible solution or implementation

A new class following the standard OOP structure — only the nuisance functions and the score components are specified, everything else inherited from the abstract base class.

**Linear Neyman-orthogonal score**, evaluated on the selected subpopulation `{D = 1}`:

```
psi(W; theta, eta) = (Y - g(X) - theta*h) * (h - m(X))
psi_a = -h * (h - m)
psi_b = (Y - g) * (h - m)
```

**Nuisance functions:**

| learner | type | target |
| --- | --- | --- |
| `ml_pi` | classifier | `pi(Z) = P(D = 1 \| Z)` |
| `ml_g` | regressor | `g(X) = E[Y \| X, D = 1]` |
| `ml_m` | regressor | `m(X) = E[h \| X, D = 1]` |

Three implementation points that I think are the interesting ones:

**(a) The inverse Mills ratio is link free.** Since `pi_i = Phi(Z_i'beta_0)`, the index is recoverable from the propensity score through the probit link, giving `h_i = phi(Phi^{-1}(pi_i)) / pi_i`. Under joint normality this is exactly the inverse Mills ratio, so *any* probabilistic classifier can be used for the selection equation and a probit is just the special case. This is what makes a control function correction usable inside the ML-learner interface at all. Note that `h` diverges as `pi -> 0`, so strict overlap is a genuine identification requirement here rather than only numerical hygiene; the class uses the standard `PSProcessorConfig` / `init_ps_processor` machinery for trimming.

**(b) Nested cross-fitting.** `h_hat` is a *generated* regressor that is simultaneously the *target* of `m`. Each training fold is therefore split in half: the first half estimates the selection equation, the second half (restricted to `D = 1`) estimates `g` and `m`. This is the same pattern as the `nonignorable` score of `DoubleMLSSM`. Exposed as `nested_cross_fitting=True` (default), which can be switched off.

**(c) Zero padding gives the effective sample size for free.** The score elements are set to zero for `D_i = 0`. With `n = sum_i D_i` observed outcomes, the base class variance estimator `mean(psi^2) / J^2 / N` then reduces algebraically to `sigma^2_theta / n` with `-J_0 = E_n[(h - m)^2 | D = 1]`. No override of the variance estimation is needed. There is a unit test asserting this identity against a manual implementation, in both languages.

**Status of the implementation.** Working and tested against `DoubleML 0.11.4`, with an R6 twin against `DoubleML 1.0.2`:

- 60 passing pytest tests. The core tests compare the coefficient, standard error, both score elements, all four cross-fitted nuisance vectors and the score-test statistic against a manual line-by-line implementation of the algorithm, across `n_folds in {3, 5}` and with/without nested cross-fitting. Plus the zero-padding variance identity, a `theta_0 = 0` centering check, and a full exception battery.
- `black --check` clean, `ruff` clean.
- Monte Carlo evidence that the orthogonalization is doing the work: bias falls as `-0.208 -> -0.108 -> -0.056` for `N = 1000, 2000, 4000` (roughly `O(1/N)`, consistent with the second-order product-of-errors term), while the naive two-step plug-in sits at `+0.33, +0.31, +0.30` and does not improve. At `theta_0 = 0` the estimator is centred (bias `-0.004`) and the score test has empirical size `0.08` against a nominal 5%, with power `0.84` at `theta_0 = -0.25` and `1.00` at `-0.50`.
- An executed example-gallery notebook in the style of `py_double_ml_ssm.ipynb`, built entirely on a simulated DGP (see the last field).

**Open questions I'd like your view on before writing PR code:**

1. **Naming.** I have been using `DoubleMLSSCF` (Sample Selection, Control Function). I'm not attached to it, and `SSCF` vs the existing `SSM` may be easy to misread in review. `DoubleMLHSM` (Heckman selection model) or `DoubleMLSSB` (sample selection bias) are alternatives. Your call.
2. **Data backend.** Currently plain `DoubleMLData` with `d_cols` holding the selection indicator, `x_cols` the outcome covariates and `z_cols` the exclusion restrictions. This needs no new data class, but it overloads "treatment" to mean "selection indicator". The alternative is extending `DoubleMLSSMData` to allow `s_col` with no treatment variable, which is semantically honest but touches an existing class. Which do you prefer?
3. **Scope.** Should the two extra methods above stay on the class or move to the example gallery?
4. **Sensitivity analysis.** Currently `_sensitivity_implemented = False`. I'm genuinely unsure whether a Chernozhukov-style sensitivity analysis is well defined for this parameter and would welcome a view.
5. **Tuning.** `_nuisance_tuning` is implemented; `_nuisance_tuning_optuna` raises `NotImplementedError`. Is that acceptable given the nested structure, as with `DoubleMLSSM`?
6. **R twin.** Open the `doubleml-for-r` PR at the same time, or after the Python class is settled?

### Did you consider alternatives to the proposed solution. If yes, please describe

Yes, four.

**Reusing `DoubleMLSSM` with a new score string.** Rejected. The nuisance set is different (this model needs the projection `m(X) = E[h | X, D = 1]`, which has no counterpart in an IPW score), the estimand is a structural coefficient rather than an ATE, and the identifying assumption is joint normality of the error terms rather than a conditional independence or instrument condition. Bolting a third score onto `DoubleMLSSM` would make that class harder to reason about for the model it already serves.

**A callable score on `DoubleMLPLR`.** Technically possible — `h` can be treated as an endogenous-looking regressor and the score is linear. Rejected because the user would have to construct the cross-fitted inverse Mills ratio themselves, get the nested cross-fitting right by hand, and remember to zero-pad off `{D = 1}`. Those are precisely the three things that are easy to get wrong, and they are the reason a dedicated class is worth having.

**Estimating `theta_0` by naive two-step Heckman with ML first stages.** This is the obvious baseline and it is what the class exists to replace. It is *inconsistent* here: without partialling out `E[h | X]`, the regularization bias of the high-dimensional first stages contaminates the estimate. In the simulation it is attenuated to less than half the true value and, unlike the orthogonal estimator, does not improve with sample size (bias `+0.33, +0.31, +0.30` at `N = 1000, 2000, 4000`). It also has a *tighter* sampling distribution than the orthogonal estimator — tightly wrong — which is a good cautionary illustration for the docs.

**A new data class vs. reusing `DoubleMLData`.** Considered both; this is open question 2 above. I implemented the lighter option so that the contribution touches as little existing code as possible, but I don't think it is obviously the right one.

**On scope:** the paper also contains an adaptive kernel group lasso for post-selection estimation of the outcome equation. I deliberately propose to **leave it out of this contribution**. It is a variable selection device applied *after* `theta_hat` is available, not part of the orthogonal score, and bundling it would roughly double the review surface for no benefit to the model class. I'd ship it with the example notebook and propose it separately only if there's appetite.

### Comments, context or references

**Data availability.** The method is developed in a working paper (Chen, Lioui and Scaillet, *Green Silence: Double machine learning carbon emissions under sample selection bias*). The empirical data are licensed from a commercial vendor and cannot be redistributed, so the contribution ships a simulated DGP and the example notebook is built entirely on it — in the same spirit as the existing `py_double_ml_ssm` example, which uses the Bia, Huber and Lafférs design. The DGP generator (characteristics with Toeplitz covariance, Gaussian kernel features on characteristic-sorted portfolio grids, group-sparse coefficients, an active set in the selection equation disjoint from the outcome active set, and an intercept calibrated to a target participation rate) would be contributed alongside the class.

**Compatibility.** Developed against `DoubleML 0.11.4` (Python) and `1.0.2` (R). Happy to rebase onto `main`.

**Proposed file layout**, if the design is agreed (sample selection models under `irm/`, next to `ssm.py`):

```
doubleml/irm/sscf.py # model class
doubleml/irm/datasets/dgp_green_silence_data.py
doubleml/irm/tests/test_sscf.py
doubleml/irm/tests/test_sscf_exceptions.py
doubleml/irm/tests/_utils_sscf_manual.py # manual reference implementation
```

plus exports and API doc stubs; the notebook goes to `doubleml-docs` as a separate PR.

**References**

- Heckman, J. J. (1979), Sample selection bias as a specification error, *Econometrica* 47(1), 153-161.
- Chernozhukov, Chetverikov, Demirer, Duflo, Hansen, Newey and Robins (2018), Double/debiased machine learning for treatment and structural parameters, *The Econometrics Journal* 21(1), C1-C68.
- Bia, Huber and Lafferty (2023), Double machine learning for sample selection models, *JBES* — the basis of the existing `DoubleMLSSM`, included here for contrast.
- Yuan and Lin (2006), Model selection and estimation in regression with grouped variables, *JRSS-B* 68(1), 49-67 — background for the post-selection step, not part of this proposal.

I'm happy to open a draft PR straight away if the design looks reasonable, or to adjust it first based on your answers to the questions in the previous field.

Contributor guide

Open the contributing guide

Research direction

Start by reading doubleml/irm/ssm.py and the proposed layout for doubleml/irm/sscf.py, then review the open questions about naming, data handling, scope, and tuning. Run doubleml/irm/tests/test_sscf.py and doubleml/irm/tests/test_sscf_exceptions.py, comparing results with _utils_sscf_manual.py. Done means the design decisions are settled, the Python class and exports are integrated, and the listed tests pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, scikit-learn
Domain
backend-api-design, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.