oi-workforce-dml 0.2.0__py3-none-any.whl

dml_code/__init__.py

@@ -0,0 +1,5 @@
+"""Cross-fitted program evaluation from processed inputs."""
+
+from dml_code.estimation import estimate
+
+__all__ = ["estimate"]

dml_code/cli.py

@@ -0,0 +1,17 @@
+"""Terminal entry point for YAML experiments."""
+
+import argparse
+from pathlib import Path
+
+from dml_code.estimation import estimate
+
+
+def cli() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("config_path", type=Path, help="YAML experiment configuration file")
+    args = parser.parse_args()
+    estimate(config_path=args.config_path)
+
+
+if __name__ == "__main__":
+    cli()

dml_code/estimation.py

@@ -0,0 +1,128 @@
+"""Run AIPW estimation for a YAML experiment config.
+
+Loads an experiment, expands its registry names into runs, and for each run:
+prepares data, fits cross-fitted nuisance models, computes ATT/AIPW estimates
+(plus control placebos), and saves logs. Groups are separate runs: a group_col
+in the config fans out into one run per group value, each fit and logged on only
+its own rows, plus the all-data run.
+"""
+
+import time
+from pathlib import Path
+
+from dml_code.src.config import Run, get_experiments, unpack_runs, validate_runs
+from dml_code.src.nuisance import fit_nuisances
+from dml_code.src.estimate_helpers import (
+    calculate_control_placebo_estimates, calculate_control_placebo_predictions,
+    calculate_estimates, prepare_estimation_data,
+)
+from dml_code.src.results import (
+    build_summary_row, delete_df, prediction_columns, rebuild_master_summary,
+    save_df, save_models, clear_control_placebo_outputs,
+)
+from dml_code.src.plotting import generate_plots
+from dml_code.src.utils import log_process_resources, time_elapsed, trim_memory
+
+
+def estimate(config_path: str | Path) -> None:
+    """Run all expanded estimation runs from one external YAML config."""
+
+    config_path = Path(config_path).resolve()
+    experiment_name = config_path.stem
+    experiments, registries, output_dir = get_experiments(config_path)
+    runs = unpack_runs(experiments, registries)
+    validate_runs(runs)
+
+    stop_resource_logging = log_process_resources(interval=5)
+    failures: list[tuple[int, str]] = []
+    try:
+        for run_number, run in enumerate(runs, start=1):
+            print(
+                f"\nRun {run_number}/{len(runs)}: {run.program_name}, "
+                f"outcome={run.outcome}, group={run.group_value or 'all'}"
+            )
+            try:
+                run_one(experiment_name, run, run_number, output_dir=output_dir)
+                print(f"Run {run_number}/{len(runs)} complete.\n")
+            except Exception as error:
+                # Isolate failures so one bad run doesn't discard the others.
+                print(f"Run {run_number}/{len(runs)} FAILED: {error}\n")
+                failures.append((run_number, str(error)))
+            trim_memory()
+    finally:
+        stop_resource_logging.set()
+
+    # Rebuild the Excel master once from summaries in this output directory.
+    rebuild_master_summary(output_dir=output_dir)
+
+    if failures:
+        print(f"\n{len(failures)} run(s) failed:")
+        for run_number, error in failures:
+            print(f"  Run {run_number}: {error}")
+        raise RuntimeError(f"{len(failures)} estimation run(s) failed; see errors above")
+
+
+def run_one(experiment_name: str, run: Run, run_number: int, *, output_dir: Path) -> None:
+    """Estimate, summarize, and save one run."""
+
+    start = time.time()
+    df, num_nulls_dropped = prepare_estimation_data(run)
+
+    print("Begin nuisance prediction")
+    estimate_start = time.time()
+    nuisance = fit_nuisances(run, df, run_number)
+    estimates = calculate_estimates(nuisance.predictions, run)
+
+    save_models(nuisance.models, experiment_name, run_number, output_dir=output_dir)
+    save_df("predictions", prediction_columns(nuisance.predictions, run),
+            experiment_name, run_number, output_dir=output_dir)
+
+    placebo_ready = False
+    rct_status = "not_applicable"
+    rct_error = None
+    # Explicit nulls replace any estimates from an earlier successful RCT run.
+    estimates.update({name: None for name in (
+        "placebo_aipw_att", "placebo_ipw_att", "placebo_outcomes_att", "placebo_treated_mean"
+    )})
+    try:
+        clear_control_placebo_outputs(experiment_name, run_number, output_dir=output_dir)
+        if run.calculate_on_control_placebo:
+            placebo_predictions = calculate_control_placebo_predictions(
+                run, nuisance.models, run_number
+            )
+            placebo_estimates = calculate_control_placebo_estimates(placebo_predictions, estimates, run)
+            save_df(
+                "control_placebo_predictions", placebo_predictions,
+                experiment_name, run_number, output_dir=output_dir,
+            )
+            estimates.update(placebo_estimates)
+            placebo_ready = True
+            rct_status = "success"
+    except Exception as error:
+        rct_status = "failed"
+        rct_error = str(error)
+        print(f"RCT-control placebo evaluation failed: {error}")
+
+    end = time.time()
+    summary = build_summary_row(
+        experiment_name, run, run_number,
+        estimates=estimates,
+        predictions=nuisance.predictions,
+        nuisance=nuisance,
+        num_nulls_dropped=num_nulls_dropped,
+        estimate_run_time=time_elapsed(estimate_start, end),
+        total_run_time=time_elapsed(start, end),
+        output_dir=output_dir,
+        rct_status=rct_status,
+        rct_error=rct_error,
+    )
+    save_df("summary", summary, experiment_name, output_dir=output_dir)
+    print(f"AIPW ATT: {estimates['aipw_att']:,.2f}")
+
+    generate_plots(experiment_name, run_number, run.covariate_set, output_dir=output_dir,
+                   include_control_placebo=placebo_ready, predictions=nuisance.predictions,
+                   control_placebo_predictions=placebo_predictions if placebo_ready else None)
+    if run.delete_raw_predictions:
+        delete_df("predictions", experiment_name, run_number, output_dir=output_dir)
+        if run.calculate_on_control_placebo:
+            delete_df("control_placebo_predictions", experiment_name, run_number, output_dir=output_dir)

dml_code/src/__init__.py

@@ -0,0 +1 @@
+"""Configuration, estimation helpers, persistence and diagnostics."""

dml_code/src/config.py

@@ -0,0 +1,337 @@
+"""YAML experiments, Python registries, run expansion and validation."""
+
+from dataclasses import dataclass, fields
+import importlib
+from itertools import product
+from pathlib import Path
+import sys
+import types
+from typing import Any
+
+import polars as pl
+from sklearn.base import clone
+import yaml
+
+from dml_code.src.utils import JOIN_KEY, TREATMENT_COL
+from dml_code.src.preprocessing import model_family, validate_model
+
+EXPERIMENT_DEFAULTS: dict[str, list] = {
+    "outcomes_tune": [False],
+    "outcomes_tuning_trials": [None],
+    "propensity_tune": [False],
+    "propensity_tuning_trials": [None],
+    "group_col": [None],
+    "delete_raw_predictions": [False],
+}
+
+
+@dataclass(frozen=True)
+class Experiment:
+    """One YAML config block of registry name lists; runs are their cross product."""
+
+    program_name: list[str]
+    outcomes: list[str]
+    covariate_set_name: list[str]
+    filter_set_name: list[str]
+    outcomes_model_name: list[str]
+    propensity_model_name: list[str]
+    outcomes_tune: list[bool]
+    outcomes_tuning_trials: list[int | None]
+    propensity_tune: list[bool]
+    propensity_tuning_trials: list[int | None]
+    group_col: list[str | None]
+    delete_raw_predictions: list[bool]
+
+
+@dataclass(frozen=True)
+class Run:
+    """Concrete estimation run after resolving experiment names."""
+
+    program_name: str
+    data_path: Path
+    control_placebo_path: Path | None
+    join_key: str
+    treatment: str
+    outcome: str
+    covariate_set: list[str]
+    filter_set: list[Any]
+    covariate_set_name: str
+    filter_set_name: str
+    outcomes_model_name: str
+    propensity_model_name: str
+    outcomes_tune: bool
+    outcomes_tuning_trials: int | None
+    propensity_tune: bool
+    propensity_tuning_trials: int | None
+    outcomes_model: Any
+    propensity_model: Any
+    calculate_on_control_placebo: bool
+    # group_col is the column we split on; group_value is this run's slice of it
+    # (None means the all-data run). Each group value is its own run, fit and
+    # logged on only its rows.
+    group_col: str | None
+    group_value: Any
+    delete_raw_predictions: bool
+
+
+def get_experiments(config_path: str | Path) -> tuple[list[Experiment], dict, Path]:
+    """Load a YAML experiment and the fixed Python registry dictionaries."""
+    config_path = Path(config_path).resolve()
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Experiment config does not exist: {config_path}")
+    if config_path.suffix.lower() not in {".yaml", ".yml"}:
+        raise ValueError("Experiment config must be a YAML file (.yaml or .yml)")
+    with config_path.open() as stream:
+        config = yaml.safe_load(stream)
+    if not isinstance(config, dict):
+        raise ValueError("YAML config must be a dictionary")
+    unknown = sorted(set(config) - {"registry_dir", "output_dir", "experiments"})
+    if unknown:
+        raise ValueError(f"Unknown top-level YAML keys: {unknown}")
+    for key in ("registry_dir", "output_dir"):
+        if not isinstance(config.get(key), str) or not config[key].strip():
+            raise ValueError(f"YAML config must define a non-empty {key} path")
+    registry_dir = Path(config["registry_dir"]).expanduser()
+    if not registry_dir.is_absolute():
+        registry_dir = config_path.parent / registry_dir
+    registry_dir = registry_dir.resolve()
+    registries = _load_registries(registry_dir)
+    output_dir = Path(config["output_dir"]).expanduser()
+    if not output_dir.is_absolute():
+        output_dir = config_path.parent / output_dir
+
+    raw_blocks = config.get("experiments")
+    if not isinstance(raw_blocks, list) or not raw_blocks:
+        raise ValueError("experiments must be a non-empty list of dictionaries")
+    field_names = {field.name for field in fields(Experiment)}
+    experiments = []
+    for raw in raw_blocks:
+        if not isinstance(raw, dict):
+            raise ValueError("Each experiments block must be a dictionary")
+        unknown = sorted(set(raw) - field_names)
+        missing = sorted(field_names - set(raw) - set(EXPERIMENT_DEFAULTS))
+        if unknown or missing:
+            raise ValueError(
+                f"Experiment '{config_path.stem}': missing required keys {missing}, "
+                f"unknown keys {unknown}"
+            )
+        values = {**EXPERIMENT_DEFAULTS, **raw}
+        if values["group_col"] == []:
+            values["group_col"] = [None]
+        for name, value in values.items():
+            if not isinstance(value, list) or not value:
+                raise ValueError(f"{name} must be a non-empty list")
+        for learner in ("outcomes", "propensity"):
+            switches = values[f"{learner}_tune"]
+            budgets = values[f"{learner}_tuning_trials"]
+            if any(type(value) is not bool for value in switches):
+                raise ValueError(f"{learner}_tune must contain booleans")
+            if True in switches:
+                if any(type(value) is not int or value <= 0 for value in budgets):
+                    raise ValueError(f"{learner}_tuning_trials must contain positive integers when tuning")
+            elif f"{learner}_tuning_trials" in raw:
+                raise ValueError(f"{learner}_tuning_trials is unused when {learner}_tune is false")
+        experiments.append(Experiment(**values))
+    return experiments, registries, output_dir.resolve()
+
+
+def _tuning_options(switches: list[bool], budgets: list[int | None]) -> list[tuple]:
+    """Expand tuning budgets only for tuned runs, not repeated preset runs."""
+    return list(dict.fromkeys(
+        (tune, budget)
+        for tune in switches
+        for budget in (budgets if tune else [None])
+    ))
+
+
+def unpack_runs(experiments: list[Experiment], registries: dict) -> list[Run]:
+    """Resolve supplied registry names and expand experiments into concrete runs."""
+    programs = registries["PROGRAM_REGISTRY"]
+    covariate_sets = registries["COVARIATE_SET_REGISTRY"]
+    filter_sets = registries["FILTER_SET_REGISTRY"]
+    outcomes_models = registries["OUTCOMES_MODEL_REGISTRY"]
+    propensity_models = registries["PROPENSITY_MODEL_REGISTRY"]
+
+    runs = []
+    for experiment in experiments:
+        for names, registry, label in (
+            (experiment.program_name, programs, "program"),
+            (experiment.covariate_set_name, covariate_sets, "covariate set"),
+            (experiment.filter_set_name, filter_sets, "filter set"),
+            (experiment.outcomes_model_name, outcomes_models, "outcomes model"),
+            (experiment.propensity_model_name, propensity_models, "propensity model"),
+        ):
+            missing = [name for name in names if name not in registry]
+            if missing:
+                raise ValueError(f"Unknown {label} names: {missing}")
+        combos = product(
+            experiment.program_name, experiment.outcomes,
+            experiment.covariate_set_name, experiment.filter_set_name,
+            experiment.outcomes_model_name, experiment.propensity_model_name,
+            _tuning_options(experiment.outcomes_tune, experiment.outcomes_tuning_trials),
+            _tuning_options(experiment.propensity_tune, experiment.propensity_tuning_trials),
+            experiment.group_col, experiment.delete_raw_predictions,
+        )
+        for (
+            program_name, outcome, covariate_set_name, filter_set_name,
+            outcomes_model_name, propensity_model_name, outcomes_tuning,
+            propensity_tuning,
+            group_col, delete_raw_predictions,
+        ) in combos:
+            program = programs[program_name]
+            if not isinstance(program, dict):
+                raise ValueError(f"Program {program_name!r} must be a dictionary")
+            if not isinstance(program.get("program_is_rct"), bool):
+                raise ValueError(f"Program {program_name!r} needs program_is_rct=True or False")
+            if "program_output_path" not in program:
+                raise ValueError(f"Program {program_name!r} needs program_output_path")
+            folder = Path(program["program_output_path"]).expanduser()
+            if not folder.is_absolute():
+                raise ValueError(
+                    f"Program {program_name!r}: program_output_path must be absolute; "
+                    "use Path(__file__).resolve().parent in the shared registry for relative locations"
+                )
+            data_path = folder.resolve() / "processed_data.parquet"
+            control_placebo_path = (
+                folder.resolve() / "rct_control_set_aside.parquet"
+                if program["program_is_rct"] else None
+            )
+            group_col = None if group_col in (None, "", False) else str(group_col)
+            required = [JOIN_KEY, TREATMENT_COL, outcome, *covariate_sets[covariate_set_name]]
+            if group_col:
+                required.append(group_col)
+            _validate_input_columns(data_path, required)
+            # Preserve current group discovery and the pooled run.
+            group_values = [None, *discover_group_values(data_path, group_col)]
+            for group_value in group_values:
+                runs.append(Run(
+                    program_name=program_name,
+                    data_path=data_path,
+                    control_placebo_path=control_placebo_path,
+                    join_key=JOIN_KEY,
+                    treatment=TREATMENT_COL,
+                    outcome=outcome,
+                    covariate_set=covariate_sets[covariate_set_name],
+                    filter_set=filter_sets[filter_set_name],
+                    covariate_set_name=covariate_set_name,
+                    filter_set_name=filter_set_name,
+                    outcomes_model_name=outcomes_model_name,
+                    propensity_model_name=propensity_model_name,
+                    outcomes_tune=outcomes_tuning[0],
+                    outcomes_tuning_trials=outcomes_tuning[1],
+                    propensity_tune=propensity_tuning[0],
+                    propensity_tuning_trials=propensity_tuning[1],
+                    outcomes_model=clone(outcomes_models[outcomes_model_name]),
+                    propensity_model=clone(propensity_models[propensity_model_name]),
+                    calculate_on_control_placebo=program["program_is_rct"],
+                    group_col=group_col,
+                    group_value=group_value,
+                    delete_raw_predictions=delete_raw_predictions,
+                ))
+    return runs
+
+
+def discover_group_values(data_path: Path, group_col: str | None) -> list[Any]:
+    """Non-null groups with at least 30 treated rows; pooled runs are separate."""
+    if not group_col:
+        return []
+    return (
+        pl.scan_parquet(data_path)
+        .filter(pl.col(TREATMENT_COL) == 1)
+        .select(group_col)
+        .drop_nulls()
+        .group_by(group_col)
+        .agg(pl.len().alias("n_obs"))
+        .filter(pl.col("n_obs") >= 30)
+        .select(group_col)
+        .collect().to_series().sort().to_list()
+    )
+
+
+def _validate_input_columns(path: Path, required: list[str]) -> None:
+    """Check files and columns before group discovery or model fitting."""
+    if not path.is_file():
+        raise FileNotFoundError(f"Input file does not exist: {path}")
+    schema = pl.scan_parquet(path).collect_schema()
+    missing = [name for name in dict.fromkeys(required) if name not in schema]
+    if missing:
+        raise ValueError(f"{path}: missing required columns {missing}")
+
+
+def validate_runs(runs: list[Run]) -> None:
+    """Validate required model types, input files and referenced columns."""
+    errors = []
+    for i, run in enumerate(runs):
+        label = f"Run #{i + 1} ({run.program_name}, {run.covariate_set_name}, {run.filter_set_name})"
+        for learner in ("outcomes", "propensity"):
+            model = getattr(run, f"{learner}_model")
+            try:
+                validate_model(model, learner)
+                if getattr(run, f"{learner}_tune") and model_family(model) == "ensemble":
+                    raise ValueError("Ensembles use registry presets; ensemble tuning is not supported")
+            except ValueError as error:
+                errors.append(f"{label}: {error}")
+        required = [run.join_key, run.treatment, run.outcome, *run.covariate_set]
+        if run.group_col:
+            required.append(run.group_col)
+        paths = [run.data_path]
+        for path in paths:
+            try:
+                _validate_input_columns(path, required)
+            except (ValueError, FileNotFoundError) as error:
+                errors.append(f"{label}: {error}")
+    if errors:
+        raise ValueError("Run validation failed:\n\n" + "\n".join(errors))
+
+
+def _load_registries(registry_dir: Path) -> dict:
+    """Load fixed registry files, isolating sibling imports between projects.
+
+    Both bare sibling imports and package-relative imports are supported.
+    Restore the caller's module cache and search path even if loading fails.
+    """
+    if not registry_dir.is_dir():
+        raise FileNotFoundError(f"Registry directory does not exist: {registry_dir}")
+    registry_files = {
+        "programs": ("PROGRAM_REGISTRY",),
+        "covariate_sets": ("COVARIATE_SET_REGISTRY",),
+        "filter_sets": ("FILTER_SET_REGISTRY",),
+        "models": ("OUTCOMES_MODEL_REGISTRY", "PROPENSITY_MODEL_REGISTRY"),
+    }
+    package_name = registry_dir.name if registry_dir.name.isidentifier() else "_dml_registries"
+    local_names = {
+        path.stem for path in registry_dir.iterdir()
+        if path.suffix == ".py" or (path.is_dir() and (path / "__init__.py").is_file())
+    } | {package_name}
+    saved_modules = {
+        name: module for name, module in sys.modules.copy().items()
+        if name.split(".")[0] in local_names
+    }
+    old_path = sys.path.copy()
+    for name in saved_modules:
+        del sys.modules[name]
+    try:
+        sys.path.insert(0, str(registry_dir))
+        # An explicit package supports relative imports even without __init__.py.
+        package = types.ModuleType(package_name)
+        package.__path__ = [str(registry_dir)]
+        package.__package__ = package_name
+        sys.modules[package_name] = package
+        registries = {}
+        for filename, names in registry_files.items():
+            path = registry_dir / f"{filename}.py"
+            if not path.is_file():
+                raise FileNotFoundError(f"Missing registry file: {path}")
+            module = importlib.import_module(f"{package_name}.{filename}")
+            for name in names:
+                value = getattr(module, name, None)
+                if not isinstance(value, dict):
+                    raise ValueError(f"{path.name} must expose a dictionary named {name}")
+                registries[name] = value
+        return registries
+    finally:
+        sys.path[:] = old_path
+        for name in list(sys.modules):
+            if name.split(".")[0] in local_names:
+                del sys.modules[name]
+        sys.modules.update(saved_modules)