ENH: Optional sequential reproducibility control for repeated cross-fitting

.gitignore

@@ -37,3 +37,4 @@ doubleml/_version.py
 # uv
 .venv/
 uv.lock
+doc/api/generated/

benchmarks/reproducibility.py

@@ -0,0 +1,99 @@
+"""Fixed-data calibration and cost benchmark for scalar PLR aggregation.
+
+Run: python benchmarks/reproducibility.py --pairs 1000 --n-jobs 4
+No assertion of universal calibration is made from a single fixed dataset.
+"""
+
+import argparse
+import json
+from copy import deepcopy
+from time import perf_counter
+
+import numpy as np
+from joblib import Parallel, delayed
+from scipy.stats import binomtest
+from sklearn.linear_model import LinearRegression
+from sklearn.model_selection import KFold
+from threadpoolctl import threadpool_limits
+
+from doubleml import DoubleMLData, DoubleMLPLR
+from doubleml.utils import fit_until_reproducible, reproducibility_check
+
+
+def make_model():
+    rng = np.random.RandomState(412)
+    x = rng.normal(size=(81, 4))
+    d = x[:, 0] + rng.normal(size=len(x))
+    y = 0.5 * d + x[:, 1] + rng.normal(size=len(x))
+    data = DoubleMLData.from_arrays(x, y, d)
+    return DoubleMLPLR(data, LinearRegression(), LinearRegression(), n_folds=3, draw_sample_splitting=False)
+
+
+def run_pair(model, pair, tolerance, max_reps):
+    with threadpool_limits(limits=1):
+        start = perf_counter()
+        a, b = [fit_until_reproducible(model, tolerance, max_reps=max_reps, random_state=2 * pair + i) for i in range(2)]
+        return {
+            "exceeded": abs(a["mean"] - b["mean"]) > tolerance,
+            "reps": [a["repetitions"], b["repetitions"]],
+            "exhausted": [a["status"] == "budget_exhausted", b["status"] == "budget_exhausted"],
+            "seconds": perf_counter() - start,
+        }
+
+
+def fixed_run(model, seed, repetitions, tolerance):
+    with threadpool_limits(limits=1):
+        worker = deepcopy(model)
+        splitter = KFold(model.n_folds, shuffle=True, random_state=np.random.RandomState(seed))
+        start = perf_counter()
+        worker.set_sample_splitting([list(splitter.split(np.empty(model.n_obs))) for _ in range(repetitions)])
+        worker.fit(store_predictions=False)
+        report = reproducibility_check(worker.all_coef[0], tolerance)
+        return {"mean": report["mean"], "mcse": report["mcse"], "seconds": perf_counter() - start}
+
+
+def benchmark(pairs, fixed_runs, tolerance, max_reps, n_jobs):
+    model = make_model()
+    start = perf_counter()
+    rows = Parallel(n_jobs=n_jobs)(delayed(run_pair)(model, i, tolerance, max_reps) for i in range(pairs))
+    exceedances = binomtest(sum(r["exceeded"] for r in rows), pairs)
+    interval = exceedances.proportion_ci(method="wilson")
+    nuisance_fits_per_rep = model.n_folds * len(model.learner)
+    result = {
+        "pairs": pairs,
+        "data_seed": 412,
+        "tolerance": tolerance,
+        "error_rate": 0.05,
+        "min_reps": 20,
+        "max_reps": max_reps,
+        "exceedance_rate": exceedances.statistic,
+        "exceedance_interval_95": [float(interval.low), float(interval.high)],
+        "mean_repetitions": float(np.mean([r["reps"] for r in rows])),
+        "budget_exhaustion_rate": float(np.mean([r["exhausted"] for r in rows])),
+        "total_nuisance_fits": int(nuisance_fits_per_rep * np.sum([r["reps"] for r in rows])),
+        "elapsed_seconds": perf_counter() - start,
+        "mean_run_seconds": float(np.mean([r["seconds"] for r in rows]) / 2),
+        "fixed": {},
+    }
+    for count in (10, 50, 100, 500):
+        fixed = Parallel(n_jobs=n_jobs)(delayed(fixed_run)(model, i, count, tolerance) for i in range(fixed_runs))
+        result["fixed"][str(count)] = {
+            "runs": fixed_runs,
+            "mean_mcse": float(np.mean([r["mcse"] for r in fixed])),
+            "mean_run_seconds": float(np.mean([r["seconds"] for r in fixed])),
+            "nuisance_fits_per_run": nuisance_fits_per_rep * count,
+        }
+    return result
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--pairs", type=int, default=1000)
+    parser.add_argument("--fixed-runs", type=int, default=20)
+    parser.add_argument("--tolerance", type=float, default=0.02)
+    parser.add_argument("--max-reps", type=int, default=1000)
+    parser.add_argument("--n-jobs", type=int, default=1)
+    args = parser.parse_args()
+    if args.pairs < 1 or args.fixed_runs < 1:
+        parser.error("pairs and fixed-runs must be positive")
+    print(json.dumps(benchmark(args.pairs, args.fixed_runs, args.tolerance, args.max_reps, args.n_jobs), indent=2))

benchmarks/reproducibility.rst

@@ -0,0 +1,103 @@
+Reproducibility-controlled PLR aggregation
+========================================
+
+The diagnostic operates on one scalar statistic specified before inspecting
+the repetitions. For an existing fitted model::
+
+    from doubleml.utils import reproducibility_check
+
+    report = reproducibility_check(model.all_coef[0], tolerance=0.01)
+
+This reports the arithmetic mean and its split-induced Monte Carlo standard
+error. It does not describe the precision of the median in ``model.coef``.
+The input repetitions must be independent draws from the same partition law;
+repeating one fixed partition does not establish reproducibility.
+
+For automatic repetition, supply a single-treatment, nonclustered
+``DoubleMLPLR`` template constructed with ``n_rep=1``::
+
+    from doubleml.utils import fit_until_reproducible
+
+    report = fit_until_reproducible(
+        model, tolerance=0.01, error_rate=0.05,
+        min_reps=20, max_reps=1000, random_state=42,
+    )
+    print(report["mean"], report["mcse"], report["status"])
+
+The driver fits a private copy on fresh complete partitions. It stores only
+the scalar estimates across repetitions. Existing coefficients, standard
+errors, confidence intervals, and predictions on the template are preserved.
+Common hyperparameters belong on the nuisance learners; fold-specific tuned
+parameters, external predictions, custom scores, and clustered data are
+outside this initial implementation. Set learner random states to integers.
+
+The stopping criterion is MCSE <= tolerance / (sqrt(2) * z), where z is the
+standard normal quantile at 1 - error_rate/2. It is checked after every complete
+repetition following burn-in. ``budget_exhausted`` means the criterion was not
+satisfied. Even ``target_reached`` reports a satisfied numerical criterion,
+not an exact finite-sample probability guarantee. A constant observed sequence
+can conceal rare outcomes, and small stopping counts can be poorly calibrated.
+
+The nominal interpretation is asymptotic as the tolerance decreases, conditional
+on the observed data. An independent pilot can inform the tolerance; the
+minimum repetition count alone cannot establish calibration. No sampling
+standard error or p-value is attached to the new mean.
+
+Validation
+----------
+
+``reproducibility.py`` runs the entire adaptive driver independently twice per
+pair on one fixed synthetic dataset. It reports exceedance frequency, a Wilson
+95% Monte Carlo interval, budget exhaustion, actual repetition counts, nuisance
+fit counts, and wall-clock time. Separate fixed-budget runs at 10, 50, 100,
+and 500 repetitions measure MCSE and cost. They do not claim an empirical
+exceedance rate from only 20 runs.
+
+Run from the repository root::
+
+    python benchmarks/reproducibility.py --pairs 1000 --fixed-runs 20 --n-jobs 3
+    python benchmarks/reproducibility.py --pairs 1000 --fixed-runs 20 --tolerance 0.01 --n-jobs 3
+    pytest doubleml/utils/tests/test_reproducibility.py
+
+Runtime depends on concurrent work and hardware; fit counts are the more
+portable comparison. This is a calibration example, not evidence of uniform
+validity across datasets, nuisance learners, or treatment models.
+
+Reference: Ritzwoller and Romano (2026),
+https://doi.org/10.1257/aer.20231651, Algorithm 1.
+
+Recorded calibration
+--------------------
+
+On the fixed dataset with seed 412, using 1,000 independent pairs for each
+tolerance, nominal error 0.05, burn-in 20, and maximum 1,000 repetitions:
+
+.. list-table:: Fixed-data calibration
+   :header-rows: 1
+
+   * - Tolerance
+     - Mean repeats
+     - Exceedance
+     - 95% Monte Carlo CI
+   * - 0.02
+     - 29.953
+     - 6.0%
+     - 4.69%--7.65%
+   * - 0.01
+     - 122.261
+     - 6.2%
+     - 4.87%--7.87%
+
+There was no budget exhaustion. Both point estimates exceed the nominal 5%;
+both intervals include it. These runs do not establish exact calibration or
+demonstrate a monotone improvement in empirical error with tighter tolerance.
+The tolerances share the same seeds, so the two experiments are correlated.
+
+Raw results, including timings and fixed-budget comparisons, are in
+``results/plr-tolerance-0.02.json`` and ``results/plr-tolerance-0.01.json``.
+The environment used Python 3.12.12, NumPy 2.5.4, SciPy 1.18.1,
+scikit-learn 1.9.1, and DoubleML based on upstream revision
+``469daf260d4ea3520b8265acf17a9d29e800c5e0``. Three worker processes were used.
+
+This pull request includes code written with the assistance of AI.
+`Human review status <https://github.com/DoubleML/doubleml-for-py/pull/429>`_ is tracked in the pull request.

benchmarks/results/plr-tolerance-0.01.json

@@ -0,0 +1,44 @@
+{
+  "pairs": 1000,
+  "data_seed": 412,
+  "tolerance": 0.01,
+  "error_rate": 0.05,
+  "min_reps": 20,
+  "max_reps": 1000,
+  "exceedance_rate": 0.062,
+  "exceedance_interval_95": [
+    0.04866418011998063,
+    0.07868806031336564
+  ],
+  "mean_repetitions": 122.261,
+  "budget_exhaustion_rate": 0.0,
+  "total_nuisance_fits": 1467132,
+  "elapsed_seconds": 561.8347865830001,
+  "mean_run_seconds": 0.8043952413327643,
+  "fixed": {
+    "10": {
+      "runs": 20,
+      "mean_mcse": 0.011677262580110271,
+      "mean_run_seconds": 0.057370227102364876,
+      "nuisance_fits_per_run": 60
+    },
+    "50": {
+      "runs": 20,
+      "mean_mcse": 0.0055360289057134245,
+      "mean_run_seconds": 0.3992237104976084,
+      "nuisance_fits_per_run": 300
+    },
+    "100": {
+      "runs": 20,
+      "mean_mcse": 0.0039995503298418085,
+      "mean_run_seconds": 1.0490305750994593,
+      "nuisance_fits_per_run": 600
+    },
+    "500": {
+      "runs": 20,
+      "mean_mcse": 0.001794004054532988,
+      "mean_run_seconds": 5.766709445748711,
+      "nuisance_fits_per_run": 3000
+    }
+  }
+}

benchmarks/results/plr-tolerance-0.02.json

@@ -0,0 +1,44 @@
+{
+  "pairs": 1000,
+  "data_seed": 412,
+  "tolerance": 0.02,
+  "error_rate": 0.05,
+  "min_reps": 20,
+  "max_reps": 1000,
+  "exceedance_rate": 0.06,
+  "exceedance_interval_95": [
+    0.04689647267593195,
+    0.07647107479135656
+  ],
+  "mean_repetitions": 29.953,
+  "budget_exhaustion_rate": 0.0,
+  "total_nuisance_fits": 359436,
+  "elapsed_seconds": 150.95566583398613,
+  "mean_run_seconds": 0.20873066854466743,
+  "fixed": {
+    "10": {
+      "runs": 20,
+      "mean_mcse": 0.011677262580110271,
+      "mean_run_seconds": 0.052713906348799355,
+      "nuisance_fits_per_run": 60
+    },
+    "50": {
+      "runs": 20,
+      "mean_mcse": 0.0055360289057134245,
+      "mean_run_seconds": 0.33320674165443054,
+      "nuisance_fits_per_run": 300
+    },
+    "100": {
+      "runs": 20,
+      "mean_mcse": 0.0039995503298418085,
+      "mean_run_seconds": 0.6555121250465163,
+      "nuisance_fits_per_run": 600
+    },
+    "500": {
+      "runs": 20,
+      "mean_mcse": 0.001794004054532988,
+      "mean_run_seconds": 3.975092010544904,
+      "nuisance_fits_per_run": 3000
+    }
+  }
+}

benchmarks/tests/test_reproducibility_benchmark.py

@@ -0,0 +1,54 @@
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+from sklearn.linear_model import LinearRegression
+
+from doubleml import DoubleMLData, DoubleMLPLR
+
+spec = importlib.util.spec_from_file_location("plr_benchmark", Path(__file__).parents[1] / "reproducibility.py")
+benchmark = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(benchmark)
+
+
+@pytest.fixture
+def model():
+    rng = np.random.RandomState(12)
+    data = DoubleMLData.from_arrays(rng.normal(size=(53, 3)), rng.normal(size=53), rng.normal(size=53))
+    return DoubleMLPLR(data, LinearRegression(), LinearRegression(), n_folds=4, draw_sample_splitting=False)
+
+
+def test_fixed_run_uses_model_dimensions(model, monkeypatch):
+    calls = []
+    fit = DoubleMLPLR.fit
+
+    def counted_fit(self, *args, **kwargs):
+        calls.append((self.n_obs, self.n_folds, self.n_rep))
+        return fit(self, *args, **kwargs)
+
+    monkeypatch.setattr(DoubleMLPLR, "fit", counted_fit)
+    result = benchmark.fixed_run(model, seed=12, repetitions=3, tolerance=0.01)
+    assert calls == [(53, 4, 3)]
+    assert np.isfinite(result["mean"])
+    assert np.isfinite(result["mcse"])
+
+
+@pytest.mark.parametrize(
+    "exceeded, interval",
+    [(False, [0, 0.2775327998628892]), (True, [0.7224672001371107, 1])],
+)
+def test_benchmark_interval_and_work_count(model, exceeded, interval, monkeypatch):
+    def run_pair(*args):
+        return {"exceeded": exceeded, "reps": [20, 20], "exhausted": [False, False], "seconds": 1}
+
+    monkeypatch.setattr(benchmark, "make_model", lambda: model)
+    monkeypatch.setattr(benchmark, "run_pair", run_pair)
+    monkeypatch.setattr(benchmark, "fixed_run", lambda *args: {"mean": 0, "mcse": 0.01, "seconds": 1})
+    report = benchmark.benchmark(pairs=10, fixed_runs=1, tolerance=0.01, max_reps=20, n_jobs=1)
+    assert report["exceedance_rate"] == float(exceeded)
+    assert_allclose(report["exceedance_interval_95"], interval, atol=1e-15)
+    assert report["total_nuisance_fits"] == 3200
+    for count, result in report["fixed"].items():
+        assert result["nuisance_fits_per_run"] == 8 * int(count)

doc/api/api.rst

@@ -1,6 +1,17 @@
 Python API reference
 ====================
 
+Reproducibility diagnostics
+---------------------------
+
+.. currentmodule:: doubleml.utils
+
+.. autosummary::
+    :toctree: generated/
+
+    reproducibility_check
+    fit_until_reproducible
+
 Double machine learning data class
 ----------------------------------
 

doubleml/utils/__init__.py

@@ -9,6 +9,7 @@
 from .global_learner import GlobalClassifier, GlobalRegressor
 from .policytree import DoubleMLPolicyTree
 from .propensity_score_processing import PSProcessor, PSProcessorConfig
+from .reproducibility import fit_until_reproducible, reproducibility_check
 from .resampling import DoubleMLClusterResampling, DoubleMLResampling
 
 __all__ = [
@@ -24,4 +25,6 @@
     "GlobalRegressor",
     "PSProcessor",
     "PSProcessorConfig",
+    "reproducibility_check",
+    "fit_until_reproducible",
 ]