Metadata-Version: 2.4
Name: oi-workforce-dml
Version: 0.2.0
Summary: Cross-fitted estimation and plotting for program evaluation.
Author: DML Pipeline Contributors
Project-URL: Repository, https://github.com/coreygb1/dml-dev
Keywords: administrative-data,causal-inference,doubleml,observational-data,program-evaluation
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Requires-Python: >=3.12
Description-Content-Type: text/markdown
Requires-Dist: doubleml
Requires-Dist: joblib
Requires-Dist: lightgbm
Requires-Dist: optuna
Requires-Dist: oi-tools[figures]
Requires-Dist: plotnine
Requires-Dist: polars
Requires-Dist: PyYAML
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: matplotlib
Requires-Dist: xlsxwriter
Requires-Dist: pyarrow
Requires-Dist: psutil
Requires-Dist: scikit-learn<1.10,>=1.2
Requires-Dist: shap
Requires-Dist: threadpoolctl
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"

# OI Workforce DML

Cross-fitted outcome, IPW, and AIPW estimation with saved models and diagnostics.
Supply prebuilt data using the input contract below.

## Version 0.2 migration

Version 0.2 runs estimation on prebuilt inputs. It no longer includes the old
build command or bundled project configuration. Copy `example_set_up/` from the
source distribution or repository, fill in the registries, and pass the YAML
experiment path to `dml-estimate`.

Replace `tune_nuisance_models` / `n_tuning_trials` with independent
`outcomes_tune` / `outcomes_tuning_trials` and
`propensity_tune` / `propensity_tuning_trials`. Remove `drop_nulls` and
`calculate_on_control_placebo`; covariate missingness is handled by the models,
and RCT-control evaluation follows the program registry. Unknown settings are
rejected. The examples below use the current format.

## Install and run

The PyPI distribution is `oi-workforce-dml`; the Python import remains
`dml_code`, and the command remains `dml-estimate`. The GitHub repository is
`dml-dev`. If the old `dml-dev` distribution is installed in this environment,
uninstall it first (`python -m pip uninstall dml-dev`): both distributions
provide the same import package and command.

```bash
# Install the published package:
python -m pip install oi-workforce-dml==0.2.0
# Or install a local checkout:
python -m pip install /path/to/dml-dev
# Or for development: python -m pip install -e /path/to/dml-dev

dml-estimate /path/to/my_analysis/experiment.yaml
# Equivalent:
python -m dml_code.cli /path/to/my_analysis/experiment.yaml
```

```python
from dml_code import estimate

estimate("/path/to/my_analysis/experiment.yaml")
```

The experiment name comes from the YAML filename (`experiment` above).
The command can run from any working directory. YAML is parsed safely; Python
experiment configurations are no longer accepted. `example_set_up/` is a
copyable checkout example, outside the installed package. It contains sibling
`experiments/` and `registries/` directories; supply your own processed inputs.

## YAML experiment configuration

```yaml
registry_dir: ../registries
output_dir: ../results
experiments:
  - program_name: [my_program]
    outcomes: [wages_2]
    covariate_set_name: [baseline]
    filter_set_name: [eligible]
    outcomes_model_name: [lightgbm]
    outcomes_tune: [false]
    propensity_model_name: [lightgbm]
    propensity_tune: [false]
    group_col: [site]
    delete_raw_predictions: [false]
```

`registry_dir` and `output_dir` resolve relative to the YAML file, or can be
absolute. The registry directory contains these fixed Python files:

| File | Dictionary |
| --- | --- |
| `programs.py` | `PROGRAM_REGISTRY` |
| `covariate_sets.py` | `COVARIATE_SET_REGISTRY` |
| `filter_sets.py` | `FILTER_SET_REGISTRY` |
| `models.py` | `OUTCOMES_MODEL_REGISTRY`, `PROPENSITY_MODEL_REGISTRY` |

Registry dictionaries map names to program definitions, covariate lists,
lists of Polars filter functions, and supported sklearn/LightGBM model instances. They remain
executable Python; load registry files you trust. Bare sibling imports,
package imports such as `from registries.models import ...`, and relative
imports such as `from .models import ...` are supported. Sibling modules are
isolated between configs loaded sequentially in one Python process.

Experiment fields remain lists and expand into combinations. Multiple blocks
are allowed. Model templates are cloned for each run/fold. `dml_type` is not
accepted. Unknown top-level and experiment settings are rejected. Optional defaults:
no tuning, no grouping (`[]` or `[null]`), and retain predictions.
Missing covariates are retained; rows missing identity, treatment, or outcome are dropped.

Set `outcomes_tune: [true]` with `outcomes_tuning_trials: [15]` to tune outcomes;
use `propensity_tune` and `propensity_tuning_trials` independently for propensity.
False uses the registry preset. True uses the built-in tuning recipe before
cross-fitting. Omit a learner's trial budget when it is not tuned. For a mixed
`[false, true]` grid, budgets expand tuned runs only; the preset runs once.
Grouping generates a pooled run plus each non-null group with at least
30 treated rows in the input. Each group uses its treated rows and all controls.
Filters select rows only: do not transform covariates, because plotting later
reloads their values from the source file.

## Model presets and tuning

The model registries contain ordinary model instances. Supported families are
LightGBM, Lasso (outcomes) / LogisticRegression (propensity), random forests,
MLP neural networks, and VotingRegressor / soft VotingClassifier ensembles.
The example model registry contains a preset for each. Ensemble members use
fixed presets and weights; requesting ensemble tuning is an error.

Each learner has its own `*_tune` switch and `*_tuning_trials` budget. False
uses the registry preset. True clones that preset and overrides the parameters
in the built-in tuning recipe, retaining parameters outside that recipe. Tuning
happens once before cross-fitting, using the existing 80/20 split, outcome MSE
on controls, and binary log loss for propensity. The best LightGBM iteration
count is retained after early stopping. Search spaces live in `nuisance.py`:

| Family | Outcomes | Propensity |
| --- | --- | --- |
| LightGBM | Existing outcome tree search | Existing constrained propensity tree search |
| Linear | Lasso regularization | L1 logistic regularization, favoring stronger penalties |
| Forest | Deeper trees, smaller leaves | Shallower trees, larger leaves |
| Neural net | Larger hidden layers, lighter regularization | Smaller layers, stronger regularization |

These are starting ranges, not guarantees of calibration or adequate overlap.
A small trial budget explores only a few parameters. Manual presets remain the
way to choose other settings. Existing seeds and CPU detection are unchanged.

`preprocessing.py` attaches training-only pipelines automatically. Non-LightGBM
models impute numeric missing values and encode categories with sparse indicator
columns (unseen categories are ignored). Linear models and neural nets also
scale numeric features. Lasso and neural-net **outcomes** additionally standardize
the training target; their regularization is in standardized-outcome units and
predictions are transformed back to the original units. Saved fold models include
all these fitted transformations. Each ensemble member owns its preprocessing.
LightGBM keeps its native categorical and missing-value handling.

SHAP diagnostics currently support LightGBM only. Other families still get the
standard calibration, propensity, and balance plots; SHAP is explicitly skipped.

## Processed input contract

Use a plain program dictionary with exactly the output path and RCT flag:

```python
from pathlib import Path

HERE = Path(__file__).resolve().parent
PROGRAM_REGISTRY = {
    "my_program": {
        "program_output_path": HERE / "processed_inputs" / "my_program",
        "program_is_rct": True,
    },
}
```

Select programs through `experiments[*].program_name`. Multiple programs per
experiment work. Use absolute program output paths; anchoring them to
`__file__` avoids dependence on where the command is run. For a non-RCT
program set `program_is_rct=False`.

The program output folder contains:

- `processed_data.parquet`: treated people and the comparison pool.
- `rct_control_set_aside.parquet`: separate held-out RCT controls, when applicable.

RCT-control placebo evaluation runs automatically when `program_is_rct=True`.
A missing or invalid RCT file does not discard the main estimates: the summary
records `rct_status` (`success`, `failed`, or `not_applicable`) and `rct_error`.
Each rerun clears that run's previous RCT predictions and plots before attempting
the calculation, and replaces its summary values. The input must contain
`pik`, `treatment`, the selected outcome/covariates, and the configured group
column. RCT rows keep `treatment=0`. Processed inputs supply categorical types;
estimation preserves native categoricals for LightGBM and automatically prepares
other supported model families as described above.
`observation_year` is not required unless selected as a covariate/group/outcome.

## Saved results and plotting

During estimation, diagnostics reuse the full prediction frames already in
memory, avoiding prediction-file reads and covariate joins. Saved predictions
remain slim. Standalone plots retain the file-loading behavior below.
`generate_plots()` and the standard single-run diagnostic helpers also accept
optional `predictions` (a Polars DataFrame or LazyFrame with this run's full
columns). For RCT overlap, pass its distinct `control_placebo_predictions` frame;
never substitute the main sample for the held-out RCT controls.


Under `output_dir`, the existing layout is retained:

- `raw/<experiment>/`: parquet summaries, per-run predictions, and model files.
- `plots/<experiment>/run_<number>/`: diagnostic plots.
- `tables/<experiment>/run_<number>/`: diagnostic tables.
- `master_summary.xlsx`: summaries for experiments under this output root.

Reruns overwrite each run's prediction/model files and replace its summary row
by run number. Nothing is cleared at startup; other summary rows remain,
including runs not reached after an interruption. If you change the experiment
combinations, use a new filename/output root to avoid retaining obsolete run
numbers from the previous specification. Do not run the same experiment
destination concurrently or write the shared master workbook concurrently.

Prediction files keep person ID, treatment, observed outcome (original name
and `outcome` alias), model predictions, run number, configured group column,
and main-sample fold assignments. Other covariates are omitted. The parquet
summary already records `covariate_set`; it now also records `data_path` and
`control_placebo_path` as absolute paths. Existing model names and parameter
strings remain. Fold models and their feature order remain in `.joblib` files.

Balance plots and model/SHAP loading recover requested covariates from the
saved source paths, joining by unique `pik`. Duplicate source IDs, unmatched
prediction IDs, and ambiguous summary rows are rejected. Keep source files
available and unchanged for later diagnostics. Filters are not reapplied.

```python
from pathlib import Path
import polars as pl
from dml_code.src.plotting import generate_plots

output_dir = Path("/secure/results")
summary = pl.read_parquet(output_dir / "raw/experiment/log_summary.parquet")
row = summary.filter(pl.col("run_number") == 1).row(0, named=True)
generate_plots("experiment", 1, row["covariate_set"], output_dir=output_dir)
```

The balance plot keeps its existing preferred list, skips covariates absent
from the run, and skips entirely when none remain. Categorical covariates use
category indicators. Calibration includes full-range and adaptive low-propensity
views (90th percentile cutoff, 20 bins); outcome axes adapt to the bin means and
confidence intervals without selecting on realized outcomes. Only small bin
summaries are converted to pandas. Diagnostics fail independently and return a
failure dictionary. Set `include_control_placebo=True` only when its saved
outputs are available.

Propensity SHAP uses each person's held-out model in repetition 1. Its overall
sample is capped at 50,000 across all folds; treated summaries include everyone.
Outcome SHAP uses repetition 1's fold 1 model, all treated people for the treated
summary, and held-out fold 1 rows for its overall plots. Magnitude tables are
written before plot rendering. Setting
`delete_raw_predictions=[True]` deletes prediction logs after current plots;
later model/SHAP plotting requires those logs and will not reconstruct them.

## Package layout and checks

- `dml_code/cli.py`: terminal argument parsing.
- `dml_code/estimation.py`: run orchestration, exposed as `estimate()`.
- `dml_code/src/config.py`: YAML, registries, Experiment/Run definitions, expansion and validation.
- `dml_code/src/nuisance.py`: tuning, folds, fitting and nuisance predictions.
- `dml_code/src/preprocessing.py`: model validation and training-only transformations.
- `dml_code/src/estimate_helpers.py`: preparation, ATT/variance and RCT-control evaluation.
- `dml_code/src/results.py`: summaries, saved predictions/models and upserts.
- `dml_code/src/plotting.py`: `generate_plots()` and individual diagnostics.
- `dml_code/src/paths.py`, `utils.py`: paths and shared utilities.

From a checkout, after installing the package:

```bash
python -m unittest discover -s tests/integration -v
```

Tests generate their own synthetic processed inputs and exercise real LightGBM
fits, estimator formulas, summary upserts, registry isolation and diagnostics.
