source/train.py

from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import trackio
from causal import estimate_effects, fit_nuisance, generate_scm

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "causal-forge"
DATA_DIR = PROJECT_DIR / "data"
REGIMES = {
    "both_correct": (True, True),
    "propensity_misspecified": (False, True),
    "outcome_misspecified": (True, False),
    "both_misspecified": (False, False),
}
ESTIMATORS = ["naive", "ipw", "outcome_regression", "aipw"]


def summarize(records: list[dict]) -> dict:
    summary = {}
    for regime in REGIMES:
        subset = [record for record in records if record["regime"] == regime]
        truth = np.asarray([record["truth"] for record in subset])
        regime_summary = {}
        for estimator in ESTIMATORS:
            estimates = np.asarray([record[estimator] for record in subset])
            errors = estimates - truth
            regime_summary[estimator] = {
                "mean_estimate": float(estimates.mean()),
                "bias": float(errors.mean()),
                "mean_absolute_error": float(np.abs(errors).mean()),
                "rmse": float(np.sqrt(np.mean(errors**2))),
            }
        coverage = np.mean(
            [
                record["aipw_ci_low"] <= record["truth"] <= record["aipw_ci_high"]
                for record in subset
            ]
        )
        regime_summary["aipw"]["confidence_interval_coverage"] = float(coverage)
        summary[regime] = regime_summary
    return summary


def save_dataset(sample) -> None:
    frame = pd.DataFrame(sample.x, columns=["x1", "x2", "x3"])
    frame["treatment"] = sample.treatment
    frame["outcome"] = sample.outcome
    frame["true_propensity"] = sample.propensity
    frame["potential_outcome_control"] = sample.y0
    frame["potential_outcome_treated"] = sample.y1
    frame["individual_treatment_effect"] = sample.ite
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    frame.to_parquet(DATA_DIR / "causal_benchmark.parquet", index=False)


def main() -> None:
    replications = 100
    samples_per_replication = 3_000
    trackio.init(
        project="causal-forge",
        name="double-robustness-benchmark-v1",
        config={
            "replications": replications,
            "samples_per_replication": samples_per_replication,
            "cross_fitting_folds": 5,
            "regimes": list(REGIMES),
        },
    )
    records = []
    for replication in range(replications):
        sample = generate_scm(samples_per_replication, seed=9100 + replication)
        for regime, (propensity_correct, outcome_correct) in REGIMES.items():
            estimates = estimate_effects(
                sample,
                propensity_correct=propensity_correct,
                outcome_correct=outcome_correct,
                folds=5,
                seed=replication,
            )
            records.append(
                {"replication": replication, "regime": regime, **estimates}
            )
        if (replication + 1) % 10 == 0:
            recent = records[-40:]
            trackio.log(
                {
                    "replication": replication + 1,
                    "aipw_recent_mae": float(
                        np.mean(
                            [
                                abs(record["aipw"] - record["truth"])
                                for record in recent
                            ]
                        )
                    ),
                }
            )
    summary = summarize(records)
    benchmark = generate_scm(20_000, seed=12043)
    final_models = fit_nuisance(
        benchmark.x,
        benchmark.treatment,
        benchmark.outcome,
        propensity_correct=True,
        outcome_correct=True,
    )
    save_dataset(benchmark)
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(final_models, ARTIFACT_DIR / "nuisance_models.joblib")
    results = {
        "benchmark": "Causal Forge",
        "replications": replications,
        "samples_per_replication": samples_per_replication,
        "cross_fitting_folds": 5,
        "structural_truth": "known heterogeneous individual treatment effects",
        "summary": summary,
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2), encoding="utf-8"
    )
    pd.DataFrame(records).to_parquet(
        ARTIFACT_DIR / "replication_estimates.parquet", index=False
    )
    trackio.log(
        {
            "both_correct_aipw_mae": summary["both_correct"]["aipw"][
                "mean_absolute_error"
            ],
            "both_correct_naive_mae": summary["both_correct"]["naive"][
                "mean_absolute_error"
            ],
            "propensity_misspecified_aipw_mae": summary[
                "propensity_misspecified"
            ]["aipw"]["mean_absolute_error"],
            "outcome_misspecified_aipw_mae": summary["outcome_misspecified"][
                "aipw"
            ]["mean_absolute_error"],
        }
    )
    trackio.finish()
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()