diff --git a/.gitignore b/.gitignore index bda36077..57faec11 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,4 @@ doubleml/_version.py # uv .venv/ uv.lock +doc/api/generated/ diff --git a/benchmarks/reproducibility.py b/benchmarks/reproducibility.py new file mode 100644 index 00000000..f4edf4ba --- /dev/null +++ b/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)) diff --git a/benchmarks/reproducibility.rst b/benchmarks/reproducibility.rst new file mode 100644 index 00000000..8532336e --- /dev/null +++ b/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 `_ is tracked in the pull request. diff --git a/benchmarks/results/plr-tolerance-0.01.json b/benchmarks/results/plr-tolerance-0.01.json new file mode 100644 index 00000000..944d90a1 --- /dev/null +++ b/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 + } + } +} diff --git a/benchmarks/results/plr-tolerance-0.02.json b/benchmarks/results/plr-tolerance-0.02.json new file mode 100644 index 00000000..b2d3be3c --- /dev/null +++ b/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 + } + } +} diff --git a/benchmarks/tests/test_reproducibility_benchmark.py b/benchmarks/tests/test_reproducibility_benchmark.py new file mode 100644 index 00000000..7c4963c8 --- /dev/null +++ b/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) diff --git a/doc/api/api.rst b/doc/api/api.rst index 65175981..5c9df66d 100644 --- a/doc/api/api.rst +++ b/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 ---------------------------------- diff --git a/doubleml/utils/__init__.py b/doubleml/utils/__init__.py index 868429da..395b33bd 100644 --- a/doubleml/utils/__init__.py +++ b/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", ] diff --git a/doubleml/utils/reproducibility.py b/doubleml/utils/reproducibility.py new file mode 100644 index 00000000..0744d35f --- /dev/null +++ b/doubleml/utils/reproducibility.py @@ -0,0 +1,184 @@ +"""Conditional Monte Carlo diagnostics for means across independent partitions.""" + +from copy import deepcopy +from numbers import Integral, Real + +import numpy as np +from scipy.stats import norm +from sklearn.model_selection import KFold +from sklearn.utils import check_random_state + + +def _check_target(tolerance, error_rate): + for name, value in (("tolerance", tolerance), ("error_rate", error_rate)): + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{name} must be a real number.") + if not np.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be finite and positive.") + if error_rate >= 1: + raise ValueError("error_rate must be less than 1.") + threshold = tolerance / (np.sqrt(2) * norm.isf(error_rate / 2)) + if not np.isfinite(threshold) or threshold <= 0: + raise ValueError("The requested target is outside floating-point precision.") + return threshold + + +def reproducibility_check(values, tolerance, error_rate=0.05): + """Assess split-induced Monte Carlo error of a scalar mean. + + Parameters + ---------- + values : array-like of shape (n_rep,) + At least two finite scalar estimates from independent, identically + distributed partitions of the same data. For a fitted DoubleML model, + use ``model.all_coef[index, :]`` for one prespecified coefficient. + Independence cannot be checked from these values. + tolerance : float + Positive absolute tolerance for the difference between two independent + mean aggregates, in the units of the coefficient. + error_rate : float, default=0.05 + Nominal probability of exceeding the tolerance, strictly between 0 and 1. + + Returns + ------- + report : dict + ``mean``, ``mcse`` (sample standard deviation divided by sqrt(n_rep)), + ``threshold`` (tolerance / (sqrt(2) * normal critical value)), + ``criterion_met``, ``repetitions``, ``tolerance``, and ``error_rate``. + + Notes + ----- + This is a plug-in diagnostic for the arithmetic mean, not DoubleML's + median coefficient or its sampling standard error, confidence interval, + or p-value. A satisfied criterion is not a finite-sample probability + guarantee. In particular, zero observed variance need not imply zero + conditional variance. Repeatedly inspecting this diagnostic alone does + not implement a complete sequential procedure. + + The stopping rule in Algorithm 1 of Ritzwoller and Romano [1]_ uses this + criterion after a prespecified burn-in. Its reproducibility interpretation + is asymptotic as the tolerance decreases, conditional on the observed data. + + References + ---------- + .. [1] Ritzwoller, D. M. and Romano, J. P. (2026), Reproducible Aggregation + of Sample-Split Statistics. https://doi.org/10.1257/aer.20231651 + + Examples + -------- + >>> from doubleml.utils import reproducibility_check + >>> report = reproducibility_check([0.10, 0.12, 0.11], tolerance=0.02) + >>> report["criterion_met"] + True + """ + threshold = _check_target(tolerance, error_rate) + values = np.asarray(values) + if values.ndim != 1 or values.size < 2 or values.dtype.kind not in "fiu": + raise ValueError("values must be a one-dimensional numeric array with at least two estimates.") + values = values.astype(float) + if not np.isfinite(values).all(): + raise ValueError("values must contain only finite estimates.") + with np.errstate(over="ignore", invalid="ignore"): + mean = float(np.mean(values)) + mcse = float(np.std(values, ddof=1) / np.sqrt(values.size)) + if not np.isfinite(mean) or not np.isfinite(mcse): + raise ValueError("The estimates overflow the mean or variance calculation; rescale them.") + return { + "mean": mean, + "mcse": mcse, + "threshold": float(threshold), + "criterion_met": bool(mcse <= threshold), + "repetitions": values.size, + "tolerance": float(tolerance), + "error_rate": float(error_rate), + } + + +def fit_until_reproducible(model, tolerance, error_rate=0.05, min_reps=20, max_reps=1000, random_state=None, n_jobs_cv=None): + """Sequentially average one PLR coefficient over fresh K-fold partitions. + + Parameters + ---------- + model : DoubleMLPLR + Template with one treatment, independent observations, ``n_rep=1``, + at least two folds, and a built-in score. Existing splits and fitted + results are not reused. Fold-specific tuned parameters are unsupported; + set common hyperparameters on the learners instead. All learner + ``random_state`` parameters must be integers so learner randomness is + fixed and only partition randomness is assessed. Custom learners must + also be deterministic conditional on their inputs; this cannot be + established automatically from their parameter names. + tolerance : float + Positive absolute tolerance for the scalar arithmetic mean. + error_rate : float, default=0.05 + Nominal reproducibility error, strictly between zero and one. + min_reps : int, default=20 + Prespecified burn-in, at least two. This is not a calibration guarantee. + max_reps : int, default=1000 + Maximum number of complete repetitions, at least ``min_reps``. + random_state : int, RandomState instance or None, default=None + Random state used to generate successive independent shuffled K-fold + partitions. An integer makes the partition sequence reproducible. + n_jobs_cv : int or None, default=None + Number of parallel jobs for fitting nuisance learners within each fold. + + Returns + ------- + report : dict + Fields from :func:`reproducibility_check`, plus ``status`` (either + ``'target_reached'`` or ``'budget_exhausted'``) and ``values``, the + per-repetition coefficient estimates. No sampling inference is attached + to this mean. The input model and its existing inference are unchanged. + + Notes + ----- + Fits a private copy once per complete partition and checks after every + repetition following burn-in. Earlier repetitions are never refitted. + Memory for nuisance estimates is bounded by one repetition. This first + implementation deliberately excludes clustered data, multiple treatments, + custom scores, and externally supplied predictions or partitions. + + The nominal error is asymptotic, not a finite-sample guarantee; see + :func:`reproducibility_check`. Choose the tolerance before running the + algorithm, preferably informed by an independent pilot. Small repetition + counts can be poorly calibrated. Budget exhaustion does not certify the + target. A diagnostic for the mean does not certify median-based inference. + """ + from doubleml.plm.plr import DoubleMLPLR + + _check_target(tolerance, error_rate) + for name, value in (("min_reps", min_reps), ("max_reps", max_reps)): + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral): + raise TypeError(f"{name} must be an integer.") + if min_reps < 2 or max_reps < min_reps: + raise ValueError("Require 2 <= min_reps <= max_reps.") + if not isinstance(model, DoubleMLPLR): + raise TypeError("model must be a DoubleMLPLR instance.") + if model._is_cluster_data or model._dml_data.n_treat != 1 or model.n_rep != 1 or model.n_folds < 2: + raise ValueError("Require independent observations, one treatment, n_rep=1, and n_folds >= 2.") + if not isinstance(model.score, str): + raise ValueError("Only built-in PLR scores are supported.") + if any(value is not None for params in model.params.values() for reps in params.values() for value in reps): + raise ValueError("Fold-specific parameters are unsupported; configure the learners directly.") + for learner in model.learner.values(): + for name, value in learner.get_params(deep=True).items(): + if name.split("__")[-1] == "random_state" and (isinstance(value, bool) or not isinstance(value, Integral)): + raise ValueError("Set every learner random_state to an integer to isolate partition randomness.") + rng = check_random_state(random_state) + worker = deepcopy(model) + splitter = KFold(n_splits=model.n_folds, shuffle=True, random_state=rng) + values = [] + for _ in range(max_reps): + worker.set_sample_splitting(list(splitter.split(np.empty(model._dml_data.n_obs)))) + worker.fit(n_jobs_cv=n_jobs_cv, store_predictions=False, store_models=False) + value = float(worker.all_coef[0, 0]) + if not np.isfinite(value): + raise ValueError("A repetition produced a non-finite coefficient.") + values.append(value) + if len(values) >= min_reps: + report = reproducibility_check(values, tolerance, error_rate) + if report["criterion_met"]: + break + report["status"] = "target_reached" if report["criterion_met"] else "budget_exhausted" + report["values"] = np.array(values) + return report diff --git a/doubleml/utils/tests/test_reproducibility.py b/doubleml/utils/tests/test_reproducibility.py new file mode 100644 index 00000000..98cbb8ba --- /dev/null +++ b/doubleml/utils/tests/test_reproducibility.py @@ -0,0 +1,206 @@ +from copy import deepcopy + +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal +from pandas.testing import assert_frame_equal +from scipy.stats import norm +from sklearn.ensemble import RandomForestRegressor +from sklearn.linear_model import LinearRegression +from sklearn.model_selection import KFold + +from doubleml import DoubleMLData, DoubleMLPLR +from doubleml.utils import fit_until_reproducible, reproducibility_check + +pytestmark = pytest.mark.ci + + +@pytest.fixture +def model(): + rng = np.random.RandomState(412) + x = rng.normal(size=(81, 4)) + d = x[:, 0] + rng.normal(size=81) + y = 0.5 * d + x[:, 1] + rng.normal(size=81) + data = DoubleMLData.from_arrays(x, y, d) + return DoubleMLPLR(data, LinearRegression(), LinearRegression(), n_folds=3, draw_sample_splitting=False) + + +def test_numerical_oracle(): + result = reproducibility_check([1, 2, 3], 2, 0.05) + assert result["mean"] == 2 + assert_allclose(result["mcse"], np.sqrt(1 / 3)) + assert_allclose(result["threshold"], 2 / (np.sqrt(2) * norm.isf(0.025))) + assert result["criterion_met"] + assert result["repetitions"] == 3 + assert reproducibility_check([4, 4], 0.01)["mcse"] == 0 + assert not reproducibility_check([0, 0, 1e6], 0.01)["criterion_met"] + + +@pytest.mark.parametrize("values", [[], [1], [[1, 2]], [1, np.nan], [1, np.inf], [1j, 2j], ["1", "2"], [True, False]]) +def test_invalid_values(values): + with pytest.raises(ValueError): + reproducibility_check(values, 0.1) + + +@pytest.mark.parametrize("name", ["tolerance", "error_rate"]) +@pytest.mark.parametrize("value", [0, -1, np.inf, np.nan, True, "0.1", None]) +def test_invalid_targets(name, value): + kwargs = {"tolerance": 0.1, "error_rate": 0.05, name: value} + with pytest.raises((ValueError, TypeError)): + reproducibility_check([1, 2], **kwargs) + + +def test_extreme_targets_and_overflow(): + with pytest.raises(ValueError, match="less than 1"): + reproducibility_check([1, 2], 0.1, 1) + assert np.isfinite(reproducibility_check([1, 2], 0.1, 1e-20)["threshold"]) + with pytest.raises(ValueError, match="floating-point precision"): + reproducibility_check([1, 2], 0.1, np.nextafter(0.0, 1.0)) + with pytest.raises(ValueError, match="overflow"): + reproducibility_check([-1e308, 1e308], 0.1) + + +def test_matches_fixed_repetitions_without_mutation(model): + rng = np.random.RandomState(17) + splits = [list(KFold(3, shuffle=True, random_state=rng).split(np.empty(81))) for _ in range(5)] + model.set_sample_splitting(splits[:1]).fit() + before = deepcopy(model) + fixed = deepcopy(model) + fixed.set_sample_splitting(splits).fit() + report = fit_until_reproducible(model, 1e-14, min_reps=2, max_reps=5, random_state=17) + assert report["status"] == "budget_exhausted" + assert_allclose(report["values"], fixed.all_coef[0], rtol=1e-12) + assert_allclose(report["mean"], fixed.all_coef.mean()) + assert_array_equal(model.all_coef, before.all_coef) + assert_array_equal(model.se, before.se) + assert_array_equal(model.pval, before.pval) + assert_frame_equal(model.summary, before.summary) + assert_array_equal(model.predictions["ml_l"], before.predictions["ml_l"]) + for actual, expected in zip(model.smpls[0], before.smpls[0]): + assert_array_equal(actual[0], expected[0]) + assert_array_equal(actual[1], expected[1]) + + +def test_minimum_and_fit_count(model, monkeypatch): + original_fit = LinearRegression.fit + calls = [] + + def counted(self, x, y, **kwargs): + calls.append(len(y)) + return original_fit(self, x, y, **kwargs) + + monkeypatch.setattr(LinearRegression, "fit", counted) + report = fit_until_reproducible(model, 100, min_reps=4, max_reps=9, random_state=10) + assert report["status"] == "target_reached" + assert report["repetitions"] == 4 + assert len(calls) == 4 * 3 * 2 + assert report["mcse"] <= report["threshold"] + + +def test_repeatability_and_tighter_target(model): + loose = fit_until_reproducible(model, 1, min_reps=3, max_reps=8, random_state=10) + tight = fit_until_reproducible(model, 1e-12, min_reps=3, max_reps=8, random_state=10) + again = fit_until_reproducible(model, 1e-12, min_reps=3, max_reps=8, random_state=10) + assert loose["repetitions"] < tight["repetitions"] + assert_array_equal(loose["values"], tight["values"][:3]) + assert_array_equal(tight["values"], again["values"]) + + +@pytest.mark.parametrize("kwargs", [{"min_reps": 1}, {"min_reps": True}, {"max_reps": 2}, {"max_reps": 2.5}]) +def test_invalid_budget(model, kwargs): + with pytest.raises((TypeError, ValueError)): + fit_until_reproducible(model, 0.1, **kwargs) + + +def test_plr_subclass(model): + class PLRSubclass(DoubleMLPLR): + pass + + subclass = PLRSubclass(model._dml_data, LinearRegression(), LinearRegression(), n_folds=3, draw_sample_splitting=False) + expected = fit_until_reproducible(model, 0.1, min_reps=3, max_reps=3, random_state=12) + actual = fit_until_reproducible(subclass, 0.1, min_reps=3, max_reps=3, random_state=12) + assert_array_equal(actual["values"], expected["values"]) + + +def test_unsupported_models(model): + with pytest.raises(TypeError, match="DoubleMLPLR"): + fit_until_reproducible(object(), 0.1) + model._n_rep = 2 + with pytest.raises(ValueError, match="n_rep=1"): + fit_until_reproducible(model, 0.1) + + +def test_reject_fold_params_and_unseeded_learner(model): + model.set_ml_nuisance_params("ml_l", "d", {"fit_intercept": False}) + with pytest.raises(ValueError, match="Fold-specific"): + fit_until_reproducible(model, 0.1) + model = DoubleMLPLR(model._dml_data, RandomForestRegressor(), LinearRegression()) + with pytest.raises(ValueError, match="random_state"): + fit_until_reproducible(model, 0.1) + + +@pytest.mark.parametrize("attribute,value", [("_is_cluster_data", True), ("_n_folds", 1), ("_score", lambda *args: None)]) +def test_unsupported_configuration(model, attribute, value): + setattr(model, attribute, value) + with pytest.raises(ValueError): + fit_until_reproducible(model, 0.1) + + +def test_multiple_treatments_rejected(model): + data = model._dml_data + multiple = DoubleMLData.from_arrays(data.x, data.y, np.column_stack([data.d, data.x[:, 2]])) + model = DoubleMLPLR(multiple, LinearRegression(), LinearRegression()) + with pytest.raises(ValueError, match="one treatment"): + fit_until_reproducible(model, 0.1) + + +def test_nonfinite_fit_fails_immediately(model, monkeypatch): + calls = [] + + def invalid_fit(self, **kwargs): + calls.append(1) + self._all_coef[:] = np.nan + return self + + monkeypatch.setattr(DoubleMLPLR, "fit", invalid_fit) + with pytest.raises(ValueError, match="non-finite coefficient"): + fit_until_reproducible(model, 0.1) + assert len(calls) == 1 + + +def test_iv_type_score(model): + template = DoubleMLPLR( + model._dml_data, + LinearRegression(), + LinearRegression(), + ml_g=LinearRegression(), + score="IV-type", + n_folds=3, + draw_sample_splitting=False, + ) + report = fit_until_reproducible(template, 1, min_reps=3, max_reps=3, random_state=7) + assert report["status"] == "target_reached" + assert np.isfinite(report["values"]).all() + + +def test_report_contract(model): + report = fit_until_reproducible(model, 1, min_reps=3, max_reps=3, random_state=7) + assert isinstance(report, dict) + assert set(report) == { + "mean", + "mcse", + "threshold", + "criterion_met", + "repetitions", + "tolerance", + "error_rate", + "status", + "values", + } + for name in ("mean", "mcse", "threshold", "tolerance", "error_rate"): + assert isinstance(report[name], float) + assert isinstance(report["criterion_met"], bool) + assert isinstance(report["repetitions"], int) + assert isinstance(report["values"], np.ndarray) + assert report["values"].shape == (report["repetitions"],) + assert report["values"].dtype.kind == "f"