oi-workforce-dml 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- oi_workforce_dml-0.2.0/MANIFEST.in +19 -0
- oi_workforce_dml-0.2.0/PKG-INFO +298 -0
- oi_workforce_dml-0.2.0/README.md +262 -0
- oi_workforce_dml-0.2.0/dml_code/__init__.py +5 -0
- oi_workforce_dml-0.2.0/dml_code/cli.py +17 -0
- oi_workforce_dml-0.2.0/dml_code/estimation.py +128 -0
- oi_workforce_dml-0.2.0/dml_code/src/__init__.py +1 -0
- oi_workforce_dml-0.2.0/dml_code/src/config.py +337 -0
- oi_workforce_dml-0.2.0/dml_code/src/estimate_helpers.py +292 -0
- oi_workforce_dml-0.2.0/dml_code/src/nuisance.py +345 -0
- oi_workforce_dml-0.2.0/dml_code/src/paths.py +45 -0
- oi_workforce_dml-0.2.0/dml_code/src/plotting.py +1438 -0
- oi_workforce_dml-0.2.0/dml_code/src/preprocessing.py +115 -0
- oi_workforce_dml-0.2.0/dml_code/src/results.py +205 -0
- oi_workforce_dml-0.2.0/dml_code/src/utils.py +184 -0
- oi_workforce_dml-0.2.0/example_set_up/experiments/experiment.yaml +13 -0
- oi_workforce_dml-0.2.0/example_set_up/registries/__init__.py +1 -0
- oi_workforce_dml-0.2.0/example_set_up/registries/covariate_sets.py +22 -0
- oi_workforce_dml-0.2.0/example_set_up/registries/filter_sets.py +66 -0
- oi_workforce_dml-0.2.0/example_set_up/registries/models.py +61 -0
- oi_workforce_dml-0.2.0/example_set_up/registries/programs.py +11 -0
- oi_workforce_dml-0.2.0/oi_workforce_dml.egg-info/SOURCES.txt +21 -0
- oi_workforce_dml-0.2.0/pyproject.toml +69 -0
- oi_workforce_dml-0.2.0/setup.cfg +4 -0
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
include README.md
|
|
2
|
+
include pyproject.toml
|
|
3
|
+
|
|
4
|
+
recursive-include dml_code *.py
|
|
5
|
+
recursive-include example_set_up *.py *.yaml
|
|
6
|
+
|
|
7
|
+
prune data
|
|
8
|
+
prune outputs
|
|
9
|
+
prune project_scripts
|
|
10
|
+
prune build
|
|
11
|
+
prune dist
|
|
12
|
+
prune *.egg-info
|
|
13
|
+
|
|
14
|
+
global-exclude __pycache__
|
|
15
|
+
global-exclude *.py[cod]
|
|
16
|
+
global-exclude .DS_Store
|
|
17
|
+
|
|
18
|
+
# Private planning notes are never distributed.
|
|
19
|
+
exclude TODO.md
|
|
@@ -0,0 +1,298 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: oi-workforce-dml
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Cross-fitted estimation and plotting for program evaluation.
|
|
5
|
+
Author: DML Pipeline Contributors
|
|
6
|
+
Project-URL: Repository, https://github.com/coreygb1/dml-dev
|
|
7
|
+
Keywords: administrative-data,causal-inference,doubleml,observational-data,program-evaluation
|
|
8
|
+
Classifier: Development Status :: 3 - Alpha
|
|
9
|
+
Classifier: Intended Audience :: Science/Research
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering
|
|
14
|
+
Requires-Python: >=3.12
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
Requires-Dist: doubleml
|
|
17
|
+
Requires-Dist: joblib
|
|
18
|
+
Requires-Dist: lightgbm
|
|
19
|
+
Requires-Dist: optuna
|
|
20
|
+
Requires-Dist: oi-tools[figures]
|
|
21
|
+
Requires-Dist: plotnine
|
|
22
|
+
Requires-Dist: polars
|
|
23
|
+
Requires-Dist: PyYAML
|
|
24
|
+
Requires-Dist: numpy
|
|
25
|
+
Requires-Dist: pandas
|
|
26
|
+
Requires-Dist: matplotlib
|
|
27
|
+
Requires-Dist: xlsxwriter
|
|
28
|
+
Requires-Dist: pyarrow
|
|
29
|
+
Requires-Dist: psutil
|
|
30
|
+
Requires-Dist: scikit-learn<1.10,>=1.2
|
|
31
|
+
Requires-Dist: shap
|
|
32
|
+
Requires-Dist: threadpoolctl
|
|
33
|
+
Provides-Extra: dev
|
|
34
|
+
Requires-Dist: build; extra == "dev"
|
|
35
|
+
Requires-Dist: twine; extra == "dev"
|
|
36
|
+
|
|
37
|
+
# OI Workforce DML
|
|
38
|
+
|
|
39
|
+
Cross-fitted outcome, IPW, and AIPW estimation with saved models and diagnostics.
|
|
40
|
+
Supply prebuilt data using the input contract below.
|
|
41
|
+
|
|
42
|
+
## Version 0.2 migration
|
|
43
|
+
|
|
44
|
+
Version 0.2 runs estimation on prebuilt inputs. It no longer includes the old
|
|
45
|
+
build command or bundled project configuration. Copy `example_set_up/` from the
|
|
46
|
+
source distribution or repository, fill in the registries, and pass the YAML
|
|
47
|
+
experiment path to `dml-estimate`.
|
|
48
|
+
|
|
49
|
+
Replace `tune_nuisance_models` / `n_tuning_trials` with independent
|
|
50
|
+
`outcomes_tune` / `outcomes_tuning_trials` and
|
|
51
|
+
`propensity_tune` / `propensity_tuning_trials`. Remove `drop_nulls` and
|
|
52
|
+
`calculate_on_control_placebo`; covariate missingness is handled by the models,
|
|
53
|
+
and RCT-control evaluation follows the program registry. Unknown settings are
|
|
54
|
+
rejected. The examples below use the current format.
|
|
55
|
+
|
|
56
|
+
## Install and run
|
|
57
|
+
|
|
58
|
+
The PyPI distribution is `oi-workforce-dml`; the Python import remains
|
|
59
|
+
`dml_code`, and the command remains `dml-estimate`. The GitHub repository is
|
|
60
|
+
`dml-dev`. If the old `dml-dev` distribution is installed in this environment,
|
|
61
|
+
uninstall it first (`python -m pip uninstall dml-dev`): both distributions
|
|
62
|
+
provide the same import package and command.
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
# Install the published package:
|
|
66
|
+
python -m pip install oi-workforce-dml==0.2.0
|
|
67
|
+
# Or install a local checkout:
|
|
68
|
+
python -m pip install /path/to/dml-dev
|
|
69
|
+
# Or for development: python -m pip install -e /path/to/dml-dev
|
|
70
|
+
|
|
71
|
+
dml-estimate /path/to/my_analysis/experiment.yaml
|
|
72
|
+
# Equivalent:
|
|
73
|
+
python -m dml_code.cli /path/to/my_analysis/experiment.yaml
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from dml_code import estimate
|
|
78
|
+
|
|
79
|
+
estimate("/path/to/my_analysis/experiment.yaml")
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
The experiment name comes from the YAML filename (`experiment` above).
|
|
83
|
+
The command can run from any working directory. YAML is parsed safely; Python
|
|
84
|
+
experiment configurations are no longer accepted. `example_set_up/` is a
|
|
85
|
+
copyable checkout example, outside the installed package. It contains sibling
|
|
86
|
+
`experiments/` and `registries/` directories; supply your own processed inputs.
|
|
87
|
+
|
|
88
|
+
## YAML experiment configuration
|
|
89
|
+
|
|
90
|
+
```yaml
|
|
91
|
+
registry_dir: ../registries
|
|
92
|
+
output_dir: ../results
|
|
93
|
+
experiments:
|
|
94
|
+
- program_name: [my_program]
|
|
95
|
+
outcomes: [wages_2]
|
|
96
|
+
covariate_set_name: [baseline]
|
|
97
|
+
filter_set_name: [eligible]
|
|
98
|
+
outcomes_model_name: [lightgbm]
|
|
99
|
+
outcomes_tune: [false]
|
|
100
|
+
propensity_model_name: [lightgbm]
|
|
101
|
+
propensity_tune: [false]
|
|
102
|
+
group_col: [site]
|
|
103
|
+
delete_raw_predictions: [false]
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
`registry_dir` and `output_dir` resolve relative to the YAML file, or can be
|
|
107
|
+
absolute. The registry directory contains these fixed Python files:
|
|
108
|
+
|
|
109
|
+
| File | Dictionary |
|
|
110
|
+
| --- | --- |
|
|
111
|
+
| `programs.py` | `PROGRAM_REGISTRY` |
|
|
112
|
+
| `covariate_sets.py` | `COVARIATE_SET_REGISTRY` |
|
|
113
|
+
| `filter_sets.py` | `FILTER_SET_REGISTRY` |
|
|
114
|
+
| `models.py` | `OUTCOMES_MODEL_REGISTRY`, `PROPENSITY_MODEL_REGISTRY` |
|
|
115
|
+
|
|
116
|
+
Registry dictionaries map names to program definitions, covariate lists,
|
|
117
|
+
lists of Polars filter functions, and supported sklearn/LightGBM model instances. They remain
|
|
118
|
+
executable Python; load registry files you trust. Bare sibling imports,
|
|
119
|
+
package imports such as `from registries.models import ...`, and relative
|
|
120
|
+
imports such as `from .models import ...` are supported. Sibling modules are
|
|
121
|
+
isolated between configs loaded sequentially in one Python process.
|
|
122
|
+
|
|
123
|
+
Experiment fields remain lists and expand into combinations. Multiple blocks
|
|
124
|
+
are allowed. Model templates are cloned for each run/fold. `dml_type` is not
|
|
125
|
+
accepted. Unknown top-level and experiment settings are rejected. Optional defaults:
|
|
126
|
+
no tuning, no grouping (`[]` or `[null]`), and retain predictions.
|
|
127
|
+
Missing covariates are retained; rows missing identity, treatment, or outcome are dropped.
|
|
128
|
+
|
|
129
|
+
Set `outcomes_tune: [true]` with `outcomes_tuning_trials: [15]` to tune outcomes;
|
|
130
|
+
use `propensity_tune` and `propensity_tuning_trials` independently for propensity.
|
|
131
|
+
False uses the registry preset. True uses the built-in tuning recipe before
|
|
132
|
+
cross-fitting. Omit a learner's trial budget when it is not tuned. For a mixed
|
|
133
|
+
`[false, true]` grid, budgets expand tuned runs only; the preset runs once.
|
|
134
|
+
Grouping generates a pooled run plus each non-null group with at least
|
|
135
|
+
30 treated rows in the input. Each group uses its treated rows and all controls.
|
|
136
|
+
Filters select rows only: do not transform covariates, because plotting later
|
|
137
|
+
reloads their values from the source file.
|
|
138
|
+
|
|
139
|
+
## Model presets and tuning
|
|
140
|
+
|
|
141
|
+
The model registries contain ordinary model instances. Supported families are
|
|
142
|
+
LightGBM, Lasso (outcomes) / LogisticRegression (propensity), random forests,
|
|
143
|
+
MLP neural networks, and VotingRegressor / soft VotingClassifier ensembles.
|
|
144
|
+
The example model registry contains a preset for each. Ensemble members use
|
|
145
|
+
fixed presets and weights; requesting ensemble tuning is an error.
|
|
146
|
+
|
|
147
|
+
Each learner has its own `*_tune` switch and `*_tuning_trials` budget. False
|
|
148
|
+
uses the registry preset. True clones that preset and overrides the parameters
|
|
149
|
+
in the built-in tuning recipe, retaining parameters outside that recipe. Tuning
|
|
150
|
+
happens once before cross-fitting, using the existing 80/20 split, outcome MSE
|
|
151
|
+
on controls, and binary log loss for propensity. The best LightGBM iteration
|
|
152
|
+
count is retained after early stopping. Search spaces live in `nuisance.py`:
|
|
153
|
+
|
|
154
|
+
| Family | Outcomes | Propensity |
|
|
155
|
+
| --- | --- | --- |
|
|
156
|
+
| LightGBM | Existing outcome tree search | Existing constrained propensity tree search |
|
|
157
|
+
| Linear | Lasso regularization | L1 logistic regularization, favoring stronger penalties |
|
|
158
|
+
| Forest | Deeper trees, smaller leaves | Shallower trees, larger leaves |
|
|
159
|
+
| Neural net | Larger hidden layers, lighter regularization | Smaller layers, stronger regularization |
|
|
160
|
+
|
|
161
|
+
These are starting ranges, not guarantees of calibration or adequate overlap.
|
|
162
|
+
A small trial budget explores only a few parameters. Manual presets remain the
|
|
163
|
+
way to choose other settings. Existing seeds and CPU detection are unchanged.
|
|
164
|
+
|
|
165
|
+
`preprocessing.py` attaches training-only pipelines automatically. Non-LightGBM
|
|
166
|
+
models impute numeric missing values and encode categories with sparse indicator
|
|
167
|
+
columns (unseen categories are ignored). Linear models and neural nets also
|
|
168
|
+
scale numeric features. Lasso and neural-net **outcomes** additionally standardize
|
|
169
|
+
the training target; their regularization is in standardized-outcome units and
|
|
170
|
+
predictions are transformed back to the original units. Saved fold models include
|
|
171
|
+
all these fitted transformations. Each ensemble member owns its preprocessing.
|
|
172
|
+
LightGBM keeps its native categorical and missing-value handling.
|
|
173
|
+
|
|
174
|
+
SHAP diagnostics currently support LightGBM only. Other families still get the
|
|
175
|
+
standard calibration, propensity, and balance plots; SHAP is explicitly skipped.
|
|
176
|
+
|
|
177
|
+
## Processed input contract
|
|
178
|
+
|
|
179
|
+
Use a plain program dictionary with exactly the output path and RCT flag:
|
|
180
|
+
|
|
181
|
+
```python
|
|
182
|
+
from pathlib import Path
|
|
183
|
+
|
|
184
|
+
HERE = Path(__file__).resolve().parent
|
|
185
|
+
PROGRAM_REGISTRY = {
|
|
186
|
+
"my_program": {
|
|
187
|
+
"program_output_path": HERE / "processed_inputs" / "my_program",
|
|
188
|
+
"program_is_rct": True,
|
|
189
|
+
},
|
|
190
|
+
}
|
|
191
|
+
```
|
|
192
|
+
|
|
193
|
+
Select programs through `experiments[*].program_name`. Multiple programs per
|
|
194
|
+
experiment work. Use absolute program output paths; anchoring them to
|
|
195
|
+
`__file__` avoids dependence on where the command is run. For a non-RCT
|
|
196
|
+
program set `program_is_rct=False`.
|
|
197
|
+
|
|
198
|
+
The program output folder contains:
|
|
199
|
+
|
|
200
|
+
- `processed_data.parquet`: treated people and the comparison pool.
|
|
201
|
+
- `rct_control_set_aside.parquet`: separate held-out RCT controls, when applicable.
|
|
202
|
+
|
|
203
|
+
RCT-control placebo evaluation runs automatically when `program_is_rct=True`.
|
|
204
|
+
A missing or invalid RCT file does not discard the main estimates: the summary
|
|
205
|
+
records `rct_status` (`success`, `failed`, or `not_applicable`) and `rct_error`.
|
|
206
|
+
Each rerun clears that run's previous RCT predictions and plots before attempting
|
|
207
|
+
the calculation, and replaces its summary values. The input must contain
|
|
208
|
+
`pik`, `treatment`, the selected outcome/covariates, and the configured group
|
|
209
|
+
column. RCT rows keep `treatment=0`. Processed inputs supply categorical types;
|
|
210
|
+
estimation preserves native categoricals for LightGBM and automatically prepares
|
|
211
|
+
other supported model families as described above.
|
|
212
|
+
`observation_year` is not required unless selected as a covariate/group/outcome.
|
|
213
|
+
|
|
214
|
+
## Saved results and plotting
|
|
215
|
+
|
|
216
|
+
During estimation, diagnostics reuse the full prediction frames already in
|
|
217
|
+
memory, avoiding prediction-file reads and covariate joins. Saved predictions
|
|
218
|
+
remain slim. Standalone plots retain the file-loading behavior below.
|
|
219
|
+
`generate_plots()` and the standard single-run diagnostic helpers also accept
|
|
220
|
+
optional `predictions` (a Polars DataFrame or LazyFrame with this run's full
|
|
221
|
+
columns). For RCT overlap, pass its distinct `control_placebo_predictions` frame;
|
|
222
|
+
never substitute the main sample for the held-out RCT controls.
|
|
223
|
+
|
|
224
|
+
|
|
225
|
+
Under `output_dir`, the existing layout is retained:
|
|
226
|
+
|
|
227
|
+
- `raw/<experiment>/`: parquet summaries, per-run predictions, and model files.
|
|
228
|
+
- `plots/<experiment>/run_<number>/`: diagnostic plots.
|
|
229
|
+
- `tables/<experiment>/run_<number>/`: diagnostic tables.
|
|
230
|
+
- `master_summary.xlsx`: summaries for experiments under this output root.
|
|
231
|
+
|
|
232
|
+
Reruns overwrite each run's prediction/model files and replace its summary row
|
|
233
|
+
by run number. Nothing is cleared at startup; other summary rows remain,
|
|
234
|
+
including runs not reached after an interruption. If you change the experiment
|
|
235
|
+
combinations, use a new filename/output root to avoid retaining obsolete run
|
|
236
|
+
numbers from the previous specification. Do not run the same experiment
|
|
237
|
+
destination concurrently or write the shared master workbook concurrently.
|
|
238
|
+
|
|
239
|
+
Prediction files keep person ID, treatment, observed outcome (original name
|
|
240
|
+
and `outcome` alias), model predictions, run number, configured group column,
|
|
241
|
+
and main-sample fold assignments. Other covariates are omitted. The parquet
|
|
242
|
+
summary already records `covariate_set`; it now also records `data_path` and
|
|
243
|
+
`control_placebo_path` as absolute paths. Existing model names and parameter
|
|
244
|
+
strings remain. Fold models and their feature order remain in `.joblib` files.
|
|
245
|
+
|
|
246
|
+
Balance plots and model/SHAP loading recover requested covariates from the
|
|
247
|
+
saved source paths, joining by unique `pik`. Duplicate source IDs, unmatched
|
|
248
|
+
prediction IDs, and ambiguous summary rows are rejected. Keep source files
|
|
249
|
+
available and unchanged for later diagnostics. Filters are not reapplied.
|
|
250
|
+
|
|
251
|
+
```python
|
|
252
|
+
from pathlib import Path
|
|
253
|
+
import polars as pl
|
|
254
|
+
from dml_code.src.plotting import generate_plots
|
|
255
|
+
|
|
256
|
+
output_dir = Path("/secure/results")
|
|
257
|
+
summary = pl.read_parquet(output_dir / "raw/experiment/log_summary.parquet")
|
|
258
|
+
row = summary.filter(pl.col("run_number") == 1).row(0, named=True)
|
|
259
|
+
generate_plots("experiment", 1, row["covariate_set"], output_dir=output_dir)
|
|
260
|
+
```
|
|
261
|
+
|
|
262
|
+
The balance plot keeps its existing preferred list, skips covariates absent
|
|
263
|
+
from the run, and skips entirely when none remain. Categorical covariates use
|
|
264
|
+
category indicators. Calibration includes full-range and adaptive low-propensity
|
|
265
|
+
views (90th percentile cutoff, 20 bins); outcome axes adapt to the bin means and
|
|
266
|
+
confidence intervals without selecting on realized outcomes. Only small bin
|
|
267
|
+
summaries are converted to pandas. Diagnostics fail independently and return a
|
|
268
|
+
failure dictionary. Set `include_control_placebo=True` only when its saved
|
|
269
|
+
outputs are available.
|
|
270
|
+
|
|
271
|
+
Propensity SHAP uses each person's held-out model in repetition 1. Its overall
|
|
272
|
+
sample is capped at 50,000 across all folds; treated summaries include everyone.
|
|
273
|
+
Outcome SHAP uses repetition 1's fold 1 model, all treated people for the treated
|
|
274
|
+
summary, and held-out fold 1 rows for its overall plots. Magnitude tables are
|
|
275
|
+
written before plot rendering. Setting
|
|
276
|
+
`delete_raw_predictions=[True]` deletes prediction logs after current plots;
|
|
277
|
+
later model/SHAP plotting requires those logs and will not reconstruct them.
|
|
278
|
+
|
|
279
|
+
## Package layout and checks
|
|
280
|
+
|
|
281
|
+
- `dml_code/cli.py`: terminal argument parsing.
|
|
282
|
+
- `dml_code/estimation.py`: run orchestration, exposed as `estimate()`.
|
|
283
|
+
- `dml_code/src/config.py`: YAML, registries, Experiment/Run definitions, expansion and validation.
|
|
284
|
+
- `dml_code/src/nuisance.py`: tuning, folds, fitting and nuisance predictions.
|
|
285
|
+
- `dml_code/src/preprocessing.py`: model validation and training-only transformations.
|
|
286
|
+
- `dml_code/src/estimate_helpers.py`: preparation, ATT/variance and RCT-control evaluation.
|
|
287
|
+
- `dml_code/src/results.py`: summaries, saved predictions/models and upserts.
|
|
288
|
+
- `dml_code/src/plotting.py`: `generate_plots()` and individual diagnostics.
|
|
289
|
+
- `dml_code/src/paths.py`, `utils.py`: paths and shared utilities.
|
|
290
|
+
|
|
291
|
+
From a checkout, after installing the package:
|
|
292
|
+
|
|
293
|
+
```bash
|
|
294
|
+
python -m unittest discover -s tests/integration -v
|
|
295
|
+
```
|
|
296
|
+
|
|
297
|
+
Tests generate their own synthetic processed inputs and exercise real LightGBM
|
|
298
|
+
fits, estimator formulas, summary upserts, registry isolation and diagnostics.
|
|
@@ -0,0 +1,262 @@
|
|
|
1
|
+
# OI Workforce DML
|
|
2
|
+
|
|
3
|
+
Cross-fitted outcome, IPW, and AIPW estimation with saved models and diagnostics.
|
|
4
|
+
Supply prebuilt data using the input contract below.
|
|
5
|
+
|
|
6
|
+
## Version 0.2 migration
|
|
7
|
+
|
|
8
|
+
Version 0.2 runs estimation on prebuilt inputs. It no longer includes the old
|
|
9
|
+
build command or bundled project configuration. Copy `example_set_up/` from the
|
|
10
|
+
source distribution or repository, fill in the registries, and pass the YAML
|
|
11
|
+
experiment path to `dml-estimate`.
|
|
12
|
+
|
|
13
|
+
Replace `tune_nuisance_models` / `n_tuning_trials` with independent
|
|
14
|
+
`outcomes_tune` / `outcomes_tuning_trials` and
|
|
15
|
+
`propensity_tune` / `propensity_tuning_trials`. Remove `drop_nulls` and
|
|
16
|
+
`calculate_on_control_placebo`; covariate missingness is handled by the models,
|
|
17
|
+
and RCT-control evaluation follows the program registry. Unknown settings are
|
|
18
|
+
rejected. The examples below use the current format.
|
|
19
|
+
|
|
20
|
+
## Install and run
|
|
21
|
+
|
|
22
|
+
The PyPI distribution is `oi-workforce-dml`; the Python import remains
|
|
23
|
+
`dml_code`, and the command remains `dml-estimate`. The GitHub repository is
|
|
24
|
+
`dml-dev`. If the old `dml-dev` distribution is installed in this environment,
|
|
25
|
+
uninstall it first (`python -m pip uninstall dml-dev`): both distributions
|
|
26
|
+
provide the same import package and command.
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
# Install the published package:
|
|
30
|
+
python -m pip install oi-workforce-dml==0.2.0
|
|
31
|
+
# Or install a local checkout:
|
|
32
|
+
python -m pip install /path/to/dml-dev
|
|
33
|
+
# Or for development: python -m pip install -e /path/to/dml-dev
|
|
34
|
+
|
|
35
|
+
dml-estimate /path/to/my_analysis/experiment.yaml
|
|
36
|
+
# Equivalent:
|
|
37
|
+
python -m dml_code.cli /path/to/my_analysis/experiment.yaml
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
from dml_code import estimate
|
|
42
|
+
|
|
43
|
+
estimate("/path/to/my_analysis/experiment.yaml")
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
The experiment name comes from the YAML filename (`experiment` above).
|
|
47
|
+
The command can run from any working directory. YAML is parsed safely; Python
|
|
48
|
+
experiment configurations are no longer accepted. `example_set_up/` is a
|
|
49
|
+
copyable checkout example, outside the installed package. It contains sibling
|
|
50
|
+
`experiments/` and `registries/` directories; supply your own processed inputs.
|
|
51
|
+
|
|
52
|
+
## YAML experiment configuration
|
|
53
|
+
|
|
54
|
+
```yaml
|
|
55
|
+
registry_dir: ../registries
|
|
56
|
+
output_dir: ../results
|
|
57
|
+
experiments:
|
|
58
|
+
- program_name: [my_program]
|
|
59
|
+
outcomes: [wages_2]
|
|
60
|
+
covariate_set_name: [baseline]
|
|
61
|
+
filter_set_name: [eligible]
|
|
62
|
+
outcomes_model_name: [lightgbm]
|
|
63
|
+
outcomes_tune: [false]
|
|
64
|
+
propensity_model_name: [lightgbm]
|
|
65
|
+
propensity_tune: [false]
|
|
66
|
+
group_col: [site]
|
|
67
|
+
delete_raw_predictions: [false]
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
`registry_dir` and `output_dir` resolve relative to the YAML file, or can be
|
|
71
|
+
absolute. The registry directory contains these fixed Python files:
|
|
72
|
+
|
|
73
|
+
| File | Dictionary |
|
|
74
|
+
| --- | --- |
|
|
75
|
+
| `programs.py` | `PROGRAM_REGISTRY` |
|
|
76
|
+
| `covariate_sets.py` | `COVARIATE_SET_REGISTRY` |
|
|
77
|
+
| `filter_sets.py` | `FILTER_SET_REGISTRY` |
|
|
78
|
+
| `models.py` | `OUTCOMES_MODEL_REGISTRY`, `PROPENSITY_MODEL_REGISTRY` |
|
|
79
|
+
|
|
80
|
+
Registry dictionaries map names to program definitions, covariate lists,
|
|
81
|
+
lists of Polars filter functions, and supported sklearn/LightGBM model instances. They remain
|
|
82
|
+
executable Python; load registry files you trust. Bare sibling imports,
|
|
83
|
+
package imports such as `from registries.models import ...`, and relative
|
|
84
|
+
imports such as `from .models import ...` are supported. Sibling modules are
|
|
85
|
+
isolated between configs loaded sequentially in one Python process.
|
|
86
|
+
|
|
87
|
+
Experiment fields remain lists and expand into combinations. Multiple blocks
|
|
88
|
+
are allowed. Model templates are cloned for each run/fold. `dml_type` is not
|
|
89
|
+
accepted. Unknown top-level and experiment settings are rejected. Optional defaults:
|
|
90
|
+
no tuning, no grouping (`[]` or `[null]`), and retain predictions.
|
|
91
|
+
Missing covariates are retained; rows missing identity, treatment, or outcome are dropped.
|
|
92
|
+
|
|
93
|
+
Set `outcomes_tune: [true]` with `outcomes_tuning_trials: [15]` to tune outcomes;
|
|
94
|
+
use `propensity_tune` and `propensity_tuning_trials` independently for propensity.
|
|
95
|
+
False uses the registry preset. True uses the built-in tuning recipe before
|
|
96
|
+
cross-fitting. Omit a learner's trial budget when it is not tuned. For a mixed
|
|
97
|
+
`[false, true]` grid, budgets expand tuned runs only; the preset runs once.
|
|
98
|
+
Grouping generates a pooled run plus each non-null group with at least
|
|
99
|
+
30 treated rows in the input. Each group uses its treated rows and all controls.
|
|
100
|
+
Filters select rows only: do not transform covariates, because plotting later
|
|
101
|
+
reloads their values from the source file.
|
|
102
|
+
|
|
103
|
+
## Model presets and tuning
|
|
104
|
+
|
|
105
|
+
The model registries contain ordinary model instances. Supported families are
|
|
106
|
+
LightGBM, Lasso (outcomes) / LogisticRegression (propensity), random forests,
|
|
107
|
+
MLP neural networks, and VotingRegressor / soft VotingClassifier ensembles.
|
|
108
|
+
The example model registry contains a preset for each. Ensemble members use
|
|
109
|
+
fixed presets and weights; requesting ensemble tuning is an error.
|
|
110
|
+
|
|
111
|
+
Each learner has its own `*_tune` switch and `*_tuning_trials` budget. False
|
|
112
|
+
uses the registry preset. True clones that preset and overrides the parameters
|
|
113
|
+
in the built-in tuning recipe, retaining parameters outside that recipe. Tuning
|
|
114
|
+
happens once before cross-fitting, using the existing 80/20 split, outcome MSE
|
|
115
|
+
on controls, and binary log loss for propensity. The best LightGBM iteration
|
|
116
|
+
count is retained after early stopping. Search spaces live in `nuisance.py`:
|
|
117
|
+
|
|
118
|
+
| Family | Outcomes | Propensity |
|
|
119
|
+
| --- | --- | --- |
|
|
120
|
+
| LightGBM | Existing outcome tree search | Existing constrained propensity tree search |
|
|
121
|
+
| Linear | Lasso regularization | L1 logistic regularization, favoring stronger penalties |
|
|
122
|
+
| Forest | Deeper trees, smaller leaves | Shallower trees, larger leaves |
|
|
123
|
+
| Neural net | Larger hidden layers, lighter regularization | Smaller layers, stronger regularization |
|
|
124
|
+
|
|
125
|
+
These are starting ranges, not guarantees of calibration or adequate overlap.
|
|
126
|
+
A small trial budget explores only a few parameters. Manual presets remain the
|
|
127
|
+
way to choose other settings. Existing seeds and CPU detection are unchanged.
|
|
128
|
+
|
|
129
|
+
`preprocessing.py` attaches training-only pipelines automatically. Non-LightGBM
|
|
130
|
+
models impute numeric missing values and encode categories with sparse indicator
|
|
131
|
+
columns (unseen categories are ignored). Linear models and neural nets also
|
|
132
|
+
scale numeric features. Lasso and neural-net **outcomes** additionally standardize
|
|
133
|
+
the training target; their regularization is in standardized-outcome units and
|
|
134
|
+
predictions are transformed back to the original units. Saved fold models include
|
|
135
|
+
all these fitted transformations. Each ensemble member owns its preprocessing.
|
|
136
|
+
LightGBM keeps its native categorical and missing-value handling.
|
|
137
|
+
|
|
138
|
+
SHAP diagnostics currently support LightGBM only. Other families still get the
|
|
139
|
+
standard calibration, propensity, and balance plots; SHAP is explicitly skipped.
|
|
140
|
+
|
|
141
|
+
## Processed input contract
|
|
142
|
+
|
|
143
|
+
Use a plain program dictionary with exactly the output path and RCT flag:
|
|
144
|
+
|
|
145
|
+
```python
|
|
146
|
+
from pathlib import Path
|
|
147
|
+
|
|
148
|
+
HERE = Path(__file__).resolve().parent
|
|
149
|
+
PROGRAM_REGISTRY = {
|
|
150
|
+
"my_program": {
|
|
151
|
+
"program_output_path": HERE / "processed_inputs" / "my_program",
|
|
152
|
+
"program_is_rct": True,
|
|
153
|
+
},
|
|
154
|
+
}
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
Select programs through `experiments[*].program_name`. Multiple programs per
|
|
158
|
+
experiment work. Use absolute program output paths; anchoring them to
|
|
159
|
+
`__file__` avoids dependence on where the command is run. For a non-RCT
|
|
160
|
+
program set `program_is_rct=False`.
|
|
161
|
+
|
|
162
|
+
The program output folder contains:
|
|
163
|
+
|
|
164
|
+
- `processed_data.parquet`: treated people and the comparison pool.
|
|
165
|
+
- `rct_control_set_aside.parquet`: separate held-out RCT controls, when applicable.
|
|
166
|
+
|
|
167
|
+
RCT-control placebo evaluation runs automatically when `program_is_rct=True`.
|
|
168
|
+
A missing or invalid RCT file does not discard the main estimates: the summary
|
|
169
|
+
records `rct_status` (`success`, `failed`, or `not_applicable`) and `rct_error`.
|
|
170
|
+
Each rerun clears that run's previous RCT predictions and plots before attempting
|
|
171
|
+
the calculation, and replaces its summary values. The input must contain
|
|
172
|
+
`pik`, `treatment`, the selected outcome/covariates, and the configured group
|
|
173
|
+
column. RCT rows keep `treatment=0`. Processed inputs supply categorical types;
|
|
174
|
+
estimation preserves native categoricals for LightGBM and automatically prepares
|
|
175
|
+
other supported model families as described above.
|
|
176
|
+
`observation_year` is not required unless selected as a covariate/group/outcome.
|
|
177
|
+
|
|
178
|
+
## Saved results and plotting
|
|
179
|
+
|
|
180
|
+
During estimation, diagnostics reuse the full prediction frames already in
|
|
181
|
+
memory, avoiding prediction-file reads and covariate joins. Saved predictions
|
|
182
|
+
remain slim. Standalone plots retain the file-loading behavior below.
|
|
183
|
+
`generate_plots()` and the standard single-run diagnostic helpers also accept
|
|
184
|
+
optional `predictions` (a Polars DataFrame or LazyFrame with this run's full
|
|
185
|
+
columns). For RCT overlap, pass its distinct `control_placebo_predictions` frame;
|
|
186
|
+
never substitute the main sample for the held-out RCT controls.
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
Under `output_dir`, the existing layout is retained:
|
|
190
|
+
|
|
191
|
+
- `raw/<experiment>/`: parquet summaries, per-run predictions, and model files.
|
|
192
|
+
- `plots/<experiment>/run_<number>/`: diagnostic plots.
|
|
193
|
+
- `tables/<experiment>/run_<number>/`: diagnostic tables.
|
|
194
|
+
- `master_summary.xlsx`: summaries for experiments under this output root.
|
|
195
|
+
|
|
196
|
+
Reruns overwrite each run's prediction/model files and replace its summary row
|
|
197
|
+
by run number. Nothing is cleared at startup; other summary rows remain,
|
|
198
|
+
including runs not reached after an interruption. If you change the experiment
|
|
199
|
+
combinations, use a new filename/output root to avoid retaining obsolete run
|
|
200
|
+
numbers from the previous specification. Do not run the same experiment
|
|
201
|
+
destination concurrently or write the shared master workbook concurrently.
|
|
202
|
+
|
|
203
|
+
Prediction files keep person ID, treatment, observed outcome (original name
|
|
204
|
+
and `outcome` alias), model predictions, run number, configured group column,
|
|
205
|
+
and main-sample fold assignments. Other covariates are omitted. The parquet
|
|
206
|
+
summary already records `covariate_set`; it now also records `data_path` and
|
|
207
|
+
`control_placebo_path` as absolute paths. Existing model names and parameter
|
|
208
|
+
strings remain. Fold models and their feature order remain in `.joblib` files.
|
|
209
|
+
|
|
210
|
+
Balance plots and model/SHAP loading recover requested covariates from the
|
|
211
|
+
saved source paths, joining by unique `pik`. Duplicate source IDs, unmatched
|
|
212
|
+
prediction IDs, and ambiguous summary rows are rejected. Keep source files
|
|
213
|
+
available and unchanged for later diagnostics. Filters are not reapplied.
|
|
214
|
+
|
|
215
|
+
```python
|
|
216
|
+
from pathlib import Path
|
|
217
|
+
import polars as pl
|
|
218
|
+
from dml_code.src.plotting import generate_plots
|
|
219
|
+
|
|
220
|
+
output_dir = Path("/secure/results")
|
|
221
|
+
summary = pl.read_parquet(output_dir / "raw/experiment/log_summary.parquet")
|
|
222
|
+
row = summary.filter(pl.col("run_number") == 1).row(0, named=True)
|
|
223
|
+
generate_plots("experiment", 1, row["covariate_set"], output_dir=output_dir)
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
The balance plot keeps its existing preferred list, skips covariates absent
|
|
227
|
+
from the run, and skips entirely when none remain. Categorical covariates use
|
|
228
|
+
category indicators. Calibration includes full-range and adaptive low-propensity
|
|
229
|
+
views (90th percentile cutoff, 20 bins); outcome axes adapt to the bin means and
|
|
230
|
+
confidence intervals without selecting on realized outcomes. Only small bin
|
|
231
|
+
summaries are converted to pandas. Diagnostics fail independently and return a
|
|
232
|
+
failure dictionary. Set `include_control_placebo=True` only when its saved
|
|
233
|
+
outputs are available.
|
|
234
|
+
|
|
235
|
+
Propensity SHAP uses each person's held-out model in repetition 1. Its overall
|
|
236
|
+
sample is capped at 50,000 across all folds; treated summaries include everyone.
|
|
237
|
+
Outcome SHAP uses repetition 1's fold 1 model, all treated people for the treated
|
|
238
|
+
summary, and held-out fold 1 rows for its overall plots. Magnitude tables are
|
|
239
|
+
written before plot rendering. Setting
|
|
240
|
+
`delete_raw_predictions=[True]` deletes prediction logs after current plots;
|
|
241
|
+
later model/SHAP plotting requires those logs and will not reconstruct them.
|
|
242
|
+
|
|
243
|
+
## Package layout and checks
|
|
244
|
+
|
|
245
|
+
- `dml_code/cli.py`: terminal argument parsing.
|
|
246
|
+
- `dml_code/estimation.py`: run orchestration, exposed as `estimate()`.
|
|
247
|
+
- `dml_code/src/config.py`: YAML, registries, Experiment/Run definitions, expansion and validation.
|
|
248
|
+
- `dml_code/src/nuisance.py`: tuning, folds, fitting and nuisance predictions.
|
|
249
|
+
- `dml_code/src/preprocessing.py`: model validation and training-only transformations.
|
|
250
|
+
- `dml_code/src/estimate_helpers.py`: preparation, ATT/variance and RCT-control evaluation.
|
|
251
|
+
- `dml_code/src/results.py`: summaries, saved predictions/models and upserts.
|
|
252
|
+
- `dml_code/src/plotting.py`: `generate_plots()` and individual diagnostics.
|
|
253
|
+
- `dml_code/src/paths.py`, `utils.py`: paths and shared utilities.
|
|
254
|
+
|
|
255
|
+
From a checkout, after installing the package:
|
|
256
|
+
|
|
257
|
+
```bash
|
|
258
|
+
python -m unittest discover -s tests/integration -v
|
|
259
|
+
```
|
|
260
|
+
|
|
261
|
+
Tests generate their own synthetic processed inputs and exercise real LightGBM
|
|
262
|
+
fits, estimator formulas, summary upserts, registry isolation and diagnostics.
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""Terminal entry point for YAML experiments."""
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
from dml_code.estimation import estimate
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def cli() -> None:
|
|
10
|
+
parser = argparse.ArgumentParser()
|
|
11
|
+
parser.add_argument("config_path", type=Path, help="YAML experiment configuration file")
|
|
12
|
+
args = parser.parse_args()
|
|
13
|
+
estimate(config_path=args.config_path)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
if __name__ == "__main__":
|
|
17
|
+
cli()
|