From 78f97b9b000d6e0e2e1ffa783a0e7f7bc20b22ca Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Jul 2026 11:08:44 -0400 Subject: [PATCH 1/2] chore(lint): pin tooling + mechanical source normalization pyproject: exact pins (black==26.3.1, ruff==0.15.13, mypy==2.1.0) in the dev extra + scoped per-file-ignores for deliberate patterns. diff_diff/ changes are the verbatim output of `ruff check diff_diff --fix` + `black diff_diff` at the pinned versions. mypy error count verified unchanged (184) pre/post. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- CHANGELOG.md | 7 ++ diff_diff/__init__.py | 235 +++++++++++++++++++------------------- diff_diff/_backend.py | 56 +++++++-- diff_diff/datasets.py | 1 - diff_diff/honest_did.py | 4 +- diff_diff/local_linear.py | 13 +-- diff_diff/triple_diff.py | 4 +- diff_diff/trop.py | 2 +- diff_diff/trop_local.py | 1 - pyproject.toml | 26 ++++- 10 files changed, 198 insertions(+), 151 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a08852362..9f6ca0207 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- **Internal: repo-wide lint normalization + pinned tooling.** black/ruff/mypy are now + pinned exactly in the `dev` extra (`black==26.3.1`, `ruff==0.15.13`, `mypy==2.1.0`; + the tools require Python >= 3.10 — the library floor stays 3.9); full `black` + + `ruff --fix` pass over source and tests; scoped `per-file-ignores` for deliberate + patterns (trop logger-before-imports E402, honest_did math-notation E741, `__init__` + re-export F401); ~24 audited unused-local test fixes, two upgraded to real assertions. + No public API or numerical behavior change. - **Internal: cleaned package-source Ruff static-analysis findings** across forward-reference, stale f-string, invalid `noqa`, and audited unused-local sites. No public API or numerical behavior change. diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index df9d57991..68e9de64b 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -28,13 +28,48 @@ _rust_project_simplex, _rust_solve_ols, ) - +from diff_diff._guides_api import get_llm_guide +from diff_diff.agent_workflow import agent_workflow from diff_diff.bacon import ( BaconDecomposition, BaconDecompositionResults, Comparison2x2, bacon_decompose, ) +from diff_diff.business_report import ( + BUSINESS_REPORT_SCHEMA_VERSION, + BusinessContext, + BusinessReport, +) +from diff_diff.chaisemartin_dhaultfoeuille import ( + ChaisemartinDHaultfoeuille, + TWFEWeightsResult, + chaisemartin_dhaultfoeuille, + twowayfeweights, +) +from diff_diff.chaisemartin_dhaultfoeuille_results import ( + ChaisemartinDHaultfoeuilleResults, + DCDHBootstrapResults, +) +from diff_diff.continuous_did import ( + ContinuousDiD, + ContinuousDiDResults, + DoseResponseCurve, +) +from diff_diff.datasets import ( + clear_cache, + list_datasets, + load_card_krueger, + load_castle_doctrine, + load_dataset, + load_divorce_laws, + load_mpdta, +) +from diff_diff.diagnostic_report import ( + DIAGNOSTIC_REPORT_SCHEMA_VERSION, + DiagnosticReport, + DiagnosticReportResults, +) from diff_diff.diagnostics import ( PlaceboTestResults, leave_one_out_test, @@ -44,22 +79,16 @@ run_all_placebo_tests, run_placebo_test, ) -from diff_diff.linalg import ( - InferenceResult, - LinearRegression, +from diff_diff.efficient_did import ( + EDiDBootstrapResults, + EfficientDiD, + EfficientDiDResults, ) -from diff_diff.local_linear import ( - KERNELS, - BandwidthResult, - BiasCorrectedFit, - LocalLinearFit, - bias_corrected_local_linear, - epanechnikov_kernel, - kernel_moments, - local_linear_fit, - mse_optimal_bandwidth, - triangular_kernel, - uniform_kernel, +from diff_diff.estimators import ( + DifferenceInDifferences, + MultiPeriodDiD, + SyntheticDiD, + TwoWayFixedEffects, ) from diff_diff.had import ( HeterogeneousAdoptionDiD, @@ -80,12 +109,6 @@ stute_test, yatchew_hr_test, ) -from diff_diff.estimators import ( - DifferenceInDifferences, - MultiPeriodDiD, - SyntheticDiD, - TwoWayFixedEffects, -) from diff_diff.honest_did import ( DeltaRM, DeltaSD, @@ -96,6 +119,31 @@ compute_honest_did, sensitivity_plot, ) +from diff_diff.imputation import ( + ImputationBootstrapResults, + ImputationDiD, + ImputationDiDResults, + imputation_did, +) +from diff_diff.linalg import ( + InferenceResult, + LinearRegression, +) +from diff_diff.local_linear import ( + KERNELS, + BandwidthResult, + BiasCorrectedFit, + LocalLinearFit, + bias_corrected_local_linear, + epanechnikov_kernel, + kernel_moments, + local_linear_fit, + mse_optimal_bandwidth, + triangular_kernel, + uniform_kernel, +) +from diff_diff.lpdid import LPDiD +from diff_diff.lpdid_results import LPDiDResults from diff_diff.power import ( PowerAnalysis, PowerResults, @@ -110,22 +158,16 @@ simulate_power, simulate_sample_size, ) -from diff_diff.pretrends import ( - PreTrendsPower, - PreTrendsPowerCurve, - PreTrendsPowerResults, - compute_mdv, - compute_pretrends_power, -) +from diff_diff.practitioner import practitioner_next_steps from diff_diff.prep import ( aggregate_survey, aggregate_to_cohorts, balance_panel, create_event_time, generate_continuous_did_data, - generate_did_data, generate_ddd_data, generate_ddd_panel_data, + generate_did_data, generate_event_study_data, generate_factor_data, generate_panel_data, @@ -142,55 +184,40 @@ validate_did_data, wide_to_long, ) +from diff_diff.pretrends import ( + PreTrendsPower, + PreTrendsPowerCurve, + PreTrendsPowerResults, + compute_mdv, + compute_pretrends_power, +) +from diff_diff.profile import ( + Alert, + OutcomeShape, + PanelProfile, + TreatmentDoseShape, + profile_panel, +) from diff_diff.results import ( DiDResults, MultiPeriodDiDResults, PeriodEffect, + SpilloverDiDResults, # re-export SyntheticDiDResults, ) -from diff_diff.survey import ( - DEFFDiagnostics, - SurveyDesign, - SurveyMetadata, - compute_deff_diagnostics, - make_pweight_design, -) -from diff_diff.staggered import ( - CallawaySantAnna, - CallawaySantAnnaResults, - CSBootstrapResults, - GroupTimeEffect, -) -from diff_diff.imputation import ( - ImputationBootstrapResults, - ImputationDiD, - ImputationDiDResults, - imputation_did, -) -from diff_diff.two_stage import ( - TwoStageBootstrapResults, - TwoStageDiD, - TwoStageDiDResults, - two_stage_did, -) from diff_diff.spillover import ( SpilloverDiD, ) -from diff_diff.results import SpilloverDiDResults # re-export from diff_diff.stacked_did import ( StackedDiD, StackedDiDResults, stacked_did, ) -from diff_diff.sun_abraham import ( - SABootstrapResults, - SunAbraham, - SunAbrahamResults, -) -from diff_diff.triple_diff import ( - TripleDifference, - TripleDifferenceResults, - triple_difference, +from diff_diff.staggered import ( + CallawaySantAnna, + CallawaySantAnnaResults, + CSBootstrapResults, + GroupTimeEffect, ) from diff_diff.staggered_triple_diff import ( StaggeredTripleDifference, @@ -198,40 +225,39 @@ from diff_diff.staggered_triple_diff_results import ( StaggeredTripleDiffResults, ) -from diff_diff.continuous_did import ( - ContinuousDiD, - ContinuousDiDResults, - DoseResponseCurve, +from diff_diff.sun_abraham import ( + SABootstrapResults, + SunAbraham, + SunAbrahamResults, ) -from diff_diff.efficient_did import ( - EfficientDiD, - EfficientDiDResults, - EDiDBootstrapResults, +from diff_diff.survey import ( + DEFFDiagnostics, + SurveyDesign, + SurveyMetadata, + compute_deff_diagnostics, + make_pweight_design, ) -from diff_diff.chaisemartin_dhaultfoeuille import ( - ChaisemartinDHaultfoeuille, - TWFEWeightsResult, - chaisemartin_dhaultfoeuille, - twowayfeweights, +from diff_diff.synthetic_control import ( + SyntheticControl, + synthetic_control, ) -from diff_diff.chaisemartin_dhaultfoeuille_results import ( - ChaisemartinDHaultfoeuilleResults, - DCDHBootstrapResults, +from diff_diff.synthetic_control_results import SyntheticControlResults +from diff_diff.triple_diff import ( + TripleDifference, + TripleDifferenceResults, + triple_difference, ) from diff_diff.trop import ( TROP, TROPResults, trop, ) -from diff_diff.synthetic_control import ( - SyntheticControl, - synthetic_control, +from diff_diff.two_stage import ( + TwoStageBootstrapResults, + TwoStageDiD, + TwoStageDiDResults, + two_stage_did, ) -from diff_diff.synthetic_control_results import SyntheticControlResults -from diff_diff.wooldridge import WooldridgeDiD -from diff_diff.wooldridge_results import WooldridgeDiDResults -from diff_diff.lpdid import LPDiD -from diff_diff.lpdid_results import LPDiDResults from diff_diff.utils import ( WildBootstrapResults, check_parallel_trends, @@ -252,35 +278,8 @@ plot_staircase, plot_synth_weights, ) -from diff_diff.practitioner import practitioner_next_steps -from diff_diff.business_report import ( - BUSINESS_REPORT_SCHEMA_VERSION, - BusinessContext, - BusinessReport, -) -from diff_diff.diagnostic_report import ( - DIAGNOSTIC_REPORT_SCHEMA_VERSION, - DiagnosticReport, - DiagnosticReportResults, -) -from diff_diff._guides_api import get_llm_guide -from diff_diff.agent_workflow import agent_workflow -from diff_diff.profile import ( - Alert, - OutcomeShape, - PanelProfile, - TreatmentDoseShape, - profile_panel, -) -from diff_diff.datasets import ( - clear_cache, - list_datasets, - load_card_krueger, - load_castle_doctrine, - load_dataset, - load_divorce_laws, - load_mpdta, -) +from diff_diff.wooldridge import WooldridgeDiD +from diff_diff.wooldridge_results import WooldridgeDiDResults # Estimator aliases — short names for convenience DiD = DifferenceInDifferences diff --git a/diff_diff/_backend.py b/diff_diff/_backend.py index 2d379d8c7..346684601 100644 --- a/diff_diff/_backend.py +++ b/diff_diff/_backend.py @@ -18,27 +18,59 @@ # Try to import Rust backend for accelerated operations try: from diff_diff._rust_backend import ( - generate_bootstrap_weights_batch as _rust_bootstrap_weights, - project_simplex as _rust_project_simplex, - solve_ols as _rust_solve_ols, + bootstrap_trop_variance as _rust_bootstrap_trop_variance, + ) + from diff_diff._rust_backend import ( + bootstrap_trop_variance_global as _rust_bootstrap_trop_variance_global, + ) + from diff_diff._rust_backend import ( + compute_noise_level as _rust_compute_noise_level, + ) + from diff_diff._rust_backend import ( compute_robust_vcov as _rust_compute_robust_vcov, + ) + from diff_diff._rust_backend import ( + # SDID weights (Frank-Wolfe matching R's synthdid) + compute_sdid_unit_weights as _rust_sdid_unit_weights, + ) + from diff_diff._rust_backend import ( + compute_time_weights as _rust_compute_time_weights, + ) + from diff_diff._rust_backend import ( # TROP estimator acceleration (local method) compute_unit_distance_matrix as _rust_unit_distance_matrix, + ) + from diff_diff._rust_backend import ( + generate_bootstrap_weights_batch as _rust_bootstrap_weights, + ) + from diff_diff._rust_backend import ( loocv_grid_search as _rust_loocv_grid_search, - bootstrap_trop_variance as _rust_bootstrap_trop_variance, + ) + from diff_diff._rust_backend import ( # TROP estimator acceleration (global method) loocv_grid_search_global as _rust_loocv_grid_search_global, - bootstrap_trop_variance_global as _rust_bootstrap_trop_variance_global, - # SDID weights (Frank-Wolfe matching R's synthdid) - compute_sdid_unit_weights as _rust_sdid_unit_weights, - compute_time_weights as _rust_compute_time_weights, - compute_noise_level as _rust_compute_noise_level, + ) + from diff_diff._rust_backend import ( + project_simplex as _rust_project_simplex, + ) + from diff_diff._rust_backend import ( + # Diagnostics + rust_backend_info as _rust_backend_info, + ) + from diff_diff._rust_backend import ( sc_weight_fw as _rust_sc_weight_fw, - sc_weight_fw_with_convergence as _rust_sc_weight_fw_with_convergence, + ) + from diff_diff._rust_backend import ( sc_weight_fw_weighted as _rust_sc_weight_fw_weighted, + ) + from diff_diff._rust_backend import ( sc_weight_fw_weighted_with_convergence as _rust_sc_weight_fw_weighted_with_convergence, - # Diagnostics - rust_backend_info as _rust_backend_info, + ) + from diff_diff._rust_backend import ( + sc_weight_fw_with_convergence as _rust_sc_weight_fw_with_convergence, + ) + from diff_diff._rust_backend import ( + solve_ols as _rust_solve_ols, ) _rust_available = True diff --git a/diff_diff/datasets.py b/diff_diff/datasets.py index d3fe55601..98a14546b 100644 --- a/diff_diff/datasets.py +++ b/diff_diff/datasets.py @@ -17,7 +17,6 @@ import numpy as np import pandas as pd - # Cache directory for downloaded datasets _CACHE_DIR = Path.home() / ".cache" / "diff_diff" / "datasets" diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index c9004352c..ea6ec7bab 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -17,8 +17,8 @@ https://github.com/asheshrambachan/HonestDiD - R package implementation """ -from dataclasses import dataclass, field import warnings +from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Tuple, Union import numpy as np @@ -2194,7 +2194,7 @@ def _compute_arp_test( reject : bool True if H0 is rejected. """ - from scipy.stats import norm, truncnorm + from scipy.stats import truncnorm if kappa is None: kappa = alpha / 10.0 diff --git a/diff_diff/local_linear.py b/diff_diff/local_linear.py index 0d94dbed5..b53bec908 100644 --- a/diff_diff/local_linear.py +++ b/diff_diff/local_linear.py @@ -1088,17 +1088,14 @@ def bias_corrected_local_linear( # boundary with the same messages the port raises downstream. if weights.shape[0] != d_np.shape[0]: raise ValueError( - f"weights length ({weights.shape[0]}) does not match " - f"d/y ({d_np.shape[0]})." + f"weights length ({weights.shape[0]}) does not match " f"d/y ({d_np.shape[0]})." ) if not np.all(np.isfinite(weights)): raise ValueError("weights contains non-finite values (NaN or Inf).") if np.any(weights < 0): raise ValueError("weights must be non-negative.") if np.sum(weights) <= 0: - raise ValueError( - "weights sum to zero — no observations have positive weight." - ) + raise ValueError("weights sum to zero — no observations have positive weight.") if weights.shape[0] == d_np.shape[0]: _positive_mask_full = weights > 0.0 if not bool(_positive_mask_full.all()): @@ -1303,11 +1300,7 @@ def bias_corrected_local_linear( # / df_survey structure of the full design — the correct # subpopulation / domain-analysis convention. if_out = result.influence_function - if ( - if_out is not None - and _positive_mask_full is not None - and _n_full_for_if is not None - ): + if if_out is not None and _positive_mask_full is not None and _n_full_for_if is not None: if_full = np.zeros(_n_full_for_if, dtype=np.float64) if_full[_positive_mask_full] = if_out if_out = if_full diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index 7829c6b37..d070e9204 100644 --- a/diff_diff/triple_diff.py +++ b/diff_diff/triple_diff.py @@ -1380,9 +1380,7 @@ def _estimate_ddd_decomposition( # suppressed under rank_deficient_action="silent". if self._lstsq_fallback_tracker and self.rank_deficient_action != "silent": n_cells = len(self._lstsq_fallback_tracker) - finite_conds = [ - c for c in self._lstsq_fallback_tracker if np.isfinite(c) - ] + finite_conds = [c for c in self._lstsq_fallback_tracker if np.isfinite(c)] max_cond = max(finite_conds) if finite_conds else float("inf") warnings.warn( f"Rank-deficient covariate design encountered {n_cells} time(s) " diff --git a/diff_diff/trop.py b/diff_diff/trop.py index 467e88786..968b8adf9 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -38,8 +38,8 @@ ) from diff_diff.trop_results import ( _LAMBDA_INF, - _PrecomputedStructures, TROPResults, + _PrecomputedStructures, ) from diff_diff.utils import safe_inference, warn_if_not_converged diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index 263401dbb..2890d3774 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -1405,7 +1405,6 @@ def _bootstrap_rao_wu_local( Tuple[float, np.ndarray] (se, bootstrap_estimates). """ - import warnings from diff_diff.bootstrap_utils import generate_rao_wu_weights from diff_diff.linalg import _factorize_cluster_ids diff --git a/pyproject.toml b/pyproject.toml index a546eb103..b1c63f557 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,9 +60,13 @@ dev = [ "pytest>=7.0", "pytest-xdist>=3.0", "pytest-cov>=4.0", - "black>=23.0", - "ruff>=0.1.0", - "mypy>=1.0", + # Lint/format/type tools are pinned exactly so local results match the + # Lint CI workflow (.github/workflows/lint.yml) — update both together. + # NOTE: black/mypy require Python >= 3.10 (dev tooling only; the library + # floor stays 3.9). + "black==26.3.1", + "ruff==0.15.13", + "mypy==2.1.0", "maturin>=1.4,<2.0", "matplotlib>=3.5", "nbmake>=1.5", @@ -129,6 +133,22 @@ target-version = "py39" select = ["E", "F", "I", "W"] ignore = ["E501"] +[tool.ruff.lint.per-file-ignores] +# Package __init__ deliberately re-exports backend symbols (implicit re-export +# surface; mirrors mypy implicit_reexport = true). +"diff_diff/__init__.py" = ["F401"] +# The trop modules configure the module logger before intra-package imports on +# purpose; import order is load-bearing. +"diff_diff/trop.py" = ["E402"] +"diff_diff/trop_global.py" = ["E402"] +"diff_diff/trop_local.py" = ["E402"] +# Rambachan-Roth / FLCI math notation uses `l` for the linear-combination vector. +"diff_diff/honest_did.py" = ["E741"] +# conftest sets sys.path and MPLBACKEND before imports (order-sensitive). +"tests/conftest.py" = ["E402"] +# CdH event-study horizon notation `l` is quoted in the tutorial narrative. +"tests/test_t19_marketing_pulse_drift.py" = ["E741"] + [tool.mypy] python_version = "3.9" warn_return_any = false From 6d77942b907965d21142847572a05b6494a3ee8c Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Jul 2026 11:08:44 -0400 Subject: [PATCH 2/2] test(lint): normalize + hand-fix the 16 files with non-autofixable findings Each of these files carries its mechanical black/ruff-fix changes PLUS the hand fixes: - F821: test_methodology_sdid.py missing typing imports for a local annotation (latent NameError trap, not runtime) - E402: test_rust_backend.py mid-file _rust_demean_map import hoisted - E741: test_doc_snippets.py comprehension variable l -> ln - F841 x26 audited individually: * test_chaisemartin_dhaultfoeuille trends_nonparam_basic: r_plain was computed but never asserted despite the comment promising both results finite -> added the missing finiteness assertions * test_se_accuracy influence_function_normalization: deleted the abandoned R-style comparison block and rewrote the docstring to state what the test actually asserts * test_imputation always_treated_warning: removed the dead first two data-construction attempts (kept the explanatory note) * remaining sites: dropped dead locals, keeping side-effectful calls Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- tests/test_chaisemartin_dhaultfoeuille.py | 2610 +++++++++++++-------- tests/test_doc_snippets.py | 5 +- tests/test_estimators.py | 4 +- tests/test_imputation.py | 41 +- tests/test_linalg.py | 17 +- tests/test_methodology_sdid.py | 1565 +++++++----- tests/test_methodology_triple_diff.py | 111 +- tests/test_openai_review.py | 792 ++++--- tests/test_rust_backend.py | 61 +- tests/test_se_accuracy.py | 194 +- tests/test_staggered.py | 218 +- tests/test_sun_abraham.py | 6 +- tests/test_survey.py | 4 - tests/test_survey_phase6.py | 741 ++++-- tests/test_triple_diff.py | 2 +- tests/test_wooldridge.py | 2 - 16 files changed, 3950 insertions(+), 2423 deletions(-) diff --git a/tests/test_chaisemartin_dhaultfoeuille.py b/tests/test_chaisemartin_dhaultfoeuille.py index f6a314ad4..4f42e6c3e 100644 --- a/tests/test_chaisemartin_dhaultfoeuille.py +++ b/tests/test_chaisemartin_dhaultfoeuille.py @@ -1891,8 +1891,12 @@ def test_L_max_1_bootstrap_overall_matches_es1(self, data, ci_params): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - data, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + data, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) es1 = r.event_study_effects[1] assert r.overall_att == es1["effect"] @@ -1910,8 +1914,12 @@ def test_L_max_1_suppresses_joiner_leaver_decomposition(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - data, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + data, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) # Joiners/leavers suppressed for L_max=1 assert r.joiners_available is False @@ -1936,8 +1944,12 @@ def test_L_max_1_bootstrap_results_overall_synced(self, data, ci_params): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - data, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + data, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) br = r.bootstrap_results assert br is not None @@ -2147,9 +2159,7 @@ def test_placebo_se_finite_multi_horizon(self, data): def test_placebo_bootstrap_se_multi_horizon(self, data, ci_params): """Multi-horizon placebo bootstrap SE should be finite.""" n_boot = ci_params.bootstrap(199) - est = ChaisemartinDHaultfoeuille( - twfe_diagnostic=False, n_bootstrap=n_boot, seed=42 - ) + est = ChaisemartinDHaultfoeuille(twfe_diagnostic=False, n_bootstrap=n_boot, seed=42) with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( @@ -2365,9 +2375,7 @@ def _make_panel_with_covariates(seed=42, n_groups=40, n_periods=6): # Outcome depends on group FE, time trend, covariate, # and treatment effect y = group_fe + 2.0 * t + 3.0 * x + 5.0 * d + rng.normal(0, 0.5) - rows.append( - {"group": g, "period": t, "treatment": d, "outcome": y, "X1": x} - ) + rows.append({"group": g, "period": t, "treatment": d, "outcome": y, "X1": x}) return pd.DataFrame(rows) def test_controls_requires_lmax(self): @@ -2383,8 +2391,13 @@ def test_controls_missing_column(self): df = self._make_panel_with_covariates() with pytest.raises(ValueError, match="not found in data"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["nonexistent"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["nonexistent"], + L_max=1, ) def test_covariate_residualization_basic(self): @@ -2396,8 +2409,13 @@ def test_covariate_residualization_basic(self): r_plain = est.fit(df, "outcome", "group", "period", "treatment", L_max=1) # Covariate-adjusted r_x = est.fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=1, ) # Results should differ (covariate is confounding) @@ -2416,8 +2434,13 @@ def test_multiple_covariates(self): df["X2"] = np.random.RandomState(99).normal(0, 1, len(df)) est = ChaisemartinDHaultfoeuille(seed=1) r = est.fit( - df, "outcome", "group", "period", "treatment", - controls=["X1", "X2"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1", "X2"], + L_max=1, ) assert r.covariate_residuals is not None # Should have rows for each (baseline, covariate) combination @@ -2427,8 +2450,13 @@ def test_covariate_residuals_diagnostics(self): """Diagnostics DataFrame has expected structure.""" df = self._make_panel_with_covariates() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=2, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=2, ) diag = r.covariate_residuals assert diag is not None @@ -2459,8 +2487,13 @@ def test_controls_with_nonbinary_treatment(self): rows.append({"group": g, "period": t, "treatment": d, "outcome": y, "X1": x}) df = pd.DataFrame(rows) r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=1, ) assert np.isfinite(r.overall_att) assert np.isfinite(r.overall_se) @@ -2469,8 +2502,13 @@ def test_controls_with_multi_horizon(self): """Covariates work with L_max > 1 event study.""" df = self._make_panel_with_covariates() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=2, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=2, ) assert r.event_study_effects is not None assert 1 in r.event_study_effects @@ -2489,8 +2527,13 @@ def test_controls_lmax1_estimand_contract(self): r_raw = est.fit(df, "outcome", "group", "period", "treatment") # Fit with controls r_x = est.fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=1, ) # per_period_effects should be UNADJUSTED (raw Phase 1 DID_M) @@ -2504,9 +2547,7 @@ def test_controls_lmax1_estimand_contract(self): ), f"per_period_effects should be unadjusted at period {period_key}" # overall_att should come from event_study_effects[1] (DID^X_1) - assert r_x.overall_att == pytest.approx( - r_x.event_study_effects[1]["effect"], abs=1e-10 - ) + assert r_x.overall_att == pytest.approx(r_x.event_study_effects[1]["effect"], abs=1e-10) # and should differ from the raw overall_att (covariate effect) assert r_x.overall_att != r_raw.overall_att @@ -2541,7 +2582,11 @@ def test_trends_linear_requires_lmax(self): df = self._make_panel_with_trends() with pytest.raises(ValueError, match="requires L_max >= 1"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", trends_linear=True, ) @@ -2551,8 +2596,13 @@ def test_trends_linear_basic(self): est = ChaisemartinDHaultfoeuille(seed=1) r_plain = est.fit(df, "outcome", "group", "period", "treatment", L_max=2) r_fd = est.fit( - df, "outcome", "group", "period", "treatment", - L_max=2, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + trends_linear=True, ) # Results should differ (group-specific trends confound unadjusted) assert r_fd.overall_att != r_plain.overall_att @@ -2564,8 +2614,13 @@ def test_cumulated_level_effects(self): """Cumulated delta^{fd}_l = sum DID^{fd}_{l'} for l'=1..l.""" df = self._make_panel_with_trends() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=3, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=3, + trends_linear=True, ) assert r.linear_trends_effects is not None # Check cumulation: delta^{fd}_1 = DID^{fd}_1 @@ -2593,8 +2648,13 @@ def test_fg_less_than_3_warning(self): df = pd.DataFrame(rows) with pytest.warns(UserWarning, match="F_g < 3"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + trends_linear=True, ) def test_trends_with_covariates(self): @@ -2602,8 +2662,14 @@ def test_trends_with_covariates(self): df = self._make_panel_with_trends() df["X1"] = np.random.RandomState(77).normal(0, 1, len(df)) r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=2, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=2, + trends_linear=True, ) # overall_att is NaN for trends + L_max>=2 (no aggregate) assert np.isnan(r.overall_att) @@ -2619,8 +2685,13 @@ def test_trends_linear_lmax2_overall_surface(self): """ df = self._make_panel_with_trends() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=3, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=3, + trends_linear=True, ) # overall_* should be NaN (not computed in trends mode) assert np.isnan(r.overall_att) @@ -2654,8 +2725,13 @@ def test_cumulated_se_nan_propagation(self): rows.append({"group": g, "period": t, "treatment": d, "outcome": y}) df = pd.DataFrame(rows) r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + trends_linear=True, ) # If SE at horizon 1 is finite but horizon 2 is NaN, # cumulated h=2 SE must be NaN (not 0.0) @@ -2664,8 +2740,7 @@ def test_cumulated_se_nan_propagation(self): es = r.event_study_effects if es and 2 in es and not np.isfinite(es[2]["se"]): assert not np.isfinite(cum_se), ( - f"Cumulated SE should be NaN when component h=2 SE is NaN, " - f"got {cum_se}" + f"Cumulated SE should be NaN when component h=2 SE is NaN, " f"got {cum_se}" ) @@ -2684,17 +2759,26 @@ def _make_panel_with_sets(seed=42, n_groups=40, n_periods=6): for t in range(n_periods): d = 1 if (switches and t >= 3) else 0 y = group_fe + 2.0 * t + 5.0 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": state, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": state, + } + ) return pd.DataFrame(rows) def test_trends_nonparam_requires_lmax(self): df = self._make_panel_with_sets() with pytest.raises(ValueError, match="requires L_max >= 1"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", trends_nonparam="state", ) @@ -2704,11 +2788,18 @@ def test_trends_nonparam_basic(self): est = ChaisemartinDHaultfoeuille(seed=1) r_plain = est.fit(df, "outcome", "group", "period", "treatment", L_max=1) r_set = est.fit( - df, "outcome", "group", "period", "treatment", - L_max=1, trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + trends_nonparam="state", ) # With set-restricted controls, results may differ # (both should be finite and reasonable) + assert np.isfinite(r_plain.overall_att) + assert np.isfinite(r_plain.overall_se) assert np.isfinite(r_set.overall_att) assert np.isfinite(r_set.overall_se) @@ -2719,16 +2810,26 @@ def test_time_varying_set_raises(self): df.loc[(df["group"] == 0) & (df["period"] == 3), "state"] = 99 with pytest.raises(ValueError, match="time-invariant"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + trends_nonparam="state", ) def test_missing_set_column_raises(self): df = self._make_panel_with_sets() with pytest.raises(ValueError, match="not found in data"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, trends_nonparam="nonexistent", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + trends_nonparam="nonexistent", ) def test_group_level_set_rejected(self): @@ -2737,8 +2838,13 @@ def test_group_level_set_rejected(self): # Use group column itself as set (each group is its own set) with pytest.raises(ValueError, match="coarser than group"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, trends_nonparam="group", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + trends_nonparam="group", ) def test_nan_set_membership_rejected(self): @@ -2747,8 +2853,13 @@ def test_nan_set_membership_rejected(self): df.loc[df["group"] == 0, "state"] = np.nan with pytest.raises(ValueError, match="NaN/missing"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + trends_nonparam="state", ) def test_nonparam_with_covariates(self): @@ -2756,8 +2867,14 @@ def test_nonparam_with_covariates(self): df = self._make_panel_with_sets() df["X1"] = np.random.RandomState(77).normal(0, 1, len(df)) r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=1, trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=1, + trends_nonparam="state", ) assert np.isfinite(r.overall_att) assert r.covariate_residuals is not None @@ -2778,25 +2895,40 @@ def test_trends_nonparam_unequal_support(self): for t in range(n_periods): d = 1 if (switches and t >= 3) else 0 y = 10 + 2.0 * t + 5.0 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": "A", - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": "A", + } + ) # State B: groups 8-9 (both switch at t=3, NO controls in this set) for g in range(8, 10): for t in range(n_periods): d = 1 if t >= 3 else 0 y = 10 + 2.0 * t + 5.0 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": "B", - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": "B", + } + ) df = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + trends_nonparam="state", ) # Should not error; State A groups contribute, State B excluded assert np.isfinite(r.overall_att) @@ -2819,18 +2951,28 @@ def _make_panel_with_het(seed=42, n_groups=40, n_periods=6): for t in range(n_periods): d = 1 if (switches and t >= 3) else 0 y = group_fe + 2.0 * t + effect * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": x_g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": x_g, + } + ) return pd.DataFrame(rows) def test_heterogeneity_basic(self): """Detect heterogeneous effects with binary covariate.""" df = self._make_panel_with_het() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="het_x", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="het_x", ) assert r.heterogeneity_effects is not None assert 1 in r.heterogeneity_effects @@ -2850,14 +2992,24 @@ def test_heterogeneity_null(self): for t in range(6): d = 1 if (switches and t >= 3) else 0 y = 10 + 2 * t + 5 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": x_g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": x_g, + } + ) df = pd.DataFrame(rows) r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="het_x", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="het_x", ) het = r.heterogeneity_effects[1] # Not significantly different from zero @@ -2867,8 +3019,13 @@ def test_heterogeneity_multi_horizon(self): """Heterogeneity test at multiple horizons.""" df = self._make_panel_with_het() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, heterogeneity="het_x", + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + heterogeneity="het_x", ) assert 1 in r.heterogeneity_effects assert 2 in r.heterogeneity_effects @@ -2891,8 +3048,13 @@ def test_heterogeneity_inference_local_invariants(self): """ df = self._make_panel_with_het() r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, heterogeneity="het_x", + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + heterogeneity="het_x", ) assert r.heterogeneity_effects is not None checked = 0 @@ -2901,20 +3063,16 @@ def test_heterogeneity_inference_local_invariants(self): continue expected_t = het["beta"] / het["se"] assert het["t_stat"] == pytest.approx(expected_t, rel=1e-12), ( - f"l={l_h} t_stat: stored={het['t_stat']} vs " - f"beta/se={expected_t}" + f"l={l_h} t_stat: stored={het['t_stat']} vs " f"beta/se={expected_t}" ) half_low = het["beta"] - het["conf_int"][0] half_high = het["conf_int"][1] - het["beta"] assert half_low > 0, f"l={l_h} conf_int_lower not below beta" assert half_high > 0, f"l={l_h} conf_int_upper not above beta" assert half_low == pytest.approx(half_high, rel=1e-12), ( - f"l={l_h} conf_int asymmetric: " - f"below={half_low} above={half_high}" - ) - assert 0.0 <= het["p_value"] <= 1.0, ( - f"l={l_h} p_value out of [0, 1]: {het['p_value']}" + f"l={l_h} conf_int asymmetric: " f"below={half_low} above={half_high}" ) + assert 0.0 <= het["p_value"] <= 1.0, f"l={l_h} p_value out of [0, 1]: {het['p_value']}" checked += 1 assert checked >= 1, "Expected at least one populated heterogeneity horizon" @@ -2922,8 +3080,13 @@ def test_heterogeneity_missing_column(self): df = self._make_panel_with_het() with pytest.raises(ValueError, match="not found"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="nonexistent", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="nonexistent", ) def test_heterogeneity_rejects_controls(self): @@ -2932,8 +3095,14 @@ def test_heterogeneity_rejects_controls(self): df["X1"] = np.random.RandomState(42).normal(0, 1, len(df)) with pytest.raises(ValueError, match="cannot be combined with controls"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="het_x", controls=["X1"], + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="het_x", + controls=["X1"], ) def test_heterogeneity_requires_lmax(self): @@ -2941,7 +3110,11 @@ def test_heterogeneity_requires_lmax(self): df = self._make_panel_with_het() with pytest.raises(ValueError, match="requires L_max >= 1"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", heterogeneity="het_x", ) @@ -2950,8 +3123,14 @@ def test_heterogeneity_rejects_trends_linear(self): df = self._make_panel_with_het() with pytest.raises(ValueError, match="cannot be combined with trends_linear"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, heterogeneity="het_x", trends_linear=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + heterogeneity="het_x", + trends_linear=True, ) def test_heterogeneity_rejects_trends_nonparam(self): @@ -2960,8 +3139,14 @@ def test_heterogeneity_rejects_trends_nonparam(self): df["state"] = df["group"] % 3 with pytest.raises(ValueError, match="cannot be combined with trends_nonparam"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="het_x", trends_nonparam="state", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="het_x", + trends_nonparam="state", ) @@ -2994,8 +3179,13 @@ def test_design2_basic(self): df = self._make_join_then_leave_panel() # drop_larger_lower=False to keep the 2-switch groups r = ChaisemartinDHaultfoeuille(seed=1, drop_larger_lower=False).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, design2=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + design2=True, ) assert r.design2_effects is not None assert r.design2_effects["n_design2_groups"] == 10 @@ -3016,8 +3206,13 @@ def test_design2_no_eligible(self): df = pd.DataFrame(rows) # drop_larger_lower=False required for design2=True r = ChaisemartinDHaultfoeuille(seed=1, drop_larger_lower=False).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, design2=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + design2=True, ) assert r.design2_effects is None @@ -3025,7 +3220,12 @@ def test_design2_disabled_by_default(self): """design2=False (default) produces no design2_effects.""" df = self._make_join_then_leave_panel() r = ChaisemartinDHaultfoeuille(seed=1, drop_larger_lower=False).fit( - df, "outcome", "group", "period", "treatment", L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, ) assert r.design2_effects is None @@ -3034,8 +3234,13 @@ def test_design2_rejects_drop_larger_lower(self): df = self._make_join_then_leave_panel() with pytest.raises(ValueError, match="drop_larger_lower=False"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, design2=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + design2=True, ) @@ -3060,8 +3265,12 @@ def test_ordinal_treatment(self): warnings.simplefilter("ignore") # Non-binary treatment requires L_max (multi-horizon path) r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=2, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=2, ) assert np.isfinite(r.overall_att) @@ -3099,8 +3308,12 @@ def test_single_large_dose_not_flagged_multi_switch(self): warnings.simplefilter("ignore") # Non-binary treatment requires L_max >= 1 (multi-horizon path) r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) # All 20 switcher groups should be kept (0 dropped as multi-switch) assert r.n_groups_dropped_crossers == 0 @@ -3183,8 +3396,12 @@ def test_mixed_binary_nonbinary_panel_lmax1(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) # overall_att should be from per-group path (includes both 0->1 and 0->2) assert np.isfinite(r.overall_att) @@ -3206,8 +3423,12 @@ def test_constant_nonbinary_treatment_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) def test_nonbinary_bootstrap(self, ci_params): @@ -3226,14 +3447,16 @@ def test_nonbinary_bootstrap(self, ci_params): y = 10 + t + np.random.randn() * 0.3 rows.append({"group": g, "period": t, "treatment": 0, "outcome": y}) df = pd.DataFrame(rows) - est = ChaisemartinDHaultfoeuille( - twfe_diagnostic=False, n_bootstrap=n_boot, seed=42 - ) + est = ChaisemartinDHaultfoeuille(twfe_diagnostic=False, n_bootstrap=n_boot, seed=42) with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) assert r.bootstrap_results is not None assert r.bootstrap_results.event_study_ses is not None @@ -3264,8 +3487,12 @@ def test_nonbinary_lmax1_renderer_contract(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) # __repr__ should say DID_1 assert "DID_1" in repr(r) @@ -3304,8 +3531,12 @@ def test_twfe_diagnostic_skipped_nonbinary(self): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=1, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=1, ) twfe_warnings = [x for x in w if "TWFE diagnostic" in str(x.message)] assert len(twfe_warnings) >= 1 @@ -3331,8 +3562,12 @@ def test_normalized_effects_general_formula(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) if r.normalized_effects is not None and 1 in r.normalized_effects: # For dose 0->2: denominator at l=1 should be ~2 (not 1) @@ -3350,9 +3585,7 @@ class TestHonestDiDIntegration: @staticmethod def _make_data(n_groups=40, n_periods=6, seed=42): - return generate_reversible_did_data( - n_groups=n_groups, n_periods=n_periods, seed=seed - ) + return generate_reversible_did_data(n_groups=n_groups, n_periods=n_periods, seed=seed) def test_honest_did_basic(self): """honest_did=True with L_max>=2 produces HonestDiDResults.""" @@ -3362,8 +3595,13 @@ def test_honest_did_basic(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) assert r.honest_did_results is not None assert isinstance(r.honest_did_results, HonestDiDResults) @@ -3375,7 +3613,11 @@ def test_honest_did_requires_lmax(self): df = self._make_data() with pytest.raises(ValueError, match="honest_did=True requires L_max"): ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", honest_did=True, ) @@ -3384,8 +3626,13 @@ def test_honest_did_rejects_placebo_false(self): df = self._make_data() with pytest.raises(ValueError, match="placebo=False"): ChaisemartinDHaultfoeuille(seed=1, placebo=False).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) def test_honest_did_standalone(self): @@ -3396,23 +3643,26 @@ def test_honest_did_standalone(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r_auto = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) r_plain = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", L_max=2, ) - r_manual = compute_honest_did( - r_plain, method="relative_magnitude", M=1.0 - ) + r_manual = compute_honest_did(r_plain, method="relative_magnitude", M=1.0) # Deterministic - bitwise identical - np.testing.assert_allclose( - r_auto.honest_did_results.ci_lb, r_manual.ci_lb, rtol=0 - ) - np.testing.assert_allclose( - r_auto.honest_did_results.ci_ub, r_manual.ci_ub, rtol=0 - ) + np.testing.assert_allclose(r_auto.honest_did_results.ci_lb, r_manual.ci_lb, rtol=0) + np.testing.assert_allclose(r_auto.honest_did_results.ci_ub, r_manual.ci_ub, rtol=0) def test_honest_did_with_controls(self): """HonestDiD runs on DID^X placebos.""" @@ -3421,8 +3671,14 @@ def test_honest_did_with_controls(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=2, + honest_did=True, ) assert r.honest_did_results is not None assert np.isfinite(r.honest_did_results.ci_lb) @@ -3433,8 +3689,14 @@ def test_honest_did_with_trends_linear(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - trends_linear=True, L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + trends_linear=True, + L_max=2, + honest_did=True, ) # Bounds should be computed on second-differenced estimand assert r.honest_did_results is not None @@ -3448,13 +3710,15 @@ def test_honest_did_sensitivity(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", L_max=2, ) honest = HonestDiD(method="relative_magnitude") - sens = honest.sensitivity_analysis( - r, M_grid=list(np.linspace(0, 2, 5)) - ) + sens = honest.sensitivity_analysis(r, M_grid=list(np.linspace(0, 2, 5))) assert sens.breakdown_M is not None or len(sens.bounds) == 5 def test_honest_did_smoothness(self): @@ -3465,7 +3729,11 @@ def test_honest_did_smoothness(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", L_max=2, ) rm_bounds = compute_honest_did(r, method="relative_magnitude", M=1.0) @@ -3479,8 +3747,13 @@ def test_honest_did_original_estimate_is_post_average(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) hd = r.honest_did_results assert hd is not None @@ -3497,7 +3770,11 @@ def test_honest_did_custom_l_vec_on_impact(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", L_max=2, ) # l_vec=[1, 0] targets only DID_1 (on-impact, R's default) @@ -3514,8 +3791,13 @@ def test_honest_did_respects_alpha(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1, alpha=0.10).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) assert r.honest_did_results is not None assert r.honest_did_results.alpha == 0.10 @@ -3526,8 +3808,13 @@ def test_honest_did_retains_period_metadata(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) hd = r.honest_did_results assert hd.pre_periods_used is not None @@ -3546,13 +3833,15 @@ def test_honest_did_custom_l_vec_summary_label(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", + df, + "outcome", + "group", + "period", + "treatment", L_max=2, ) # Attach custom-target HonestDiD to results - r.honest_did_results = compute_honest_did( - r, l_vec=np.array([1.0, 0.0]) - ) + r.honest_did_results = compute_honest_did(r, l_vec=np.array([1.0, 0.0])) text = r.summary() assert "on-impact" in text.lower() assert "Equal-weight" not in text @@ -3567,16 +3856,27 @@ def test_honest_did_with_trends_nonparam(self): for t in range(7): d = 1 if (switches and t >= 3) else 0 y = 10 + 2.0 * t + 5.0 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": state, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": state, + } + ) df = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, trends_nonparam="state", honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + trends_nonparam="state", + honest_did=True, ) assert r.honest_did_results is not None assert np.isfinite(r.honest_did_results.ci_lb) @@ -3599,28 +3899,44 @@ def test_honest_did_trends_nonparam_trimming(self): switches = g < 3 for t in range(n_periods): d = 1 if (switches and t >= 5) else 0 - y = 10 + 2.0*t + 5.0*d + rng.normal(0, 0.3) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": "A", - }) + y = 10 + 2.0 * t + 5.0 * d + rng.normal(0, 0.3) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": "A", + } + ) # State B: 4 switch at t=2, 2 "controls" switch at t=3 for g in range(7, 13): switch_t = 2 if g < 11 else 3 for t in range(n_periods): d = 1 if t >= switch_t else 0 - y = 10 + 2.0*t + 5.0*d + rng.normal(0, 0.3) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "state": "B", - }) + y = 10 + 2.0 * t + 5.0 * d + rng.normal(0, 0.3) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "state": "B", + } + ) df = pd.DataFrame(rows) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=3, trends_nonparam="state", honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=3, + trends_nonparam="state", + honest_did=True, ) # h=3 and h=-3 should be NaN (N_l=0 from support trimming) assert r.event_study_effects[3]["n_obs"] == 0 @@ -3635,8 +3951,9 @@ def test_honest_did_trends_nonparam_trimming(self): assert hd.post_periods_used == [1, 2] # The placebo-based pre-period warning should have been emitted placebo_warns = [ - x for x in w if "placebo" in str(x.message).lower() - and "pre-period" in str(x.message).lower() + x + for x in w + if "placebo" in str(x.message).lower() and "pre-period" in str(x.message).lower() ] assert len(placebo_warns) >= 1 @@ -3646,8 +3963,13 @@ def test_honest_did_with_bootstrap(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1, n_bootstrap=49).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) assert r.honest_did_results is not None assert np.isfinite(r.honest_did_results.ci_lb) @@ -3664,9 +3986,7 @@ class TestSummaryPhase3: @staticmethod def _make_data(n_groups=40, n_periods=6, seed=42): - return generate_reversible_did_data( - n_groups=n_groups, n_periods=n_periods, seed=seed - ) + return generate_reversible_did_data(n_groups=n_groups, n_periods=n_periods, seed=seed) def test_summary_renders_covariate_diagnostics(self): """Covariate Adjustment section appears in summary().""" @@ -3675,8 +3995,13 @@ def test_summary_renders_covariate_diagnostics(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - controls=["X1"], L_max=1, + df, + "outcome", + "group", + "period", + "treatment", + controls=["X1"], + L_max=1, ) text = r.summary() assert "Covariate Adjustment" in text @@ -3687,8 +4012,13 @@ def test_summary_renders_linear_trends(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - trends_linear=True, L_max=2, + df, + "outcome", + "group", + "period", + "treatment", + trends_linear=True, + L_max=2, ) text = r.summary() assert "Cumulated Level Effects" in text @@ -3703,16 +4033,26 @@ def test_summary_renders_heterogeneity(self): for t in range(6): d = 1 if (switches and t >= 3) else 0 y = 10 + 2.0 * t + 5.0 * d + 3.0 * x_g * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": x_g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": x_g, + } + ) df = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, heterogeneity="het_x", + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + heterogeneity="het_x", ) text = r.summary() assert "Heterogeneity Test" in text @@ -3730,17 +4070,25 @@ def test_summary_renders_design2(self): else: d = 0 # never switch y = 10 + t + 5.0 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, "outcome": y, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + } + ) df = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore") - r = ChaisemartinDHaultfoeuille( - seed=1, drop_larger_lower=False - ).fit( - df, "outcome", "group", "period", "treatment", - L_max=1, design2=True, + r = ChaisemartinDHaultfoeuille(seed=1, drop_larger_lower=False).fit( + df, + "outcome", + "group", + "period", + "treatment", + L_max=1, + design2=True, ) text = r.summary() assert "Design-2" in text @@ -3751,8 +4099,13 @@ def test_summary_renders_honest_did(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") r = ChaisemartinDHaultfoeuille(seed=1).fit( - df, "outcome", "group", "period", "treatment", - L_max=2, honest_did=True, + df, + "outcome", + "group", + "period", + "treatment", + L_max=2, + honest_did=True, ) text = r.summary() assert "HonestDiD Sensitivity" in text @@ -4153,9 +4506,7 @@ def test_empty_path_surface_when_no_complete_window(self): assert results.path_effects == {} # Fit-time warning surfaced - empty_warnings = [ - w for w in caught if "no observed treatment path" in str(w.message) - ] + empty_warnings = [w for w in caught if "no observed treatment path" in str(w.message)] assert empty_warnings, ( "Expected a UserWarning when by_path is requested but no " "observed path has a complete window" @@ -4208,10 +4559,18 @@ def test_degenerate_cohort_path_nan_inference_and_warning(self): "period": [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], "treatment": [0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1], "outcome": [ - 10.0, 13.0, 14.0, - 10.0, 11.0, 9.0, - 10.0, 11.0, 12.0, - 10.0, 11.0, 12.0, + 10.0, + 13.0, + 14.0, + 10.0, + 11.0, + 9.0, + 10.0, + 11.0, + 12.0, + 10.0, + 11.0, + 12.0, ], } ) @@ -4239,8 +4598,7 @@ def test_degenerate_cohort_path_nan_inference_and_warning(self): degenerate_warnings = [ w for w in caught - if "unidentified for path=" in str(w.message) - and "horizon l=" in str(w.message) + if "unidentified for path=" in str(w.message) and "horizon l=" in str(w.message) ] assert degenerate_warnings, ( "Expected a per-(path, horizon) degenerate-cohort UserWarning " @@ -4314,9 +4672,7 @@ def test_point_estimates_preserved(self): bit-identical to the analytical fit.""" data = _by_path_three_path_data() _est_a, res_a = _fit_by_path(data, by_path=3, L_max=3) - _est_b, res_b = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=100, seed=42 - ) + _est_b, res_b = self._fit_with_bootstrap(data, by_path=3, L_max=3, n_bootstrap=100, seed=42) assert res_a.path_effects is not None and res_b.path_effects is not None assert set(res_a.path_effects.keys()) == set(res_b.path_effects.keys()) for path, entry_a in res_a.path_effects.items(): @@ -4327,7 +4683,10 @@ def test_point_estimates_preserved(self): assert np.isnan(h_b["effect"]) else: np.testing.assert_allclose( - h_b["effect"], h_a["effect"], atol=1e-14, rtol=1e-14, + h_b["effect"], + h_a["effect"], + atol=1e-14, + rtol=1e-14, err_msg=f"path={path} l={l_h}: bootstrap changed effect", ) @@ -4335,9 +4694,7 @@ def test_bootstrap_se_finite_and_positive(self): """On the hand-built 3-path panel, every non-degenerate (path, horizon) produces a positive finite bootstrap SE.""" data = _by_path_three_path_data() - _est, res = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=200, seed=42 - ) + _est, res = self._fit_with_bootstrap(data, by_path=3, L_max=3, n_bootstrap=200, seed=42) assert res.path_effects is not None any_finite = False for path, entry in res.path_effects.items(): @@ -4349,8 +4706,7 @@ def test_bootstrap_se_finite_and_positive(self): ) if np.isfinite(h["se"]): assert h["se"] > 0, ( - f"path={path} l={l_h}: bootstrap SE is not " - f"positive: {h['se']}" + f"path={path} l={l_h}: bootstrap SE is not " f"positive: {h['se']}" ) any_finite = True assert any_finite, "No (path, horizon) produced a finite bootstrap SE" @@ -4365,10 +4721,7 @@ def test_bootstrap_se_close_to_analytical_on_well_conditioned(self): extra panel construction is required. """ golden_path = ( - Path(__file__).parents[1] - / "benchmarks" - / "data" - / "dcdh_dynr_golden_values.json" + Path(__file__).parents[1] / "benchmarks" / "data" / "dcdh_dynr_golden_values.json" ) if not golden_path.exists(): pytest.skip( @@ -4420,16 +4773,22 @@ def test_degenerate_cohort_still_nan(self): "period": [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], "treatment": [0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1], "outcome": [ - 10.0, 13.0, 14.0, - 10.0, 11.0, 9.0, - 10.0, 11.0, 12.0, - 10.0, 11.0, 12.0, + 10.0, + 13.0, + 14.0, + 10.0, + 11.0, + 9.0, + 10.0, + 11.0, + 12.0, + 10.0, + 11.0, + 12.0, ], } ) - _est, res = self._fit_with_bootstrap( - panel, by_path=2, L_max=1, n_bootstrap=100, seed=42 - ) + _est, res = self._fit_with_bootstrap(panel, by_path=2, L_max=1, n_bootstrap=100, seed=42) assert res.path_effects is not None any_nan = False for entry in res.path_effects.values(): @@ -4442,8 +4801,7 @@ def test_degenerate_cohort_still_nan(self): assert np.isnan(lo) and np.isnan(hi) assert np.isfinite(h["effect"]) assert any_nan, ( - "Expected at least one NaN-SE (path, horizon) entry under " - "singleton-cohort panel" + "Expected at least one NaN-SE (path, horizon) entry under " "singleton-cohort panel" ) @pytest.mark.parametrize("weights", ["rademacher", "mammen", "webb"]) @@ -4452,8 +4810,12 @@ def test_bootstrap_weights_variants(self, weights): 3-path hand-built panel.""" data = _by_path_three_path_data() _est, res = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=100, - bootstrap_weights=weights, seed=42, + data, + by_path=3, + L_max=3, + n_bootstrap=100, + bootstrap_weights=weights, + seed=42, ) assert res.path_effects is not None any_finite = False @@ -4461,19 +4823,25 @@ def test_bootstrap_weights_variants(self, weights): for h in entry["horizons"].values(): if np.isfinite(h["se"]) and h["se"] > 0: any_finite = True - assert any_finite, ( - f"bootstrap_weights={weights!r} produced no finite per-path SE" - ) + assert any_finite, f"bootstrap_weights={weights!r} produced no finite per-path SE" def test_bootstrap_seed_reproducibility(self): """Two fits with the same seed must produce bit-identical per-(path, horizon) bootstrap SE.""" data = _by_path_three_path_data() _est1, res1 = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=100, seed=2026, + data, + by_path=3, + L_max=3, + n_bootstrap=100, + seed=2026, ) _est2, res2 = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=100, seed=2026, + data, + by_path=3, + L_max=3, + n_bootstrap=100, + seed=2026, ) assert res1.path_effects is not None and res2.path_effects is not None for path, entry1 in res1.path_effects.items(): @@ -4484,7 +4852,8 @@ def test_bootstrap_seed_reproducibility(self): assert np.isnan(h2["se"]) else: np.testing.assert_array_equal( - h1["se"], h2["se"], + h1["se"], + h2["se"], err_msg=f"path={path} l={l_h}: seed reproducibility broke", ) @@ -4499,7 +4868,11 @@ def test_inference_fields_match_bootstrap_results(self): data = _by_path_three_path_data() _est_a, res_a = _fit_by_path(data, by_path=3, L_max=3) _est_b, res_b = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=200, seed=42, + data, + by_path=3, + L_max=3, + n_bootstrap=200, + seed=42, ) assert res_a.path_effects is not None assert res_b.path_effects is not None @@ -4547,7 +4920,10 @@ def test_inference_fields_match_bootstrap_results(self): if np.isfinite(h_b["se"]) and h_b["se"] > 0: expected_t = h_b["effect"] / h_b["se"] np.testing.assert_allclose( - h_b["t_stat"], expected_t, atol=1e-10, rtol=1e-10, + h_b["t_stat"], + expected_t, + atol=1e-10, + rtol=1e-10, err_msg=( f"path={path} l={l_h}: t_stat should be " f"SE-derived per anti-pattern rule" @@ -4575,7 +4951,11 @@ def test_inference_fields_equal_bootstrap_results_directly(self): """ data = _by_path_three_path_data() _est, res = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=200, seed=42, + data, + by_path=3, + L_max=3, + n_bootstrap=200, + seed=42, ) assert res.path_effects is not None br = res.bootstrap_results @@ -4594,14 +4974,16 @@ def test_inference_fields_equal_bootstrap_results_directly(self): continue if np.isfinite(se_br): np.testing.assert_array_equal( - h["se"], se_br, + h["se"], + se_br, err_msg=( f"path={path} l={l_h}: path_effects se " f"{h['se']} != bootstrap_results.path_ses {se_br}" ), ) np.testing.assert_array_equal( - h["p_value"], p_br if p_br is not None else np.nan, + h["p_value"], + p_br if p_br is not None else np.nan, err_msg=( f"path={path} l={l_h}: path_effects p_value " f"{h['p_value']} != " @@ -4612,7 +4994,8 @@ def test_inference_fields_equal_bootstrap_results_directly(self): assert ci_br is not None lo_br, hi_br = ci_br np.testing.assert_array_equal( - [lo_e, hi_e], [lo_br, hi_br], + [lo_e, hi_e], + [lo_br, hi_br], err_msg=( f"path={path} l={l_h}: path_effects conf_int " f"{(lo_e, hi_e)} != " @@ -4657,7 +5040,8 @@ def test_overflow_warning_fires_exactly_once_under_bootstrap(self): L_max=3, ) overflow_warnings = [ - w for w in caught + w + for w in caught if "exceeds the number of observed paths" in str(w.message) or "more than the observed number of paths" in str(w.message) or "requested but only" in str(w.message) @@ -4824,8 +5208,7 @@ def test_nan_contract_extends_to_placebo_event_study_horizons(self): df_es = res.to_dataframe(level="event_study") negative_rows = df_es[df_es["horizon"] < 0] if len(negative_rows) > 0: - for col in ("se", "t_stat", "p_value", - "conf_int_lower", "conf_int_upper"): + for col in ("se", "t_stat", "p_value", "conf_int_lower", "conf_int_upper"): assert negative_rows[col].isna().all(), ( f"to_dataframe(level='event_study') negative-horizon " f"column {col!r} must be NaN under n_bootstrap=1; " @@ -4850,15 +5233,17 @@ def test_summary_footer_mixed_validity_surfaces_live_targets(self): """ data = _by_path_three_path_data() _est, res = self._fit_with_bootstrap( - data, by_path=3, L_max=3, n_bootstrap=200, seed=42, + data, + by_path=3, + L_max=3, + n_bootstrap=200, + seed=42, ) # Sanity: healthy fit has finite overall and path SEs. assert np.isfinite(res.overall_se) assert res.path_effects is not None any_finite_path = any( - np.isfinite(h["se"]) - for e in res.path_effects.values() - for h in e["horizons"].values() + np.isfinite(h["se"]) for e in res.path_effects.values() for h in e["horizons"].values() ) assert any_finite_path @@ -4889,10 +5274,7 @@ def test_summary_footer_mixed_validity_surfaces_live_targets(self): # percentile). assert "multiplier-bootstrap percentile inference" not in summary_text # Must mention "per-path bootstrap inference is populated" - assert ( - "per-path" in summary_text - and "bootstrap inference is populated" in summary_text - ), ( + assert "per-path" in summary_text and "bootstrap inference is populated" in summary_text, ( "Footer must surface which targets retain finite bootstrap " "inference when overall/event-study degenerates. Summary " "tail:\n" @@ -4933,8 +5315,7 @@ def test_nan_contract_extends_to_overall_and_event_study_horizons(self): ) assert np.isnan(res.overall_se), ( - f"n_bootstrap=1: overall_se must be NaN (bootstrap " - f"contract), got {res.overall_se}" + f"n_bootstrap=1: overall_se must be NaN (bootstrap " f"contract), got {res.overall_se}" ) assert np.isnan(res.overall_t_stat) assert np.isnan(res.overall_p_value) @@ -4957,8 +5338,7 @@ def test_nan_contract_extends_to_overall_and_event_study_horizons(self): assert res.event_study_effects is not None for l_h, entry in res.event_study_effects.items(): assert np.isnan(entry["se"]), ( - f"n_bootstrap=1: event_study_effects[{l_h}].se must be " - f"NaN, got {entry['se']}" + f"n_bootstrap=1: event_study_effects[{l_h}].se must be " f"NaN, got {entry['se']}" ) assert np.isnan(entry["t_stat"]) assert np.isnan(entry["p_value"]) @@ -5187,8 +5567,7 @@ def test_path_placebo_point_estimate_within_path_mean(self): g_to_F_g[int(g)] = int(treated["period"].iloc[0]) outcome_lookup = { - (int(r["group"]), int(r["period"])): float(r["outcome"]) - for _, r in data.iterrows() + (int(r["group"]), int(r["period"])): float(r["outcome"]) for _, r in data.iterrows() } # Per-group path tuple g_to_path = {} @@ -5221,22 +5600,16 @@ def test_path_placebo_point_estimate_within_path_mean(self): # switchers in this fixture share baseline 0), not # switched by forward, observed at ref+backward+forward ctrl_groups = [ - gc - for gc in g_to_F_g - if gc != g and g_to_F_g[gc] > forward + gc for gc in g_to_F_g if gc != g and g_to_F_g[gc] > forward ] + never_treated if not ctrl_groups: continue - switcher_change = ( - outcome_lookup[(g, backward)] - outcome_lookup[(g, F_g - 1)] - ) + switcher_change = outcome_lookup[(g, backward)] - outcome_lookup[(g, F_g - 1)] ctrl_changes = [ outcome_lookup[(int(gc), backward)] - outcome_lookup[(int(gc), F_g - 1)] for gc in ctrl_groups ] - contributions.append( - switcher_change - sum(ctrl_changes) / len(ctrl_changes) - ) + contributions.append(switcher_change - sum(ctrl_changes) / len(ctrl_changes)) if contributions: expected_mean = sum(contributions) / len(contributions) np.testing.assert_allclose( @@ -5266,9 +5639,7 @@ def test_path_placebo_se_finite_or_nan(self): if np.isfinite(se): assert se > 0, f"path={path} lag={lag_key}: SE={se} not positive" else: - assert np.isnan(se), ( - f"path={path} lag={lag_key}: SE={se} not NaN-finite" - ) + assert np.isnan(se), f"path={path} lag={lag_key}: SE={se} not NaN-finite" def test_switcher_subset_mask_default_preserves_legacy_placebo_if(self): """``_compute_per_group_if_placebo_horizon(switcher_subset_mask=None)`` @@ -5331,9 +5702,7 @@ def test_path_placebo_t_stat_uses_safe_inference(self): for lag_key, entry in lag_dict.items(): if not np.isfinite(entry["se"]): continue - expected_t = safe_inference( - entry["effect"], entry["se"], alpha=0.05, df=None - )[0] + expected_t = safe_inference(entry["effect"], entry["se"], alpha=0.05, df=None)[0] np.testing.assert_allclose( entry["t_stat"], expected_t, @@ -5355,9 +5724,9 @@ def test_path_placebo_to_dataframe_emits_negative_horizons(self): data = _by_path_placebo_data() _est, res = _fit_by_path_with_placebo(data, by_path=3, L_max=3) df = res.to_dataframe(level="by_path") - assert (df["horizon"] < 0).any(), ( - "to_dataframe(level='by_path') did not emit any negative-horizon rows" - ) + assert ( + df["horizon"] < 0 + ).any(), "to_dataframe(level='by_path') did not emit any negative-horizon rows" def test_empty_path_placebo_surface_when_no_complete_window(self): """``path_placebo_event_study`` empty-state contract: ``{}`` (NOT @@ -5452,8 +5821,7 @@ def test_bootstrap_point_estimates_preserved(self): atol=1e-14, rtol=1e-14, err_msg=( - f"path={path} lag={lag_key}: bootstrap " - f"changed point estimate" + f"path={path} lag={lag_key}: bootstrap " f"changed point estimate" ), ) @@ -5481,9 +5849,7 @@ def test_n_bootstrap_1_enforces_full_nan_tuple(self): PR #364 three rounds in a row. """ data = _by_path_placebo_data() - _est, res = _fit_by_path_with_placebo( - data, by_path=3, L_max=3, n_bootstrap=1, seed=42 - ) + _est, res = _fit_by_path_with_placebo(data, by_path=3, L_max=3, n_bootstrap=1, seed=42) assert res.path_placebo_event_study is not None br = res.bootstrap_results assert br is not None @@ -5663,9 +6029,7 @@ def test_path_sup_t_bands_keys_match_path_effects_with_finite_crit(self): # a finite crit; otherwise it should be absent. for path, entry in res.path_effects.items(): n_valid = sum( - 1 - for h in entry["horizons"].values() - if np.isfinite(h["se"]) and h["se"] > 0 + 1 for h in entry["horizons"].values() if np.isfinite(h["se"]) and h["se"] > 0 ) if n_valid >= 2: # Must be present (assuming gate also passes); if it's @@ -5800,17 +6164,12 @@ def test_path_sup_t_skipped_when_path_has_only_one_valid_horizon(self): single_horizon_paths = [ path for path, entry in res.path_effects.items() - if sum( - 1 - for h in entry["horizons"].values() - if np.isfinite(h["se"]) and h["se"] > 0 - ) + if sum(1 for h in entry["horizons"].values() if np.isfinite(h["se"]) and h["se"] > 0) < 2 ] for path in single_horizon_paths: assert path not in res.path_sup_t_bands, ( - f"path={path} has <2 valid horizons; should be absent " - f"from path_sup_t_bands" + f"path={path} has <2 valid horizons; should be absent " f"from path_sup_t_bands" ) # And no horizon should have cband_conf_int populated. for l_h, h in res.path_effects[path]["horizons"].items(): @@ -5929,14 +6288,10 @@ def test_path_sup_t_bands_empty_dict_when_no_complete_window(self): for g in (1, 2, 3, 4): for t in range(4): d = 1 if t >= 3 else 0 - rows.append( - {"group": g, "period": t, "treatment": d, "outcome": rng.normal()} - ) + rows.append({"group": g, "period": t, "treatment": d, "outcome": rng.normal()}) for g in (5, 6): for t in range(4): - rows.append( - {"group": g, "period": t, "treatment": 0, "outcome": rng.normal()} - ) + rows.append({"group": g, "period": t, "treatment": 0, "outcome": rng.normal()}) data = pd.DataFrame(rows) with warnings.catch_warnings(): @@ -5960,8 +6315,7 @@ def test_path_sup_t_bands_empty_dict_when_no_complete_window(self): # Empty-state contract: requested but empty -> {} not None. assert res.path_effects == {}, ( - f"Expected path_effects == {{}} on no-complete-window panel; " - f"got {res.path_effects}" + f"Expected path_effects == {{}} on no-complete-window panel; " f"got {res.path_effects}" ) assert res.path_sup_t_bands == {}, ( f"Expected path_sup_t_bands == {{}} (not None) when " @@ -6001,9 +6355,7 @@ def test_path_sup_t_to_dataframe_emits_cband_columns(self): horizon = int(row["horizon"]) if horizon > 0 and path in res.path_sup_t_bands: # Should match the horizon's cband_conf_int. - expected_cband = res.path_effects[path]["horizons"][horizon].get( - "cband_conf_int" - ) + expected_cband = res.path_effects[path]["horizons"][horizon].get("cband_conf_int") if expected_cband is not None: np.testing.assert_allclose(row["cband_lower"], expected_cband[0]) np.testing.assert_allclose(row["cband_upper"], expected_cband[1]) @@ -6023,14 +6375,10 @@ def test_path_sup_t_to_dataframe_empty_path_fallback_has_cband_columns(self): for g in (1, 2, 3, 4): for t in range(4): d = 1 if t >= 3 else 0 - rows.append( - {"group": g, "period": t, "treatment": d, "outcome": rng.normal()} - ) + rows.append({"group": g, "period": t, "treatment": d, "outcome": rng.normal()}) for g in (5, 6): for t in range(4): - rows.append( - {"group": g, "period": t, "treatment": 0, "outcome": rng.normal()} - ) + rows.append({"group": g, "period": t, "treatment": 0, "outcome": rng.normal()}) data = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) @@ -6076,9 +6424,7 @@ def test_path_sup_t_strict_majority_gate_at_exact_50pct(self, monkeypatch): original_generator = bs_mod._generate_psu_or_group_weights - def fake_generator( - n_bootstrap, n_groups_target, weight_type, rng, group_to_psu_map - ): + def fake_generator(n_bootstrap, n_groups_target, weight_type, rng, group_to_psu_map): # Call the original to get a sane base, then inject NaN into # exactly half of the bootstrap rows. The NaN propagates # through `weights @ u_centered` -> NaN deviations -> NaN @@ -6147,9 +6493,7 @@ def fake_generator( # --------------------------------------------------------------------------- -def _by_path_three_path_data_with_controls( - seed: int = 42, x_effect: float = 3.0 -) -> pd.DataFrame: +def _by_path_three_path_data_with_controls(seed: int = 42, x_effect: float = 3.0) -> pd.DataFrame: """Three-path panel with confounding covariate X1. Extends ``_by_path_three_path_data``: same 8-group / 4-period @@ -6195,12 +6539,7 @@ def _load_by_path_controls_scenario(): file is missing (CI's isolated-install job ships only tests/, not benchmarks/, per ``feedback_golden_file_pytest_skip.md``). """ - golden_path = ( - Path(__file__).parents[1] - / "benchmarks" - / "data" - / "dcdh_dynr_golden_values.json" - ) + golden_path = Path(__file__).parents[1] / "benchmarks" / "data" / "dcdh_dynr_golden_values.json" if not golden_path.exists(): pytest.skip( f"dCDH golden values file not found at {golden_path}; " @@ -6325,9 +6664,7 @@ def test_multi_covariate_works(self): data = _by_path_three_path_data_with_controls() # Add a second covariate rng = np.random.default_rng(99) - data = data.assign( - X2=lambda d: 0.5 * d["X1"] + rng.normal(0, 0.5, size=len(d)) - ) + data = data.assign(X2=lambda d: 0.5 * d["X1"] + rng.normal(0, 0.5, size=len(d))) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3) @@ -6343,9 +6680,9 @@ def test_multi_covariate_works(self): assert res.path_effects is not None for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["effect"]), ( - f"path={path} l={l_h}: effect not finite under multi-covariate" - ) + assert np.isfinite( + vals["effect"] + ), f"path={path} l={l_h}: effect not finite under multi-covariate" # Bootstrap SE inheritance ------------------------------------------ @pytest.mark.slow @@ -6383,9 +6720,7 @@ def test_bootstrap_point_estimates_unchanged(self): data = _load_by_path_controls_scenario() with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - est_a = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, seed=42 - ) + est_a = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, seed=42) res_a = est_a.fit( data, outcome="outcome", @@ -6448,8 +6783,7 @@ def test_per_path_placebos_with_controls_present(self): any_finite = True break assert any_finite, ( - "No per-path placebo lag produced a finite effect under " - "controls + by_path + placebo" + "No per-path placebo lag produced a finite effect under " "controls + by_path + placebo" ) @pytest.mark.slow @@ -6483,9 +6817,7 @@ def test_per_path_placebos_with_controls_bootstrap(self): if np.isfinite(se) and se > 0: any_finite_se = True break - assert any_finite_se, ( - "No per-path placebo lag produced a finite > 0 bootstrap SE" - ) + assert any_finite_se, "No per-path placebo lag produced a finite > 0 bootstrap SE" # Per-path sup-t bands inheritance ---------------------------------- @pytest.mark.slow @@ -6512,8 +6844,7 @@ def test_sup_t_bands_with_controls_finite_crit(self): assert res.path_sup_t_bands is not None # At least one path should pass both gates any_finite = any( - np.isfinite(entry.get("crit_value", np.nan)) - and entry.get("crit_value", -1) > 0 + np.isfinite(entry.get("crit_value", np.nan)) and entry.get("crit_value", -1) > 0 for entry in res.path_sup_t_bands.values() ) assert any_finite, ( @@ -6553,9 +6884,9 @@ def test_per_period_effects_unadjusted_with_by_path_controls(self): if res_no.per_period_effects is not None: assert res_yes.per_period_effects is not None for t in res_no.per_period_effects: - assert t in res_yes.per_period_effects, ( - f"per_period_effects period {t} missing under controls" - ) + assert ( + t in res_yes.per_period_effects + ), f"per_period_effects period {t} missing under controls" for field in ("did_plus_t", "did_minus_t"): np.testing.assert_allclose( res_no.per_period_effects[t][field], @@ -6618,9 +6949,7 @@ def test_to_dataframe_by_path_with_controls_and_bootstrap(self): assert "cband_lower" in df_long.columns assert "cband_upper" in df_long.columns # At least one row must have a finite cband - any_finite_cband = ( - df_long["cband_lower"].notna() & df_long["cband_upper"].notna() - ).any() + any_finite_cband = (df_long["cband_lower"].notna() & df_long["cband_upper"].notna()).any() assert any_finite_cband, ( "to_dataframe(level='by_path') produced no rows with finite " "cband columns under controls + bootstrap" @@ -6646,9 +6975,7 @@ def _add(group, treatment_path): for t, d in enumerate(treatment_path): x = 0.05 * group + 0.15 * t + rng.normal(0, 0.1) y = d * 2.0 + 1.0 * x + rng.normal(0, 0.1) - rows.append( - {"group": group, "period": t, "treatment": d, "outcome": y, "X1": x} - ) + rows.append({"group": group, "period": t, "treatment": d, "outcome": y, "X1": x}) for g in (1, 2, 3): _add(g, [0, 0, 1, 1, 1, 1]) # joiner-late path 0,0,1,1,1,1 @@ -6670,9 +6997,7 @@ def _add(group, treatment_path): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2 - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2) est.fit( data, outcome="outcome", @@ -6689,10 +7014,7 @@ def _add(group, treatment_path): if issubclass(w.category, UserWarning) and "+ controls" in str(w.message) and "multi-baseline" not in str(w.message).lower() - or ( - issubclass(w.category, UserWarning) - and "switcher baselines" in str(w.message) - ) + or (issubclass(w.category, UserWarning) and "switcher baselines" in str(w.message)) ] assert deviation_msgs, ( "Expected a UserWarning mentioning 'by_path + controls' and " @@ -6721,8 +7043,7 @@ def test_single_baseline_panel_does_not_emit_r_deviation_warning(self): deviation_msgs = [ str(w.message) for w in caught - if issubclass(w.category, UserWarning) - and "switcher baselines" in str(w.message) + if issubclass(w.category, UserWarning) and "switcher baselines" in str(w.message) ] assert not deviation_msgs, ( "Multi-baseline deviation warning fired on a single-baseline " @@ -6774,9 +7095,7 @@ def test_single_baseline_heterogeneous_F_g_does_not_warn(self): f"single value, got {sorted(switcher_baselines.unique())}" ) first_treat = ( - data[ - (data["treatment"] == 1) & data["group"].isin(switcher_groups) - ] + data[(data["treatment"] == 1) & data["group"].isin(switcher_groups)] .groupby("group")["period"] .min() ) @@ -6801,8 +7120,7 @@ def test_single_baseline_heterogeneous_F_g_does_not_warn(self): deviation_msgs = [ str(w.message) for w in caught - if issubclass(w.category, UserWarning) - and "switcher baselines" in str(w.message) + if issubclass(w.category, UserWarning) and "switcher baselines" in str(w.message) ] assert not deviation_msgs, ( "Multi-baseline deviation warning fired on a single-baseline " @@ -6849,9 +7167,12 @@ def _by_path_data_with_trends_linear(seed: int = 42) -> pd.DataFrame: (0, 1, 0, 0), # path 3, on briefly ] fg_path_counts = [ - (4, 0, 20), (5, 0, 18), # path 1 = 38 - (6, 1, 13), (7, 1, 11), # path 2 = 24 - (8, 2, 11), (9, 2, 7), # path 3 = 18 + (4, 0, 20), + (5, 0, 18), # path 1 = 38 + (6, 1, 13), + (7, 1, 11), # path 2 = 24 + (8, 2, 11), + (9, 2, 7), # path 3 = 18 ] rows = [] g_id = 0 @@ -6908,17 +7229,22 @@ def _by_path_data_with_trends_nonparam(seed: int = 43) -> pd.DataFrame: ] # F_g 2/3 -> path 1 (40), F_g 4/5 -> path 2 (25), F_g 6 -> path 3 (10) fg_path_counts = [ - (2, 0, 20), (3, 0, 20), - (4, 1, 15), (5, 1, 10), + (2, 0, 20), + (3, 0, 20), + (4, 1, 15), + (5, 1, 10), (6, 2, 10), (7, 2, 5), # rank 4-equivalent absorbed into path 3 cluster (kept by path 3 by frequency) ] # Adjust to ensure top 3 paths have unique counts: # path 1 = 40, path 2 = 25, path 3 = 15 fg_path_counts = [ - (2, 0, 20), (3, 0, 20), - (4, 1, 15), (5, 1, 10), - (6, 2, 8), (7, 2, 7), + (2, 0, 20), + (3, 0, 20), + (4, 1, 15), + (5, 1, 10), + (6, 2, 8), + (7, 2, 7), ] rows = [] g_id = 0 @@ -6958,49 +7284,31 @@ def _by_path_data_with_trends_nonparam(seed: int = 43) -> pd.DataFrame: def _load_by_path_trends_lin_scenario(): """Load golden-value scenario for by_path + trends_linear.""" - golden_path = ( - Path(__file__).parents[1] - / "benchmarks" - / "data" - / "dcdh_dynr_golden_values.json" - ) + golden_path = Path(__file__).parents[1] / "benchmarks" / "data" / "dcdh_dynr_golden_values.json" if not golden_path.exists(): pytest.skip( f"dCDH golden values file not found at {golden_path}; " "run: Rscript benchmarks/R/generate_dcdh_dynr_test_values.R" ) with open(golden_path) as f: - sc = json.load(f)["scenarios"].get( - "single_baseline_multi_path_by_path_trends_lin" - ) + sc = json.load(f)["scenarios"].get("single_baseline_multi_path_by_path_trends_lin") if sc is None: - pytest.skip( - "scenario 'single_baseline_multi_path_by_path_trends_lin' absent" - ) + pytest.skip("scenario 'single_baseline_multi_path_by_path_trends_lin' absent") return pd.DataFrame(sc["data"]) def _load_by_path_trends_nonparam_scenario(): """Load golden-value scenario for by_path + trends_nonparam.""" - golden_path = ( - Path(__file__).parents[1] - / "benchmarks" - / "data" - / "dcdh_dynr_golden_values.json" - ) + golden_path = Path(__file__).parents[1] / "benchmarks" / "data" / "dcdh_dynr_golden_values.json" if not golden_path.exists(): pytest.skip( f"dCDH golden values file not found at {golden_path}; " "run: Rscript benchmarks/R/generate_dcdh_dynr_test_values.R" ) with open(golden_path) as f: - sc = json.load(f)["scenarios"].get( - "multi_path_reversible_by_path_trends_nonparam" - ) + sc = json.load(f)["scenarios"].get("multi_path_reversible_by_path_trends_nonparam") if sc is None: - pytest.skip( - "scenario 'multi_path_reversible_by_path_trends_nonparam' absent" - ) + pytest.skip("scenario 'multi_path_reversible_by_path_trends_nonparam' absent") return pd.DataFrame(sc["data"]) @@ -7056,9 +7364,7 @@ def test_path_effects_present_under_trends_linear(self): assert res.path_effects is not None and len(res.path_effects) > 0 for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["effect"]), ( - f"path={path} l={l_h}: DID^{{fd}}_l not finite" - ) + assert np.isfinite(vals["effect"]), f"path={path} l={l_h}: DID^{{fd}}_l not finite" def test_path_cumulated_event_study_present(self): """path_cumulated_event_study populated under trends_linear=True.""" @@ -7076,17 +7382,17 @@ def test_path_cumulated_event_study_present(self): L_max=3, ) assert res.path_cumulated_event_study is not None - assert set(res.path_cumulated_event_study.keys()) == set( - res.path_effects.keys() - ) + assert set(res.path_cumulated_event_study.keys()) == set(res.path_effects.keys()) for path, h_dict in res.path_cumulated_event_study.items(): - assert set(h_dict.keys()) == {1, 2, 3}, ( - f"path={path}: expected horizons 1..3, got {sorted(h_dict.keys())}" - ) + assert set(h_dict.keys()) == { + 1, + 2, + 3, + }, f"path={path}: expected horizons 1..3, got {sorted(h_dict.keys())}" for l_h, vals in h_dict.items(): - assert np.isfinite(vals["effect"]), ( - f"path={path} l={l_h}: cumulated effect not finite" - ) + assert np.isfinite( + vals["effect"] + ), f"path={path} l={l_h}: cumulated effect not finite" def test_path_cumulated_is_none_without_trends_linear(self): """path_cumulated_event_study is None when trends_linear=False.""" @@ -7266,9 +7572,9 @@ def test_per_period_effects_unaffected_by_trends_linear_by_path(self): f"vs bp={bp_pp is not None})" ) if no_bp_pp is not None and bp_pp is not None: - assert set(no_bp_pp.keys()) == set(bp_pp.keys()), ( - f"per_period_effects horizon set differs" - ) + assert set(no_bp_pp.keys()) == set( + bp_pp.keys() + ), "per_period_effects horizon set differs" for t_h in no_bp_pp: for field_name in ("did_plus_t", "did_minus_t"): if field_name not in no_bp_pp[t_h]: @@ -7310,9 +7616,7 @@ def test_bootstrap_with_trends_linear_finite_se(self): ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: bootstrap SE not finite" - ) + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: bootstrap SE not finite" def test_per_path_placebos_with_trends_linear_present(self): """``path_placebo_event_study`` populated under ``by_path + @@ -7327,9 +7631,7 @@ def test_per_path_placebos_with_trends_linear_present(self): data = _by_path_data_with_trends_linear() with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True) res = est.fit( data, outcome="outcome", @@ -7346,9 +7648,9 @@ def test_per_path_placebos_with_trends_linear_present(self): # keys mirror the placebo_event_study convention. any_finite = False for path, lag_dict in res.path_placebo_event_study.items(): - assert all(k < 0 for k in lag_dict.keys()), ( - f"path={path}: placebo lag keys must be negative ints" - ) + assert all( + k < 0 for k in lag_dict.keys() + ), f"path={path}: placebo lag keys must be negative ints" for lag_k, vals in lag_dict.items(): if np.isfinite(vals["effect"]): any_finite = True @@ -7379,9 +7681,7 @@ def test_per_path_placebos_with_trends_linear_bootstrap_inference(self): data = _by_path_data_with_trends_linear() with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - est_a = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True - ) + est_a = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True) res_a = est_a.fit( data, outcome="outcome", @@ -7470,9 +7770,9 @@ def test_per_path_placebos_with_trends_linear_bootstrap_nan_consistent(self): f"path={path} lag={lag_k}: SE finite ({vals['se']}) " "under n_bootstrap=1; expected NaN" ) - assert not np.isfinite(vals["p_value"]), ( - f"path={path} lag={lag_k}: p_value finite under n_bootstrap=1" - ) + assert not np.isfinite( + vals["p_value"] + ), f"path={path} lag={lag_k}: p_value finite under n_bootstrap=1" @pytest.mark.slow def test_sup_t_bands_with_trends_linear_finite_crit(self): @@ -7506,14 +7806,12 @@ def test_sup_t_bands_with_trends_linear_finite_crit(self): any_finite = True break assert any_finite, ( - "No path produced a finite sup-t crit value under " - "trends_linear + bootstrap" + "No path produced a finite sup-t crit value under " "trends_linear + bootstrap" ) df_bp = res.to_dataframe(level="by_path") positive = df_bp[df_bp["horizon"] > 0] assert positive["cband_lower"].notna().any(), ( - "No positive-horizon cband rows populated under " - "trends_linear + bootstrap" + "No positive-horizon cband rows populated under " "trends_linear + bootstrap" ) @pytest.mark.slow @@ -7584,9 +7882,7 @@ def test_multi_baseline_panel_emits_r_deviation_warning(self): def _add(group, treatment_path): for t, d in enumerate(treatment_path): y = d * 2.0 + rng.normal(0, 0.1) + 0.1 * t - rows.append( - {"group": group, "period": t, "treatment": d, "outcome": y} - ) + rows.append({"group": group, "period": t, "treatment": d, "outcome": y}) # F_g=3 joiners (path 0,0,1,1,1,1) for g in (1, 2, 3): @@ -7604,11 +7900,7 @@ def _add(group, treatment_path): # Sanity: switchers have both D_{g,1}=0 and D_{g,1}=1 baselines switcher_ids = data[data["group"].isin([1, 2, 3, 4, 5, 6])] - baselines = ( - switcher_ids[switcher_ids["period"] == 0] - .groupby("group")["treatment"] - .first() - ) + baselines = switcher_ids[switcher_ids["period"] == 0].groupby("group")["treatment"].first() assert sorted(baselines.unique()) == [0, 1] with warnings.catch_warnings(record=True) as caught: @@ -7685,9 +7977,7 @@ def test_F_g_three_boundary_case_emits_warning(self): def _add(group, treatment_path): for t, d in enumerate(treatment_path): y = d * 2.0 + 0.05 * group + rng.normal(0, 0.1) - rows.append( - {"group": group, "period": t, "treatment": d, "outcome": y} - ) + rows.append({"group": group, "period": t, "treatment": d, "outcome": y}) for g in (1, 2, 3, 4): _add(g, [0, 0, 1, 1, 1, 1, 1, 1]) # F_g=3 path @@ -7735,12 +8025,8 @@ def test_single_baseline_heterogeneous_F_g_does_not_warn(self): # across 3 paths; all switchers have D_{g,1}=0. data = _by_path_data_with_trends_linear() # Sanity: F_g varies across switchers - switcher_first_treat = ( - data[data["treatment"] == 1].groupby("group")["period"].min() - ) - all_groups_first_treat = ( - data.groupby("group")["treatment"].agg(lambda x: x.iloc[0]) - ) + switcher_first_treat = data[data["treatment"] == 1].groupby("group")["period"].min() + all_groups_first_treat = data.groupby("group")["treatment"].agg(lambda x: x.iloc[0]) # Drop always-treated (D_{g,1}=1) groups to isolate switchers switcher_groups = all_groups_first_treat[all_groups_first_treat == 0].index switcher_F_g = switcher_first_treat[switcher_first_treat.index.isin(switcher_groups)] @@ -7748,8 +8034,7 @@ def test_single_baseline_heterogeneous_F_g_does_not_warn(self): # treated, 20 never-treated; switchers all have D_{g,1}=0 and # F_g spans {4,5,6,7,8,9} (6 distinct values) assert switcher_F_g.nunique() >= 2, ( - "Test fixture pre-condition violated: F_g should be heterogeneous " - "across switchers" + "Test fixture pre-condition violated: F_g should be heterogeneous " "across switchers" ) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") @@ -7770,9 +8055,7 @@ def test_single_baseline_heterogeneous_F_g_does_not_warn(self): and "+ trends_linear" in str(w.message) and "switcher baselines" in str(w.message) ] - assert not deviation_msgs, ( - f"Heterogeneous F_g triggered the warning: {deviation_msgs}" - ) + assert not deviation_msgs, f"Heterogeneous F_g triggered the warning: {deviation_msgs}" # Sanity: fit produced finite per-path effects assert res.path_effects is not None and len(res.path_effects) >= 1 @@ -7813,16 +8096,16 @@ def test_bootstrap_cumulated_nan_consistent_when_n_bootstrap_one(self): f"({vals['se']}) under n_bootstrap=1; expected NaN per " "the NaN-on-invalid bootstrap contract" ) - assert not np.isfinite(vals["t_stat"]), ( - f"path={path} l={l_h}: cumulated t_stat not NaN" - ) - assert not np.isfinite(vals["p_value"]), ( - f"path={path} l={l_h}: cumulated p_value not NaN" - ) + assert not np.isfinite( + vals["t_stat"] + ), f"path={path} l={l_h}: cumulated t_stat not NaN" + assert not np.isfinite( + vals["p_value"] + ), f"path={path} l={l_h}: cumulated p_value not NaN" ci_lo, ci_hi = vals["conf_int"] - assert not (np.isfinite(ci_lo) and np.isfinite(ci_hi)), ( - f"path={path} l={l_h}: cumulated conf_int not NaN" - ) + assert not ( + np.isfinite(ci_lo) and np.isfinite(ci_hi) + ), f"path={path} l={l_h}: cumulated conf_int not NaN" class TestByPathTrendsNonparam: @@ -7859,9 +8142,7 @@ def test_set_restriction_changes_per_path_estimates(self): data = _by_path_data_with_trends_nonparam() with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - est_no_set = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3 - ) + est_no_set = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3) res_no = est_no_set.fit( data, outcome="outcome", @@ -7870,9 +8151,7 @@ def test_set_restriction_changes_per_path_estimates(self): treatment="treatment", L_max=3, ) - est_set = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3 - ) + est_set = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3) res_set = est_set.fit( data, outcome="outcome", @@ -7920,9 +8199,9 @@ def test_per_path_se_finite(self): ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["se"]) and vals["se"] > 0, ( - f"path={path} l={l_h}: SE not positive-finite" - ) + assert ( + np.isfinite(vals["se"]) and vals["se"] > 0 + ), f"path={path} l={l_h}: SE not positive-finite" def test_time_varying_set_with_by_path_raises(self): """time-varying set assignment still rejected.""" @@ -7980,9 +8259,7 @@ def test_bootstrap_with_trends_nonparam_finite_se(self): ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: bootstrap SE not finite" - ) + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: bootstrap SE not finite" def test_per_period_effects_unaffected_by_trends_nonparam_by_path(self): """``per_period_effects`` is unaffected by the by_path + @@ -8029,7 +8306,9 @@ def test_per_period_effects_unaffected_by_trends_nonparam_by_path(self): bp_v = bp_v["effect"] if no_v is not None and np.isfinite(no_v): np.testing.assert_allclose( - bp_v, no_v, rtol=1e-12, + bp_v, + no_v, + rtol=1e-12, err_msg=( f"per_period_effects[{t_h}][{field_name}] " f"differs under by_path + trends_nonparam" @@ -8082,8 +8361,7 @@ def test_sup_t_bands_with_trends_nonparam_finite_crit(self): assert "cband_upper" in df_bp.columns positive = df_bp[df_bp["horizon"] > 0] assert positive["cband_lower"].notna().any(), ( - "No positive-horizon cband rows populated under " - "trends_nonparam + bootstrap" + "No positive-horizon cband rows populated under " "trends_nonparam + bootstrap" ) @pytest.mark.slow @@ -8145,9 +8423,7 @@ def test_per_path_placebos_with_trends_nonparam_bootstrap_inference(self): if not np.isfinite(vals_set["se"]): continue any_finite = True - vals_no = res_no.path_placebo_event_study.get(path, {}).get( - lag_k - ) + vals_no = res_no.path_placebo_event_study.get(path, {}).get(lag_k) if vals_no is None or not np.isfinite(vals_no["se"]): continue # Set restriction shrinks the control pool; with the @@ -8191,9 +8467,12 @@ def _by_path_data_with_non_binary_treatment(seed: int = 44) -> pd.DataFrame: (0, 1, 2, 2), # path 3, ramp-up ] fg_path_counts = [ - (4, 0, 18), (5, 0, 14), # path 1 = 32 - (6, 1, 14), (7, 1, 12), # path 2 = 26 - (8, 2, 12), (9, 2, 8), # path 3 = 20 + (4, 0, 18), + (5, 0, 14), # path 1 = 32 + (6, 1, 14), + (7, 1, 12), # path 2 = 26 + (8, 2, 12), + (9, 2, 8), # path 3 = 20 ] rows = [] g_id = 0 @@ -8248,8 +8527,12 @@ def test_no_longer_raises_on_non_binary(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_effects is not None assert len(res.path_effects) == 2 @@ -8272,8 +8555,12 @@ def test_non_integer_D_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) def test_negative_integer_D_supported(self): @@ -8309,8 +8596,12 @@ def test_negative_integer_D_supported(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_effects is not None path_keys = set(res.path_effects.keys()) @@ -8369,19 +8660,29 @@ def test_negative_baseline_path_supported(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_effects is not None path_keys = set(res.path_effects.keys()) # Both negative-baseline paths must appear with full negative # baseline preserved in the tuple key. - assert (-1, 0, 0, 0) in path_keys, ( - f"Expected (-1, 0, 0, 0) in path keys; got {sorted(path_keys)}" - ) - assert (-1, 1, 1, 1) in path_keys, ( - f"Expected (-1, 1, 1, 1) in path keys; got {sorted(path_keys)}" - ) + assert ( + -1, + 0, + 0, + 0, + ) in path_keys, f"Expected (-1, 0, 0, 0) in path keys; got {sorted(path_keys)}" + assert ( + -1, + 1, + 1, + 1, + ) in path_keys, f"Expected (-1, 1, 1, 1) in path keys; got {sorted(path_keys)}" def test_path_effects_present_under_non_binary(self): """path_effects populated; tuple keys are non-binary.""" @@ -8392,8 +8693,12 @@ def test_path_effects_present_under_non_binary(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_effects is not None assert len(res.path_effects) == 3 @@ -8405,9 +8710,7 @@ def test_path_effects_present_under_non_binary(self): ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["effect"]), ( - f"path={path} l={l_h}: effect not finite" - ) + assert np.isfinite(vals["effect"]), f"path={path} l={l_h}: effect not finite" def test_per_period_effects_unaffected_by_non_binary_by_path(self): """per_period_effects is unchanged by the by_path lift.""" @@ -8421,12 +8724,20 @@ def test_per_period_effects_unaffected_by_non_binary_by_path(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_no = est_no.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) res_bp = est_bp.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # per_period_effects is computed before path enumeration; bit-identical. for t in res_no.per_period_effects.keys(): @@ -8435,9 +8746,9 @@ def test_per_period_effects_unaffected_by_non_binary_by_path(self): v_bp = res_bp.per_period_effects[t].get(k) if v_no is None or not np.isfinite(v_no): continue - assert np.isclose(v_no, v_bp, atol=1e-14, rtol=1e-14), ( - f"per_period_effects[{t}][{k}]: {v_no} != {v_bp}" - ) + assert np.isclose( + v_no, v_bp, atol=1e-14, rtol=1e-14 + ), f"per_period_effects[{t}][{k}]: {v_no} != {v_bp}" def test_to_dataframe_by_path_with_non_binary(self): """level='by_path' DataFrame includes non-binary path-tuple labels.""" @@ -8448,8 +8759,12 @@ def test_to_dataframe_by_path_with_non_binary(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) out = res.to_dataframe(level="by_path") assert len(out) > 0 @@ -8477,15 +8792,17 @@ def test_continuous_D_without_by_path_unaffected(self): + 2.0 * df["treatment"].values + rng.normal(0, 0.5, size=len(df)) ) - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, twfe_diagnostic=False, seed=42 - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, twfe_diagnostic=False, seed=42) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) # No by_path; the new D-integer validation does not fire. res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=2, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=2, ) assert np.isfinite(res.overall_att) @@ -8494,34 +8811,46 @@ def test_bootstrap_with_non_binary_finite_se(self): """Bootstrap SE finite on every path under non-binary D.""" df = _by_path_data_with_non_binary_treatment() est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, n_bootstrap=200, - twfe_diagnostic=False, seed=42, + drop_larger_lower=False, + by_path=3, + n_bootstrap=200, + twfe_diagnostic=False, + seed=42, ) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: bootstrap SE not finite" - ) + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: bootstrap SE not finite" @pytest.mark.slow def test_per_path_placebos_with_non_binary_present(self): """path_placebo_event_study populated under non-binary + placebo.""" df = _by_path_data_with_non_binary_treatment() est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True, - twfe_diagnostic=False, seed=42, + drop_larger_lower=False, + by_path=3, + placebo=True, + twfe_diagnostic=False, + seed=42, ) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_placebo_event_study is not None assert len(res.path_placebo_event_study) > 0 @@ -8541,14 +8870,21 @@ def test_sup_t_bands_with_non_binary_finite_crit(self): """Per-path sup-t crit_value finite under non-binary D.""" df = _by_path_data_with_non_binary_treatment() est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, n_bootstrap=400, - twfe_diagnostic=False, seed=42, + drop_larger_lower=False, + by_path=3, + n_bootstrap=400, + twfe_diagnostic=False, + seed=42, ) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_sup_t_bands is not None # At least one path passed the strict-majority gate. @@ -8557,9 +8893,9 @@ def test_sup_t_bands_with_non_binary_finite_crit(self): for entry in res.path_sup_t_bands.values() if np.isfinite(entry["crit_value"]) ] - assert len(finite_crits) > 0, ( - "Expected at least one finite crit_value under non-binary D + bootstrap" - ) + assert ( + len(finite_crits) > 0 + ), "Expected at least one finite crit_value under non-binary D + bootstrap" # --------------------------------------------------------------------------- @@ -8590,26 +8926,18 @@ def test_non_tuple_path_raises(self): def test_non_int_element_raises(self): with pytest.raises(ValueError, match="must be an int"): - ChaisemartinDHaultfoeuille( - paths_of_interest=[(0, "a", 1, 1)] - ) + ChaisemartinDHaultfoeuille(paths_of_interest=[(0, "a", 1, 1)]) def test_bool_element_raises(self): with pytest.raises(ValueError, match="must be an int"): - ChaisemartinDHaultfoeuille( - paths_of_interest=[(False, True, True, True)] - ) + ChaisemartinDHaultfoeuille(paths_of_interest=[(False, True, True, True)]) def test_np_bool_element_raises(self): with pytest.raises(ValueError, match="must be an int"): - ChaisemartinDHaultfoeuille( - paths_of_interest=[(np.bool_(True), 0, 0, 0)] - ) + ChaisemartinDHaultfoeuille(paths_of_interest=[(np.bool_(True), 0, 0, 0)]) def test_np_integer_accepted_canonicalized(self): - est = ChaisemartinDHaultfoeuille( - paths_of_interest=[(np.int64(0), np.int32(1), 1, 1)] - ) + est = ChaisemartinDHaultfoeuille(paths_of_interest=[(np.int64(0), np.int32(1), 1, 1)]) # Canonicalized to Python int tuples. assert est.paths_of_interest == [(0, 1, 1, 1)] for v in est.paths_of_interest[0]: @@ -8617,15 +8945,11 @@ def test_np_integer_accepted_canonicalized(self): def test_mixed_lengths_raise(self): with pytest.raises(ValueError, match="mixed lengths"): - ChaisemartinDHaultfoeuille( - paths_of_interest=[(0, 1), (0, 1, 1, 1)] - ) + ChaisemartinDHaultfoeuille(paths_of_interest=[(0, 1), (0, 1, 1, 1)]) def test_mutex_with_by_path_raises(self): with pytest.raises(ValueError, match="mutually exclusive"): - ChaisemartinDHaultfoeuille( - by_path=2, paths_of_interest=[(0, 1, 1, 1)] - ) + ChaisemartinDHaultfoeuille(by_path=2, paths_of_interest=[(0, 1, 1, 1)]) def test_set_params_re_validates_mutex_by_path_added(self): est = ChaisemartinDHaultfoeuille(paths_of_interest=[(0, 1, 1, 1)]) @@ -8662,9 +8986,7 @@ def test_set_params_failed_validation_is_transactional(self): assert params["paths_of_interest"] is None def test_get_params_includes_paths_of_interest(self): - est = ChaisemartinDHaultfoeuille( - paths_of_interest=[(0, 1, 1, 1), (0, 1, 0, 0)] - ) + est = ChaisemartinDHaultfoeuille(paths_of_interest=[(0, 1, 1, 1), (0, 1, 0, 0)]) params = est.get_params() assert "paths_of_interest" in params assert params["paths_of_interest"] == [(0, 1, 1, 1), (0, 1, 0, 0)] @@ -8687,8 +9009,12 @@ def test_canonicalized_duplicates_dedup_warn(self): with warnings.catch_warnings(): warnings.simplefilter("default", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # ---- fit-time validation ---- @@ -8705,8 +9031,12 @@ def test_wrong_length_raises_at_fit(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) def test_paths_of_interest_requires_L_max(self): @@ -8721,7 +9051,10 @@ def test_paths_of_interest_requires_L_max(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", + df, + outcome="outcome", + group="group", + time="period", treatment="treatment", ) @@ -8736,8 +9069,12 @@ def test_paths_of_interest_requires_drop_larger_lower_false(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # ---- behavior ---- @@ -8753,8 +9090,12 @@ def test_paths_of_interest_selects_user_paths(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert set(res.path_effects.keys()) == {(0, 1, 1, 1), (0, 1, 0, 0)} @@ -8771,8 +9112,12 @@ def test_paths_of_interest_preserves_user_order(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # Insertion order preserved. assert list(res.path_effects.keys()) == user_order @@ -8794,8 +9139,12 @@ def test_paths_of_interest_order_preserved_in_to_dataframe(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) out = res.to_dataframe(level="by_path") # First-occurrence order of the path column matches user order. @@ -8822,8 +9171,12 @@ def test_paths_of_interest_order_preserved_in_summary(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) text = res.summary() # Find the ordering of the user paths as they appear in summary. @@ -8853,8 +9206,12 @@ def test_paths_of_interest_frequency_rank_is_true_frequency(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # (0,1,1,1) has higher frequency → rank 1 # (0,1,0,0) has lower frequency → rank 2 @@ -8877,8 +9234,12 @@ def test_paths_of_interest_all_unobserved_summary_distinct_text(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) text = res.summary() assert "Every path in paths_of_interest was unobserved" in text, ( @@ -8903,8 +9264,12 @@ def test_paths_of_interest_all_unobserved_emits_distinct_warning(self): with warnings.catch_warnings(record=True) as recorded: warnings.simplefilter("always", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) # The summary empty-state warning mentions paths_of_interest, not by_path empty_state_warnings = [ @@ -8934,8 +9299,12 @@ def test_unobserved_path_warns_and_omits(self): with warnings.catch_warnings(): warnings.simplefilter("default", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert (1, 1, 1, 1) not in res.path_effects assert (0, 1, 1, 1) in res.path_effects @@ -8952,8 +9321,12 @@ def test_paths_of_interest_with_non_binary_D(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert set(res.path_effects.keys()) == {(0, 2, 2, 2)} for l_h, vals in res.path_effects[(0, 2, 2, 2)]["horizons"].items(): @@ -8975,8 +9348,13 @@ def test_paths_of_interest_with_controls(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, controls=["X1"], + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + controls=["X1"], ) assert len(res.path_effects) == 2 for path, entry in res.path_effects.items(): @@ -8995,8 +9373,13 @@ def test_paths_of_interest_with_trends_linear(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, trends_linear=True, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + trends_linear=True, ) assert len(res.path_effects) == 2 # path_cumulated_event_study populated under trends_linear. @@ -9015,8 +9398,12 @@ def test_paths_of_interest_with_trends_nonparam(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, trends_nonparam="state", ) assert len(res.path_effects) == 2 @@ -9044,8 +9431,12 @@ def test_paths_of_interest_non_binary_bootstrap_placebo(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert set(res.path_effects.keys()) == {(0, 2, 2, 2), (0, 1, 1, 1)} # 1. analytical / bootstrap on path_effects (SE finite) @@ -9074,14 +9465,11 @@ def test_paths_of_interest_non_binary_bootstrap_placebo(self): # cband_conf_int populated for positive horizons of finite-crit paths. for path in finite_crit_paths: for l_h in range(1, 4): - cband = res.path_effects[path]["horizons"][l_h].get( - "cband_conf_int" - ) + cband = res.path_effects[path]["horizons"][l_h].get("cband_conf_int") assert cband is not None, f"path={path} l={l_h}: cband missing" lo, hi = cband assert np.isfinite(lo) and np.isfinite(hi), ( - f"path={path} l={l_h}: cband endpoints not finite " - f"(lo={lo}, hi={hi})" + f"path={path} l={l_h}: cband endpoints not finite " f"(lo={lo}, hi={hi})" ) @pytest.mark.slow @@ -9106,17 +9494,20 @@ def test_paths_of_interest_trends_linear_bootstrap_placebo(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, trends_linear=True, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + trends_linear=True, ) # (i) selector restricts the path set assert set(res.path_effects.keys()) == set(user_paths) # (ii) per-path bootstrap SE finite on event study for path in res.path_effects: for l_h, vals in res.path_effects[path]["horizons"].items(): - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: bootstrap SE not finite" - ) + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: bootstrap SE not finite" # (iii) post-bootstrap path_cumulated_event_study populated for # the same paths AND derived from the post-bootstrap per-horizon # SEs (cumulated SE = sum of post-bootstrap component SEs). @@ -9125,12 +9516,10 @@ def test_paths_of_interest_trends_linear_bootstrap_placebo(self): for path in user_paths: assert len(res.path_cumulated_event_study[path]) > 0 for l_h, vals in res.path_cumulated_event_study[path].items(): - assert np.isfinite(vals["effect"]), ( - f"path={path} l={l_h}: cumulated effect not finite" - ) - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: cumulated SE not finite" - ) + assert np.isfinite( + vals["effect"] + ), f"path={path} l={l_h}: cumulated effect not finite" + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: cumulated SE not finite" # (iv) per-path placebo populated assert res.path_placebo_event_study is not None assert set(res.path_placebo_event_study.keys()) == set(user_paths) @@ -9157,8 +9546,12 @@ def test_paths_of_interest_trends_nonparam_bootstrap_placebo(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, trends_nonparam="state", ) # Selector restriction @@ -9201,8 +9594,12 @@ def test_bootstrap_with_paths_of_interest_finite_se(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) for path, entry in res.path_effects.items(): for l_h, vals in entry["horizons"].items(): @@ -9221,8 +9618,12 @@ def test_per_path_placebos_with_paths_of_interest_present(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res.path_placebo_event_study is not None assert len(res.path_placebo_event_study) == 2 @@ -9258,19 +9659,33 @@ def _by_path_survey_data(seed: int = 44) -> pd.DataFrame: else: d = 0 y = 0.5 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) for g in range(30, 60): stratum = (g - 30) % 4 weight = 1.0 + 0.1 * ((g - 30) % 5) for t in range(n_periods): y = rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) return pd.DataFrame(rows) @@ -9297,19 +9712,33 @@ def _by_path_survey_data_single_path(seed: int = 44) -> pd.DataFrame: else: d = 0 y = 0.5 * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) for g in range(30, 60): stratum = (g - 30) % 4 weight = 1.0 + 0.1 * ((g - 30) % 5) for t in range(n_periods): y = rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) return pd.DataFrame(rows) @@ -9334,8 +9763,13 @@ def test_no_longer_raises_on_survey(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_effects is not None assert len(res.path_effects) >= 1 @@ -9352,8 +9786,13 @@ def test_paths_of_interest_with_survey_no_longer_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_effects is not None assert (0, 1, 1, 1) in res.path_effects @@ -9371,8 +9810,13 @@ def test_survey_design_plus_n_bootstrap_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) def test_survey_design_plus_paths_of_interest_plus_n_bootstrap_raises(self): @@ -9390,8 +9834,13 @@ def test_survey_design_plus_paths_of_interest_plus_n_bootstrap_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) def test_global_survey_plus_n_bootstrap_still_works(self): @@ -9414,8 +9863,13 @@ def test_global_survey_plus_n_bootstrap_still_works(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert np.isfinite(res.overall_se) assert res.path_effects is None @@ -9431,8 +9885,13 @@ def test_per_path_analytical_se_finite_under_survey(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_effects is not None for path, entry in res.path_effects.items(): @@ -9460,12 +9919,22 @@ def test_per_path_se_telescope_to_global_on_single_path(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_g = est_g.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) res_p = est_p.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res_p.path_effects is not None assert len(res_p.path_effects) == 1 @@ -9510,12 +9979,21 @@ def test_per_path_se_within_envelope_of_unweighted(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_survey = est_p.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) res_plain = est_p_no_survey.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, ) assert res_survey.path_effects is not None and res_plain.path_effects is not None any_se_compared = False @@ -9534,7 +10012,9 @@ def test_per_path_se_within_envelope_of_unweighted(self): se_plain = res_plain.path_effects[path]["horizons"][l_h]["se"] if np.isfinite(se_survey) and np.isfinite(se_plain): np.testing.assert_allclose( - se_survey, se_plain, rtol=0.10, + se_survey, + se_plain, + rtol=0.10, err_msg=( f"path={path} l={l_h}: survey SE outside 10% " f"rtol envelope of plug-in SE" @@ -9570,8 +10050,13 @@ def test_per_path_replicate_se_finite(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_effects is not None any_finite = False @@ -9644,18 +10129,28 @@ def wrapped_compute_se(*args, **kwargs): return se, forced_low_n_valid return se, n_valid - with _mock.patch.object( - _cd_mod, "_compute_path_effects", - side_effect=wrapped_path_effects, - ), _mock.patch.object( - _cd_mod, "_compute_se", - side_effect=wrapped_compute_se, + with ( + _mock.patch.object( + _cd_mod, + "_compute_path_effects", + side_effect=wrapped_path_effects, + ), + _mock.patch.object( + _cd_mod, + "_compute_se", + side_effect=wrapped_compute_se, + ), ): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) # The forced low n_valid (5) at later IF sites bounds the final @@ -9682,22 +10177,30 @@ def wrapped_compute_se(*args, **kwargs): if vals["n_obs"] == 0 or not np.isfinite(vals["se"]): continue t_final, p_final, ci_final = safe_inference( - vals["effect"], vals["se"], - alpha=est.alpha, df=expected_low_df, + vals["effect"], + vals["se"], + alpha=est.alpha, + df=expected_low_df, ) np.testing.assert_allclose( - vals["t_stat"], t_final, atol=1e-12, + vals["t_stat"], + t_final, + atol=1e-12, err_msg=( f"path={path} l={l_h}: t_stat reflects stale " f"snapshot df, not final df={expected_low_df}" ), ) np.testing.assert_allclose( - vals["p_value"], p_final, atol=1e-12, + vals["p_value"], + p_final, + atol=1e-12, err_msg=f"path={path} l={l_h}: p_value stale", ) np.testing.assert_allclose( - vals["conf_int"], ci_final, atol=1e-12, + vals["conf_int"], + ci_final, + atol=1e-12, err_msg=f"path={path} l={l_h}: conf_int stale", ) any_compared = True @@ -9763,8 +10266,13 @@ def test_refresh_path_inference_called_from_final_block(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) # Helper called exactly once from the final R2 P1b block. @@ -9822,14 +10330,17 @@ def test_per_path_inference_uses_final_df_after_all_appends(self): replicate_method="JK1", replicate_scale=1.0, ) - est = ChaisemartinDHaultfoeuille( - by_path=2, drop_larger_lower=False, placebo=True - ) + est = ChaisemartinDHaultfoeuille(by_path=2, drop_larger_lower=False, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.survey_metadata is not None df_final = res.survey_metadata.df_survey @@ -9843,18 +10354,27 @@ def test_per_path_inference_uses_final_df_after_all_appends(self): if vals["n_obs"] == 0 or not np.isfinite(vals["se"]): continue t_exp, p_exp, ci_exp = safe_inference( - vals["effect"], vals["se"], alpha=est.alpha, df=df_final, + vals["effect"], + vals["se"], + alpha=est.alpha, + df=df_final, ) np.testing.assert_allclose( - vals["t_stat"], t_exp, atol=1e-12, + vals["t_stat"], + t_exp, + atol=1e-12, err_msg=f"path={path} l={l_h} t_stat stale", ) np.testing.assert_allclose( - vals["p_value"], p_exp, atol=1e-12, + vals["p_value"], + p_exp, + atol=1e-12, err_msg=f"path={path} l={l_h} p_value stale", ) np.testing.assert_allclose( - vals["conf_int"], ci_exp, atol=1e-12, + vals["conf_int"], + ci_exp, + atol=1e-12, err_msg=f"path={path} l={l_h} conf_int stale", ) any_checked = True @@ -9865,7 +10385,10 @@ def test_per_path_inference_uses_final_df_after_all_appends(self): if vals["n_obs"] == 0 or not np.isfinite(vals["se"]): continue t_exp, p_exp, ci_exp = safe_inference( - vals["effect"], vals["se"], alpha=est.alpha, df=df_final, + vals["effect"], + vals["se"], + alpha=est.alpha, + df=df_final, ) np.testing.assert_allclose(vals["t_stat"], t_exp, atol=1e-12) np.testing.assert_allclose(vals["p_value"], p_exp, atol=1e-12) @@ -9897,8 +10420,13 @@ def test_per_path_replicate_n_valid_propagates_to_df_survey(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) # df_survey reflects replicate columns; cap is 15 - 1 = 14. assert res.survey_metadata is not None @@ -9913,14 +10441,17 @@ def test_per_path_placebo_se_finite_under_survey(self): df = _by_path_survey_data() sd = SurveyDesign(weights="survey_weights", strata="strata", psu="psu") - est = ChaisemartinDHaultfoeuille( - by_path=2, drop_larger_lower=False, placebo=True - ) + est = ChaisemartinDHaultfoeuille(by_path=2, drop_larger_lower=False, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_placebo_event_study is not None any_finite = False @@ -9953,28 +10484,47 @@ def test_per_path_cumulated_se_inherits_survey(self): else: d = 0 y = 0.5 * d + trend * t + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) for g in range(30, 60): stratum = (g - 30) % 4 weight = 1.0 + 0.1 * ((g - 30) % 5) trend = 0.05 * g for t in range(n_periods): y = trend * t + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, "outcome": y, - "survey_weights": weight, "strata": stratum, "psu": g, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "survey_weights": weight, + "strata": stratum, + "psu": g, + } + ) df = pd.DataFrame(rows) sd = SurveyDesign(weights="survey_weights", strata="strata", psu="psu") est = ChaisemartinDHaultfoeuille(by_path=2, drop_larger_lower=False) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, trends_linear=True, ) # path_cumulated_event_study should populate under trends_linear @@ -9995,16 +10545,20 @@ def test_path_unobserved_under_survey_warns_omits(self): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res.path_effects is not None assert (0, 1, 1, 1) in res.path_effects assert (0, 9, 9, 9) not in res.path_effects # Unobserved-path warning must have fired assert any( - "zero observed" in str(w.message) and "(0, 9, 9, 9)" in str(w.message) - for w in caught + "zero observed" in str(w.message) and "(0, 9, 9, 9)" in str(w.message) for w in caught ) @pytest.mark.slow @@ -10185,12 +10739,22 @@ def test_telescope_analytical_TSL(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_g = est_g.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) res_p = est_p.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + survey_design=sd, ) assert res_p.path_effects is not None path = next(iter(res_p.path_effects.keys())) @@ -10236,20 +10800,30 @@ def _by_path_het_data(seed=44, n_switchers=90, n_controls=30, n_periods=10): else: d = 0 y = 0.5 * t + effect * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": het_x, + } + ) # Never-treated controls (D=0 throughout), het_x balanced for k in range(n_controls): het_x = 1 if k < n_controls // 2 else 0 g = n_switchers + k for t in range(n_periods): y = 0.5 * t + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "het_x": het_x, + } + ) return pd.DataFrame(rows) @@ -10271,8 +10845,13 @@ def test_no_longer_raises_on_heterogeneity(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects is not None @@ -10286,8 +10865,13 @@ def test_paths_of_interest_with_heterogeneity_no_longer_raises(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects is not None @@ -10296,41 +10880,46 @@ def test_heterogeneity_still_rejects_controls_under_by_path(self): df = _by_path_het_data() df["X1"] = np.random.RandomState(42).normal(0, 1, len(df)) with pytest.raises(ValueError, match="cannot be combined with controls"): - ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2 - ).fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, - heterogeneity="het_x", controls=["X1"], + ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2).fit( + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", + controls=["X1"], ) def test_heterogeneity_still_rejects_trends_linear_under_by_path(self): """``heterogeneity + trends_linear`` mutex still fires under by_path.""" df = _by_path_het_data() - with pytest.raises( - ValueError, match="cannot be combined with trends_linear" - ): - ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2 - ).fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, - heterogeneity="het_x", trends_linear=True, + with pytest.raises(ValueError, match="cannot be combined with trends_linear"): + ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2).fit( + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", + trends_linear=True, ) def test_heterogeneity_still_rejects_trends_nonparam_under_by_path(self): """``heterogeneity + trends_nonparam`` mutex still fires under by_path.""" df = _by_path_het_data() df["state"] = df["group"] % 3 - with pytest.raises( - ValueError, match="cannot be combined with trends_nonparam" - ): - ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2 - ).fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, - heterogeneity="het_x", trends_nonparam="state", + with pytest.raises(ValueError, match="cannot be combined with trends_nonparam"): + ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2).fit( + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", + trends_nonparam="state", ) # Behavior @@ -10343,8 +10932,13 @@ def test_per_path_heterogeneity_finite_under_known_signal(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects # At horizon 1 every path has switchers; heterogeneity beta should @@ -10376,8 +10970,13 @@ def test_per_path_heterogeneity_inference_local_invariants(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects checked = 0 @@ -10392,22 +10991,14 @@ def test_per_path_heterogeneity_inference_local_invariants(self): ) half_low = het["beta"] - het["conf_int"][0] half_high = het["conf_int"][1] - het["beta"] - assert half_low > 0, ( - f"path={path} l={l_h} conf_int_lower not below beta" - ) - assert half_high > 0, ( - f"path={path} l={l_h} conf_int_upper not above beta" - ) - assert half_low == pytest.approx(half_high, rel=1e-12), ( - f"path={path} l={l_h} conf_int asymmetric" - ) - assert 0.0 <= het["p_value"] <= 1.0, ( - f"path={path} l={l_h} p_value out of [0, 1]" - ) + assert half_low > 0, f"path={path} l={l_h} conf_int_lower not below beta" + assert half_high > 0, f"path={path} l={l_h} conf_int_upper not above beta" + assert half_low == pytest.approx( + half_high, rel=1e-12 + ), f"path={path} l={l_h} conf_int asymmetric" + assert 0.0 <= het["p_value"] <= 1.0, f"path={path} l={l_h} p_value out of [0, 1]" checked += 1 - assert checked >= 1, ( - "Expected at least one populated (path, horizon) heterogeneity entry" - ) + assert checked >= 1, "Expected at least one populated (path, horizon) heterogeneity entry" def test_per_path_heterogeneity_telescope_to_global_on_single_path(self): """On a single-path panel, per-path == global heterogeneity. @@ -10429,35 +11020,53 @@ def test_per_path_heterogeneity_telescope_to_global_on_single_path(self): d = path[-1] else: d = 0 - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": 0.5 * t + effect * d + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": 0.5 * t + effect * d + rng.normal(0, 0.5), + "het_x": het_x, + } + ) for k in range(n_controls): het_x = 1 if k < n_controls // 2 else 0 for t in range(10): - rows.append({ - "group": n_switchers + k, "period": t, "treatment": 0, - "outcome": 0.5 * t + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": n_switchers + k, + "period": t, + "treatment": 0, + "outcome": 0.5 * t + rng.normal(0, 0.5), + "het_x": het_x, + } + ) df = pd.DataFrame(rows) # Run with by_path=1 (path is observed) est_p = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=1) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_p = est_p.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) # Run global (no by_path) est_g = ChaisemartinDHaultfoeuille(drop_larger_lower=False) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_g = est_g.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res_p.path_heterogeneity_effects path_key = (0, 1, 1, 1) @@ -10469,11 +11078,17 @@ def test_per_path_heterogeneity_telescope_to_global_on_single_path(self): assert not np.isfinite(py_global["beta"]) continue np.testing.assert_allclose( - py_path["beta"], py_global["beta"], atol=1e-14, rtol=1e-14, + py_path["beta"], + py_global["beta"], + atol=1e-14, + rtol=1e-14, err_msg=f"l={l_h}: per-path beta != global beta (telescope failed)", ) np.testing.assert_allclose( - py_path["se"], py_global["se"], atol=1e-14, rtol=1e-14, + py_path["se"], + py_global["se"], + atol=1e-14, + rtol=1e-14, err_msg=f"l={l_h}: per-path se != global se (telescope failed)", ) @@ -10497,27 +11112,40 @@ def test_per_path_heterogeneity_zero_signal_yields_small_beta(self): else: d = 0 # Effect is constant 5.0 — no heterogeneity by het_x - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), + "het_x": het_x, + } + ) for k in range(n_controls): # Draw het_x ONCE per group (must be time-invariant) het_x = rng.normal(0, 1) for t in range(10): - rows.append({ - "group": n_switchers + k, "period": t, "treatment": 0, - "outcome": 0.5 * t + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": n_switchers + k, + "period": t, + "treatment": 0, + "outcome": 0.5 * t + rng.normal(0, 0.5), + "het_x": het_x, + } + ) df = pd.DataFrame(rows) est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects for path, horizons in res.path_heterogeneity_effects.items(): @@ -10549,11 +11177,15 @@ def test_path_with_too_few_eligible_yields_nan(self): d = path[-1] else: d = 0 - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), + "het_x": het_x, + } + ) # 2 switchers on the rare path — under-eligible for g in range(30, 32): F_g = 3 @@ -10566,18 +11198,27 @@ def test_path_with_too_few_eligible_yields_nan(self): d = path[-1] else: d = 0 - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), - "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": 0.5 * t + 5.0 * d + rng.normal(0, 0.5), + "het_x": het_x, + } + ) # Controls for k in range(15): for t in range(10): - rows.append({ - "group": 32 + k, "period": t, "treatment": 0, - "outcome": 0.5 * t + rng.normal(0, 0.5), "het_x": 0, - }) + rows.append( + { + "group": 32 + k, + "period": t, + "treatment": 0, + "outcome": 0.5 * t + rng.normal(0, 0.5), + "het_x": 0, + } + ) df = pd.DataFrame(rows) est = ChaisemartinDHaultfoeuille( drop_larger_lower=False, @@ -10586,16 +11227,19 @@ def test_path_with_too_few_eligible_yields_nan(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects rare = res.path_heterogeneity_effects[(0, 1, 0, 0)] # All horizons for the rare path should have NaN inference (n_obs < 3) for l_h, vals in rare.items(): - assert vals["n_obs"] < 3, ( - f"rare path l={l_h}: expected n_obs < 3, got {vals['n_obs']}" - ) + assert vals["n_obs"] < 3, f"rare path l={l_h}: expected n_obs < 3, got {vals['n_obs']}" assert not np.isfinite(vals["beta"]) assert not np.isfinite(vals["se"]) assert not np.isfinite(vals["t_stat"]) @@ -10630,8 +11274,7 @@ def _multi_baseline_het_data(seed=44): else: d = 0 y = 0.5 * t + effect * d + rng.normal(0, 0.5) - rows.append({"group": g, "period": t, "treatment": d, - "outcome": y, "het_x": het_x}) + rows.append({"group": g, "period": t, "treatment": d, "outcome": y, "het_x": het_x}) # Leavers: baseline=1, path (1,0,0,0) for g_offset in range(n_per_baseline): g = n_per_baseline + g_offset @@ -10647,8 +11290,7 @@ def _multi_baseline_het_data(seed=44): else: d = 1 # baseline=1 — treated pre-window y = 0.5 * t + effect * d + rng.normal(0, 0.5) - rows.append({"group": g, "period": t, "treatment": d, - "outcome": y, "het_x": het_x}) + rows.append({"group": g, "period": t, "treatment": d, "outcome": y, "het_x": het_x}) return pd.DataFrame(rows) def test_per_path_heterogeneity_no_multi_baseline_warning(self): @@ -10670,9 +11312,10 @@ def test_per_path_heterogeneity_no_multi_baseline_warning(self): df = self._multi_baseline_het_data() # Sanity check: panel actually has both baselines among switchers baselines = df.groupby("group")["treatment"].first().unique() - assert set(baselines) >= {0, 1}, ( - f"fixture must include both baselines; got {sorted(baselines)}" - ) + assert set(baselines) >= { + 0, + 1, + }, f"fixture must include both baselines; got {sorted(baselines)}" with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") @@ -10699,12 +11342,10 @@ def test_per_path_heterogeneity_no_multi_baseline_warning(self): for path in [(0, 1, 1, 1), (1, 0, 0, 0)]: horizons = res.path_heterogeneity_effects[path] finite_count = sum( - 1 for v in horizons.values() - if np.isfinite(v["beta"]) and np.isfinite(v["se"]) + 1 for v in horizons.values() if np.isfinite(v["beta"]) and np.isfinite(v["se"]) ) assert finite_count >= 1, ( - f"path={path}: expected ≥1 finite per-(path, l) entry, " - f"got {finite_count}" + f"path={path}: expected ≥1 finite per-(path, l) entry, " f"got {finite_count}" ) # No multi-baseline UserWarning. Match the controls / trends_lin @@ -10712,9 +11353,9 @@ def test_per_path_heterogeneity_no_multi_baseline_warning(self): # paths_of_interest + controls/trends_linear" R-divergence text). # Be strict — both fragments must appear in the same warning. multi_baseline = [ - w for w in caught - if "baseline" in str(w.message).lower() - and "multi" in str(w.message).lower() + w + for w in caught + if "baseline" in str(w.message).lower() and "multi" in str(w.message).lower() ] assert not multi_baseline, ( f"Unexpected multi-baseline warning(s) under heterogeneity: " @@ -10723,7 +11364,8 @@ def test_per_path_heterogeneity_no_multi_baseline_warning(self): # Also check no controls/trends-linear divergence verbatim text controls_divergence = [ - w for w in caught + w + for w in caught if "by_path / paths_of_interest + controls" in str(w.message) or "by_path / paths_of_interest + trends_linear" in str(w.message) ] @@ -10749,11 +11391,7 @@ def _by_path_het_data_with_survey(seed=44, n_replicates=0): rng = np.random.RandomState(seed) n_switchers, n_controls, n_periods = 90, 30, 10 n_groups_total = n_switchers + n_controls - H = ( - rng.choice([-1, 1], size=(n_groups_total, n_replicates)) - if n_replicates > 0 - else None - ) + H = rng.choice([-1, 1], size=(n_groups_total, n_replicates)) if n_replicates > 0 else None rows = [] paths = [(0, 1, 1, 1), (0, 1, 0, 0), (0, 1, 1, 0)] for g in range(n_switchers): @@ -10838,16 +11476,14 @@ def test_per_path_heterogeneity_under_survey_finite(self): for path, horizons in res.path_heterogeneity_effects.items(): for l_h, vals in horizons.items(): if vals["n_obs"] >= 3: - assert np.isfinite(vals["beta"]), ( - f"path={path} l={l_h}: beta is NaN under survey TSL" - ) - assert np.isfinite(vals["se"]) and vals["se"] > 0, ( - f"path={path} l={l_h}: se non-positive under survey TSL" - ) + assert np.isfinite( + vals["beta"] + ), f"path={path} l={l_h}: beta is NaN under survey TSL" + assert ( + np.isfinite(vals["se"]) and vals["se"] > 0 + ), f"path={path} l={l_h}: se non-positive under survey TSL" finite_count += 1 - assert finite_count >= 4, ( - f"Expected ≥4 finite (path, l) entries, got {finite_count}" - ) + assert finite_count >= 4, f"Expected ≥4 finite (path, l) entries, got {finite_count}" @pytest.mark.slow def test_per_path_heterogeneity_replicate_weights_propagates_n_valid(self): @@ -10888,21 +11524,18 @@ def test_per_path_heterogeneity_replicate_weights_propagates_n_valid(self): # df_survey ≤ n_replicates - 1 per Rao-Wu replicate convention. # With well-formed BRR weights and n_obs >= 3 per (path, l), we # expect every replicate fit to produce finite SE → df = 7. - assert res.survey_metadata.df_survey is not None, ( - "df_survey must be populated under replicate-weight survey" - ) + assert ( + res.survey_metadata.df_survey is not None + ), "df_survey must be populated under replicate-weight survey" assert res.survey_metadata.df_survey == n_replicates - 1, ( - f"df_survey={res.survey_metadata.df_survey}, " - f"expected {n_replicates - 1}" + f"df_survey={res.survey_metadata.df_survey}, " f"expected {n_replicates - 1}" ) # Every populated (path, l) should have finite inference under # replicate weights too. for path, horizons in res.path_heterogeneity_effects.items(): for l_h, vals in horizons.items(): if vals["n_obs"] >= 3: - assert np.isfinite(vals["se"]), ( - f"path={path} l={l_h}: replicate SE non-finite" - ) + assert np.isfinite(vals["se"]), f"path={path} l={l_h}: replicate SE non-finite" # Verify the final df_survey is actually USED to refresh the # inference fields on path_heterogeneity_effects (not the @@ -10922,16 +11555,12 @@ def test_per_path_heterogeneity_replicate_weights_propagates_n_valid(self): expected_t, expected_p, expected_ci = safe_inference( vals["beta"], vals["se"], df=df_final ) - assert vals["t_stat"] == pytest.approx( - expected_t, rel=1e-12, nan_ok=True - ), ( + assert vals["t_stat"] == pytest.approx(expected_t, rel=1e-12, nan_ok=True), ( f"path={path} l={l_h}: t_stat not refreshed at " f"df={df_final} (have {vals['t_stat']}, expected " f"{expected_t})" ) - assert vals["p_value"] == pytest.approx( - expected_p, rel=1e-12, nan_ok=True - ), ( + assert vals["p_value"] == pytest.approx(expected_p, rel=1e-12, nan_ok=True), ( f"path={path} l={l_h}: p_value not refreshed at " f"df={df_final} (have {vals['p_value']}, expected " f"{expected_p})" @@ -10989,9 +11618,10 @@ def test_paths_of_interest_heterogeneity_survey_design_analytical(self): # Selector keys are preserved in the user-specified order # (not frequency-ranked like by_path). keys = list(res.path_heterogeneity_effects.keys()) - assert keys == [(0, 1, 1, 0), (0, 1, 1, 1)], ( - f"paths_of_interest order not preserved: got {keys}" - ) + assert keys == [ + (0, 1, 1, 0), + (0, 1, 1, 1), + ], f"paths_of_interest order not preserved: got {keys}" # Every populated (path, l) entry yields finite analytical SE. for path, horizons in res.path_heterogeneity_effects.items(): for l_h, vals in horizons.items(): @@ -11046,9 +11676,7 @@ def test_survey_design_plus_n_bootstrap_with_heterogeneity_still_raises( df = self._by_path_het_data_with_survey() sd = SurveyDesign(weights="survey_weights", strata="strata", psu="psu") - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2, n_bootstrap=10, seed=1 - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2, n_bootstrap=10, seed=1) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) with pytest.raises(NotImplementedError, match="multiplier"): @@ -11071,14 +11699,17 @@ def test_to_dataframe_by_path_includes_heterogeneity_columns(self): ``placebo=True`` and ``heterogeneity=`` are both set (post-2026-05-15 #422).""" df = _by_path_het_data() - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=2, placebo=True - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=2, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) out = res.to_dataframe(level="by_path") assert "het_beta" in out.columns @@ -11113,13 +11744,18 @@ def test_per_path_heterogeneity_renders_in_summary(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) report = res.summary() - assert "Heterogeneity Test (Section 1.5, partial)" in report, ( - "summary() must render the per-path heterogeneity sub-block" - ) + assert ( + "Heterogeneity Test (Section 1.5, partial)" in report + ), "summary() must render the per-path heterogeneity sub-block" # The header appears in BOTH the global and per-path blocks; check # that at least one populated path's beta value is rendered. We # use a small float comparison rather than a full string match @@ -11155,15 +11791,20 @@ def test_path_unobserved_under_heterogeneity_warns_omits(self): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) # Unobserved path warning should have fired (at least once; # may fire from path_effects + path_heterogeneity_effects) unobs = [ - w for w in caught - if "(1, 1, 1, 0)" in str(w.message) - and "zero observed groups" in str(w.message) + w + for w in caught + if "(1, 1, 1, 0)" in str(w.message) and "zero observed groups" in str(w.message) ] assert unobs, "expected unobserved-path UserWarning" assert res.path_heterogeneity_effects is not None @@ -11194,19 +11835,29 @@ def _single_path_het_data(seed=44, n_switchers=30, n_controls=15, n_periods=10): else: d = 0 y = 0.5 * t + effect * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": het_x, + } + ) for k in range(n_controls): het_x = 1 if k < n_controls // 2 else 0 g = n_switchers + k for t in range(n_periods): y = 0.5 * t + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "het_x": het_x, + } + ) return pd.DataFrame(rows) @@ -11227,14 +11878,17 @@ class TestByPathPredictHetPlacebo: def test_to_dataframe_by_path_emits_het_columns_on_placebo_rows(self): """`to_dataframe(level="by_path")` placebo rows now have het_*.""" df = _by_path_het_data() - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) out = res.to_dataframe(level="by_path") placebo_rows = out[out["horizon"] < 0] @@ -11268,22 +11922,26 @@ def test_predict_het_placebo_survey_design_warns_and_skips_backward(self): df["stratum"] = df["group"] % 4 df["psu_id"] = df["group"] sd = SurveyDesign( - weights="sw", strata="stratum", psu="psu_id", - ) - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True + weights="sw", + strata="stratum", + psu="psu_id", ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, - heterogeneity="het_x", survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", + survey_design=sd, ) # Warning fired with the expected substring het_warnings = [ - w for w in caught - if "backward-horizon (placebo) predict_het" in str(w.message) + w for w in caught if "backward-horizon (placebo) predict_het" in str(w.message) ] assert het_warnings, ( "Expected UserWarning about backward-horizon survey gate. " @@ -11293,16 +11951,16 @@ def test_predict_het_placebo_survey_design_warns_and_skips_backward(self): assert res.heterogeneity_effects is not None # Only positive-int keys (forward); no negative (placebo) keys het_keys = sorted(res.heterogeneity_effects.keys()) - assert all(h > 0 for h in het_keys), ( - f"Expected only positive horizons under survey gate, got: {het_keys}" - ) + assert all( + h > 0 for h in het_keys + ), f"Expected only positive horizons under survey gate, got: {het_keys}" # Per-path heterogeneity also forward-only assert res.path_heterogeneity_effects is not None for path, horizons in res.path_heterogeneity_effects.items(): path_keys = sorted(horizons.keys()) - assert all(h > 0 for h in path_keys), ( - f"path={path}: expected only positive horizons, got {path_keys}" - ) + assert all( + h > 0 for h in path_keys + ), f"path={path}: expected only positive horizons, got {path_keys}" def test_compute_heterogeneity_test_direct_call_raises_on_backward_survey( self, @@ -11325,7 +11983,9 @@ def test_compute_heterogeneity_test_direct_call_raises_on_backward_survey( df["stratum"] = df["group"] % 4 df["psu_id"] = df["group"] sd = SurveyDesign( - weights="sw", strata="stratum", psu="psu_id", + weights="sw", + strata="stratum", + psu="psu_id", ) # Build a minimal valid obs_survey_info dict matching the # function's contract. SurveyDesign.resolve() takes only the @@ -11376,25 +12036,29 @@ def test_predict_het_placebo_survey_forward_only_still_works(self): df["stratum"] = df["group"] % 4 df["psu_id"] = df["group"] sd = SurveyDesign( - weights="sw", strata="stratum", psu="psu_id", - ) - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=False + weights="sw", + strata="stratum", + psu="psu_id", ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=False) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, - heterogeneity="het_x", survey_design=sd, + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", + survey_design=sd, ) assert res.path_heterogeneity_effects is not None for path, horizons in res.path_heterogeneity_effects.items(): for h in [1, 2, 3]: if h in horizons: assert np.isfinite(horizons[h]["beta"]), ( - f"path={path} h={h} beta non-finite under " - f"forward+survey path" + f"path={path} h={h} beta non-finite under " f"forward+survey path" ) def test_predict_het_placebo_eligible_filter(self): @@ -11422,30 +12086,46 @@ def test_predict_het_placebo_eligible_filter(self): else: d = 0 y = 0.5 * t + effect * d + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": d, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": d, + "outcome": y, + "het_x": het_x, + } + ) for k in range(n_controls): het_x = 1 if k < n_controls // 2 else 0 g = n_switchers + k for t in range(n_periods): y = 0.5 * t + rng.normal(0, 0.5) - rows.append({ - "group": g, "period": t, "treatment": 0, - "outcome": y, "het_x": het_x, - }) + rows.append( + { + "group": g, + "period": t, + "treatment": 0, + "outcome": y, + "het_x": het_x, + } + ) df = pd.DataFrame(rows) est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, paths_of_interest=[(0, 1, 1, 1)], + drop_larger_lower=False, + paths_of_interest=[(0, 1, 1, 1)], placebo=True, ) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) assert res.path_heterogeneity_effects is not None @@ -11455,8 +12135,7 @@ def test_predict_het_placebo_eligible_filter(self): # -3 has out_idx=-2 → all groups filtered, n_obs=0, NaN-consistent if -2 in path_het: assert path_het[-2]["n_obs"] == 0, ( - f"placebo -2 should be filtered (out_idx<0): " - f"got n_obs={path_het[-2]['n_obs']}" + f"placebo -2 should be filtered (out_idx<0): " f"got n_obs={path_het[-2]['n_obs']}" ) assert np.isnan(path_het[-2]["beta"]) assert np.isnan(path_het[-2]["se"]) @@ -11476,14 +12155,17 @@ def test_path_heterogeneity_telescopes_to_global_on_single_path_panel( the global regression. """ df = _single_path_het_data() - est_g = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, placebo=True - ) + est_g = ChaisemartinDHaultfoeuille(drop_larger_lower=False, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_g = est_g.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) est_p = ChaisemartinDHaultfoeuille( drop_larger_lower=False, @@ -11493,8 +12175,13 @@ def test_path_heterogeneity_telescopes_to_global_on_single_path_panel( with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res_p = est_p.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) path_het = res_p.path_heterogeneity_effects[(0, 1, 1, 1)] global_het = res_g.heterogeneity_effects @@ -11507,11 +12194,17 @@ def test_path_heterogeneity_telescopes_to_global_on_single_path_panel( assert not np.isfinite(p_h["beta"]) continue np.testing.assert_allclose( - p_h["beta"], g_h["beta"], atol=1e-14, rtol=1e-14, + p_h["beta"], + g_h["beta"], + atol=1e-14, + rtol=1e-14, err_msg=f"horizon {h} beta telescope failed", ) np.testing.assert_allclose( - p_h["se"], g_h["se"], atol=1e-14, rtol=1e-14, + p_h["se"], + g_h["se"], + atol=1e-14, + rtol=1e-14, err_msg=f"horizon {h} se telescope failed", ) assert int(p_h["n_obs"]) == int(g_h["n_obs"]) @@ -11519,14 +12212,17 @@ def test_path_heterogeneity_telescopes_to_global_on_single_path_panel( def test_summary_renders_placebo_het_rows(self): """`result.summary()` renders without error after #422.""" df = _by_path_het_data() - est = ChaisemartinDHaultfoeuille( - drop_larger_lower=False, by_path=3, placebo=True - ) + est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) res = est.fit( - df, outcome="outcome", group="group", time="period", - treatment="treatment", L_max=3, heterogeneity="het_x", + df, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + heterogeneity="het_x", ) s = res.summary() assert isinstance(s, str) @@ -11611,35 +12307,39 @@ def test_heterogeneity_df_uses_post_drop_rank(self): assert n_obs == 5, f"expected n_obs=5, got {n_obs}" # df = n_obs - rank = 5 - 4 = 1. safe_inference at df=1 # reproduces stored t/p/CI bit-exactly. - expected_t, expected_p, expected_ci = safe_inference( - h["beta"], h["se"], df=1 - ) + expected_t, expected_p, expected_ci = safe_inference(h["beta"], h["se"], df=1) np.testing.assert_allclose( - h["t_stat"], expected_t, atol=1e-12, rtol=1e-12, + h["t_stat"], + expected_t, + atol=1e-12, + rtol=1e-12, err_msg="t_stat does not match safe_inference(df=1)", ) np.testing.assert_allclose( - h["p_value"], expected_p, atol=1e-12, rtol=1e-12, + h["p_value"], + expected_p, + atol=1e-12, + rtol=1e-12, err_msg="p_value does not match safe_inference(df=1)", ) np.testing.assert_allclose( - h["conf_int"], expected_ci, atol=1e-12, rtol=1e-12, + h["conf_int"], + expected_ci, + atol=1e-12, + rtol=1e-12, err_msg="conf_int does not match safe_inference(df=1)", ) # safe_inference(df=n_obs - n_params=0) would produce different # p_value/conf_int. Pin the asymmetry so a regression that # reverts to pre-drop n_params is caught here. - wrong_t, wrong_p, wrong_ci = safe_inference( - h["beta"], h["se"], df=n_obs - 5 - ) + wrong_t, wrong_p, wrong_ci = safe_inference(h["beta"], h["se"], df=n_obs - 5) if np.isfinite(wrong_p): # When df=0, safe_inference NaN-fills; the asymmetry check # only fires when wrong_p is finite (which it isn't at df=0). # We still pin that the stored p_value is NOT equal to the # pre-drop result. assert not np.isclose(h["p_value"], wrong_p, atol=1e-10), ( - "stored p_value matches pre-drop n_params df; " - "rank-threading may have reverted" + "stored p_value matches pre-drop n_params df; " "rank-threading may have reverted" ) def test_heterogeneity_underidentified_nan_fills(self): @@ -11681,9 +12381,7 @@ def test_heterogeneity_underidentified_nan_fills(self): ) assert 1 in result h = result[1] - assert np.isnan(h["beta"]), ( - f"beta should be NaN when n_obs <= rank; got {h}" - ) + assert np.isnan(h["beta"]), f"beta should be NaN when n_obs <= rank; got {h}" assert np.isnan(h["se"]) assert np.isnan(h["t_stat"]) assert np.isnan(h["p_value"]) diff --git a/tests/test_doc_snippets.py b/tests/test_doc_snippets.py index a8e50a2fa..7e346f460 100644 --- a/tests/test_doc_snippets.py +++ b/tests/test_doc_snippets.py @@ -113,7 +113,9 @@ def _should_skip(code: str) -> Optional[str]: if re.search(pat, code, re.MULTILINE): return f"matches skip pattern: {pat}" # Skip if no actual Python statements (just comments / blank) - lines = [l.strip() for l in code.splitlines() if l.strip() and not l.strip().startswith("#")] + lines = [ + ln.strip() for ln in code.splitlines() if ln.strip() and not ln.strip().startswith("#") + ] if not lines: return "no executable statements" return None @@ -338,6 +340,7 @@ def _mock_list_datasets(): def _restore_datasets_module(): """Restore diff_diff.datasets after each test to prevent mock leaking.""" import sys as _sys + import diff_diff as _dd orig_mod = _sys.modules.get("diff_diff.datasets") diff --git a/tests/test_estimators.py b/tests/test_estimators.py index 2c61b8d0a..2494b6b5f 100644 --- a/tests/test_estimators.py +++ b/tests/test_estimators.py @@ -259,7 +259,7 @@ def test_rank_deficient_action_warn_default(self, simple_2x2_data): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - results = did.fit( + did.fit( data, outcome="outcome", treatment="treated", @@ -1043,6 +1043,7 @@ def test_rank_deficient_action_error_raises(self, twfe_panel_data): def test_rank_deficient_action_silent_no_warning(self, twfe_panel_data): """Test that rank_deficient_action='silent' produces no warning.""" import warnings + from diff_diff.estimators import TwoWayFixedEffects # Add a covariate that is perfectly collinear with another @@ -3976,7 +3977,6 @@ def test_sun_abraham_joint_span_covariate_dropped_att_stable(self): from diff_diff import SunAbraham df = self._frame(seed=8) - rng = np.random.default_rng(9) first = np.where(np.arange(df["unit"].nunique()) % 3 == 0, 0, 15)[ df["unit"].values % df["unit"].nunique() ] diff --git a/tests/test_imputation.py b/tests/test_imputation.py index d2288d7ca..1a2e868c5 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -674,7 +674,6 @@ def test_pretrend_with_violation(self): data = generate_test_data(seed=88, n_units=200) # Add a pre-treatment trend for treated units - rng = np.random.default_rng(88) for idx in data.index: if data.loc[idx, "first_treat"] > 0: t = data.loc[idx, "time"] @@ -1415,7 +1414,7 @@ def test_no_never_treated(self): data = generate_test_data(never_treated_frac=0.0, seed=42) est = ImputationDiD() - with warnings.catch_warnings(record=True) as w: + with warnings.catch_warnings(record=True): warnings.simplefilter("always") results = est.fit( data, @@ -1527,38 +1526,10 @@ def test_always_treated_warning(self): n_periods = 6 units = np.repeat(np.arange(n_units), n_periods) - times = np.tile(np.arange(n_periods), n_units) - - first_treat = np.zeros(n_units, dtype=int) - first_treat[10:20] = 0 # period 0 = always treated - first_treat[20:] = 3 # treated at period 3 - - # Make some units treated in all periods - first_treat[10:20] = 0 # Never treated (actually) - # To make always-treated: first_treat <= min_time (0) - first_treat[0:5] = 0 # These are never-treated - first_treat[5:10] = -1 # Treated before panel starts! - - first_treat_exp = np.repeat(first_treat, n_periods) - post = (times >= first_treat_exp) & (first_treat_exp > 0) & (first_treat_exp != np.inf) - - outcomes = ( - np.repeat(rng.standard_normal(n_units) * 2, n_periods) - + 2.0 * post - + rng.standard_normal(len(units)) * 0.5 - ) - # Fix: first_treat with -1 won't trigger the never_treated check properly - # Let's use first_treat = 0 for some units to trigger always-treated - first_treat_2 = np.zeros(n_units, dtype=int) - first_treat_2[:10] = 0 # never treated - first_treat_2[10:15] = 0 # also never treated (we need >= 1 always-treated) - first_treat_2[15:] = 3 - # Actually, to trigger always-treated, we need first_treat <= min(time) = 0 - # But first_treat == 0 means never-treated in the code - # We need first_treat > 0 but <= min(time) - # min(time) = 0, so first_treat must be <= 0 and > 0, impossible - # Let's start times at 1 + # To trigger the always-treated check we need first_treat > 0 but + # <= min(time). first_treat == 0 means never-treated in the code, so + # with times starting at 0 that is impossible — start times at 1. times_shifted = np.tile(np.arange(1, n_periods + 1), n_units) first_treat_3 = np.zeros(n_units, dtype=int) @@ -1587,7 +1558,7 @@ def test_always_treated_warning(self): est = ImputationDiD() with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - results = est.fit( + est.fit( data, outcome="outcome", unit="unit", @@ -3144,6 +3115,8 @@ def test_identified_leads_unchanged_no_snap_warning(self): ] assert len(finite_leads) >= 4 assert not any("collinear with the absorbed fixed effects" in str(x.message) for x in w) + + class TestLSMRFallbackParity: """The sparse LSMR fallback replaces dense lstsq on the (possibly singular) normal equations. Solver choice cannot change the estimator: diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 68a16ef5c..8cf026ba4 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -263,7 +263,7 @@ def test_rank_deficient_produces_nan_for_dropped_columns(self): assert np.all(np.isfinite(vcov_kept)), "VCoV for kept coefficients should be finite" # Residuals should be finite (computed using only identified coefficients) - assert np.all(np.isfinite(resid)), f"Residuals contain non-finite values" + assert np.all(np.isfinite(resid)), "Residuals contain non-finite values" def test_rank_deficient_error_mode(self): """Test that rank_deficient_action='error' raises ValueError.""" @@ -319,8 +319,8 @@ def test_skip_rank_check_bypasses_qr_decomposition(self): When skip_rank_check=True, the function should skip QR decomposition and go directly to SVD solving, even in Python backend. """ - import warnings import os + import warnings np.random.seed(42) n = 100 @@ -408,7 +408,7 @@ def test_multiperiod_like_design_full_rank(self): # If no rank deficiency, all coefficients should be finite if len(rank_warnings) == 0: - assert np.all(np.isfinite(coef)), f"Full-rank matrix: coefficients should be finite" + assert np.all(np.isfinite(coef)), "Full-rank matrix: coefficients should be finite" assert np.all(np.abs(coef) < 1e6), f"Coefficients are unreasonably large: {coef}" # The treatment effect coefficient (last one) should be close to true effect assert ( @@ -417,7 +417,7 @@ def test_multiperiod_like_design_full_rank(self): else: # If rank-deficient, check that identified coefficients are valid finite_coef = coef[~np.isnan(coef)] - assert np.all(np.isfinite(finite_coef)), f"Identified coefficients should be finite" + assert np.all(np.isfinite(finite_coef)), "Identified coefficients should be finite" # If treatment effect is identified, check it if not np.isnan(coef[-1]): assert ( @@ -489,7 +489,6 @@ def ols_data(self): np.random.seed(42) n = 200 X = np.column_stack([np.ones(n), np.random.randn(n)]) - beta = np.array([1.0, 2.0]) residuals = np.random.randn(n) return X, residuals @@ -573,8 +572,8 @@ def test_cluster_count_check_normalizes_series_cluster_ids(self): def test_numerical_instability_fallback_warns(self, ols_data): """Test that numerical instability in Rust backend triggers warning and fallback.""" - from unittest.mock import patch import warnings + from unittest.mock import patch from diff_diff import HAS_RUST_BACKEND @@ -1192,6 +1191,7 @@ def test_rank_deficient_degrees_of_freedom(self): def test_rank_deficient_inference_uses_correct_df(self): """Test that p-values and CIs use the correct df for rank-deficient matrices.""" import warnings + from scipy import stats np.random.seed(42) @@ -1528,7 +1528,6 @@ def test_did_finite_att_with_large_scale_covariate(self): def test_solve_ols_rank_zero_returns_nan_not_indexerror(self): """A design that collapses to rank 0 returns all-NaN coefficients with a warning, not a cryptic IndexError (empty float index array).""" - import warnings n = 50 X = np.zeros((n, 3)) # rank 0 @@ -2296,10 +2295,10 @@ class TestCheckPropensityDiagnostics: def test_no_warning_normal_scores(self): """No warning when all scores are within bounds.""" - from diff_diff.linalg import _check_propensity_diagnostics - import warnings + from diff_diff.linalg import _check_propensity_diagnostics + pscore = np.array([0.3, 0.5, 0.7, 0.4, 0.6]) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") diff --git a/tests/test_methodology_sdid.py b/tests/test_methodology_sdid.py index 26cdaa8fc..99133867d 100644 --- a/tests/test_methodology_sdid.py +++ b/tests/test_methodology_sdid.py @@ -7,12 +7,12 @@ """ import warnings +from typing import Any, Dict from unittest.mock import patch import numpy as np -import pytest - import pandas as pd +import pytest from diff_diff.synthetic_did import SyntheticDiD from diff_diff.utils import ( @@ -30,14 +30,12 @@ safe_inference, ) - # ============================================================================= # Test Helpers # ============================================================================= -def _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, - att=5.0, seed=42): +def _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, att=5.0, seed=42): """Create a simple panel dataset for testing.""" rng = np.random.default_rng(seed) data = [] @@ -48,13 +46,16 @@ def _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, y = 10.0 + unit_fe + t * 0.3 + rng.normal(0, 0.5) if is_treated and t >= n_pre: y += att - data.append({ - "unit": unit, - "period": t, - "treated": int(is_treated), - "outcome": y, - }) + data.append( + { + "unit": unit, + "period": t, + "treated": int(is_treated), + "outcome": y, + } + ) import pandas as pd + return pd.DataFrame(data) @@ -72,9 +73,7 @@ def test_known_values(self): # Unit 0: [1, 3, 6] -> diffs: [2, 3] # Unit 1: [2, 2, 5] -> diffs: [0, 3] # All diffs: [2, 3, 0, 3], sd(ddof=1) = std([2,3,0,3], ddof=1) - Y = np.array([[1.0, 2.0], - [3.0, 2.0], - [6.0, 5.0]]) + Y = np.array([[1.0, 2.0], [3.0, 2.0], [6.0, 5.0]]) expected = np.std([2.0, 3.0, 0.0, 3.0], ddof=1) result = _compute_noise_level(Y) assert abs(result - expected) < 1e-10 @@ -86,8 +85,7 @@ def test_single_period(self): def test_two_periods(self): """Two periods -> one diff per unit.""" - Y = np.array([[1.0, 4.0], - [3.0, 7.0]]) + Y = np.array([[1.0, 4.0], [3.0, 7.0]]) # Diffs: [2.0, 3.0], sd(ddof=1) expected = np.std([2.0, 3.0], ddof=1) assert abs(_compute_noise_level(Y) - expected) < 1e-10 @@ -99,9 +97,7 @@ class TestRegularization: def test_formula(self): """Check zeta_omega = (N1*T1)^0.25 * sigma, zeta_lambda = 1e-6 * sigma.""" # Use a simple Y_pre_control where sigma is easy to compute - Y = np.array([[1.0, 2.0], - [3.0, 4.0], - [6.0, 7.0]]) + Y = np.array([[1.0, 2.0], [3.0, 4.0], [6.0, 7.0]]) sigma = _compute_noise_level(Y) n_treated, n_post = 2, 3 @@ -115,9 +111,7 @@ def test_formula(self): def test_zero_noise(self): """Constant outcomes -> zero noise -> zero regularization.""" - Y = np.array([[5.0, 5.0], - [5.0, 5.0], - [5.0, 5.0]]) + Y = np.array([[5.0, 5.0], [5.0, 5.0], [5.0, 5.0]]) zo, zl = _compute_regularization(Y, 2, 3) assert zo == 0.0 assert zl == 0.0 @@ -142,7 +136,7 @@ def test_fw_step_descent(self): N, T0 = 10, 5 A = rng.standard_normal((N, T0)) b = rng.standard_normal(N) - eta = N * 0.1 ** 2 # eta = N * zeta^2 matching R's formulation + eta = N * 0.1**2 # eta = N * zeta^2 matching R's formulation x = np.ones(T0) / T0 @@ -309,7 +303,7 @@ def test_zero_sum(self): """Zero-sum vector -> uniform weights.""" v = np.array([0.0, 0.0, 0.0]) result = _sum_normalize(v) - np.testing.assert_allclose(result, [1.0/3, 1.0/3, 1.0/3]) + np.testing.assert_allclose(result, [1.0 / 3, 1.0 / 3, 1.0 / 3]) # ============================================================================= @@ -414,13 +408,15 @@ class TestATTFullPipeline: def test_estimation_produces_reasonable_att(self, ci_params): """Full estimation on canonical data should produce reasonable ATT.""" - df = _make_panel(n_control=20, n_treated=3, n_pre=6, n_post=3, - att=5.0, seed=42) + df = _make_panel(n_control=20, n_treated=3, n_pre=6, n_post=3, att=5.0, seed=42) n_boot = ci_params.bootstrap(50) sdid = SyntheticDiD(n_bootstrap=n_boot, seed=42, variance_method="placebo") results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(6, 9)), ) @@ -433,8 +429,11 @@ def test_results_have_regularization_info(self): df = _make_panel(seed=42) sdid = SyntheticDiD(variance_method="placebo", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) @@ -448,13 +447,13 @@ def test_results_have_regularization_info(self): def test_user_override_regularization(self): """User-specified zeta_omega/zeta_lambda should be used instead of auto.""" df = _make_panel(seed=42) - sdid = SyntheticDiD( - zeta_omega=99.0, zeta_lambda=0.5, - variance_method="placebo", seed=42 - ) + sdid = SyntheticDiD(zeta_omega=99.0, zeta_lambda=0.5, variance_method="placebo", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) @@ -473,12 +472,13 @@ class TestPlaceboSE: def test_placebo_se_formula(self): """SE should be sqrt((r-1)/r) * sd(estimates, ddof=1).""" df = _make_panel(n_control=15, n_treated=2, seed=42) - sdid = SyntheticDiD( - variance_method="placebo", n_bootstrap=100, seed=42 - ) + sdid = SyntheticDiD(variance_method="placebo", n_bootstrap=100, seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) @@ -510,12 +510,13 @@ def test_bootstrap_se_positive(self, ci_params): """Bootstrap SE should be positive.""" df = _make_panel(n_control=20, n_treated=3, seed=42) n_boot = ci_params.bootstrap(50) - sdid = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=n_boot, seed=42 - ) + sdid = SyntheticDiD(variance_method="bootstrap", n_bootstrap=n_boot, seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) @@ -528,12 +529,18 @@ def test_bootstrap_with_zeta_overrides(self, ci_params): df = _make_panel(n_control=20, n_treated=3, seed=42) n_boot = ci_params.bootstrap(50) sdid = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=n_boot, - zeta_omega=99.0, zeta_lambda=0.5, seed=42, + variance_method="bootstrap", + n_bootstrap=n_boot, + zeta_omega=99.0, + zeta_lambda=0.5, + seed=42, ) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) @@ -574,17 +581,21 @@ def test_bootstrap_se_tracks_placebo_se_exchangeable(self, ci_params): n_boot = ci_params.bootstrap(200, min_n=100) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - r_placebo = SyntheticDiD( - variance_method="placebo", n_bootstrap=n_boot, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", post_periods=[5, 6, 7], + r_placebo = SyntheticDiD(variance_method="placebo", n_bootstrap=n_boot, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=[5, 6, 7], ) - r_boot = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=n_boot, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", post_periods=[5, 6, 7], + r_boot = SyntheticDiD(variance_method="bootstrap", n_bootstrap=n_boot, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=[5, 6, 7], ) rel_diff = abs(r_boot.se - r_placebo.se) / r_placebo.se # Tolerance chosen for B=100–200 with MC noise floor ~7–10% on the @@ -603,13 +614,15 @@ def test_bootstrap_succeeds_on_pweight_survey(self): finite, SE is positive, variance_method is preserved. """ from diff_diff.survey import SurveyDesign + df = _make_panel(n_control=20, n_treated=3, seed=42) df["wt"] = 1.0 - result = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + result = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt"), ) @@ -627,17 +640,19 @@ def test_bootstrap_succeeds_on_full_design_survey(self): records the strata/PSU layout. """ from diff_diff.survey import SurveyDesign + df = _make_panel(n_control=20, n_treated=3, seed=42) df["wt"] = 1.0 df["stratum"] = df["unit"] % 2 # Globally unique PSU labels (avoids needing nest=True). Each unit # gets its own PSU id; stratum partitions the PSUs into two groups. df["psu"] = df["unit"] - result = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + result = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", strata="stratum", psu="psu"), ) @@ -660,11 +675,12 @@ def test_bootstrap_summary_shows_replications(self, ci_params): n_boot = ci_params.bootstrap(50) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - r = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=n_boot, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + r = SyntheticDiD(variance_method="bootstrap", n_bootstrap=n_boot, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) summary = r.summary() @@ -690,17 +706,19 @@ def test_bootstrap_pweight_only_retries_zero_treated_mass_draws(self): zero-mass condition are retried rather than getting a wrong τ. """ from diff_diff.survey import SurveyDesign + df = _make_panel(n_control=25, n_treated=3, seed=42) # Give one treated unit weight 0 and the other two weight 1. # Control units get unit weights (any positive). Treated_idx 25, # 26, 27; set 25 → weight 0, 26 and 27 → weight 1. df["wt"] = 1.0 df.loc[df["unit"] == 25, "wt"] = 0.0 - result = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=100, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + result = SyntheticDiD(variance_method="bootstrap", n_bootstrap=100, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt"), ) @@ -721,14 +739,16 @@ def test_bootstrap_full_design_without_explicit_weights(self): could run. The helper now returns ones for this configuration. """ from diff_diff.survey import SurveyDesign + df = _make_panel(n_control=20, n_treated=3, seed=42) df["stratum"] = df["unit"] % 2 df["psu"] = df["unit"] - result = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + result = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(strata="stratum", psu="psu"), # weights=None ) @@ -758,8 +778,8 @@ def test_bootstrap_full_design_rao_wu_boot_idx_slice(self, monkeypatch): boot_idx slicing, or the rw-then-slice order gets swapped to slice-then-rw). Both would silently produce wrong bootstrap SE. """ - from diff_diff import utils as dd_utils from diff_diff import synthetic_did as sdid_mod + from diff_diff import utils as dd_utils from diff_diff.survey import SurveyDesign df = _make_panel(n_control=15, n_treated=3, seed=42) @@ -786,16 +806,19 @@ def capturing_helper(Y_pre_c, Y_pre_t_mean, rw, *args, **kwargs): captured.append(np.array(rw, copy=True)) return real_helper(Y_pre_c, Y_pre_t_mean, rw, *args, **kwargs) - monkeypatch.setattr( - sdid_mod, "compute_sdid_unit_weights_survey", capturing_helper - ) + monkeypatch.setattr(sdid_mod, "compute_sdid_unit_weights_survey", capturing_helper) bootstrap_seed = 1 SyntheticDiD( - variance_method="bootstrap", n_bootstrap=10, seed=bootstrap_seed, + variance_method="bootstrap", + n_bootstrap=10, + seed=bootstrap_seed, ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", strata="stratum", psu="psu"), ) @@ -828,9 +851,7 @@ def capturing_helper(Y_pre_c, Y_pre_t_mean, rw, *args, **kwargs): expected_slices.append(known_rw[:n_control][boot_idx[boot_is_control]]) assert len(captured) >= 1, "no FW calls captured — survey dispatch broken" - for i, (rw_captured, rw_expected) in enumerate( - zip(captured, expected_slices) - ): + for i, (rw_captured, rw_expected) in enumerate(zip(captured, expected_slices)): np.testing.assert_array_equal( rw_captured, rw_expected, @@ -862,8 +883,11 @@ def test_fit_raises_on_zero_total_treated_survey_mass(self): df["wt"] = np.where(df["treated"] == 1, 0.0, 1.0) with pytest.raises(ValueError, match=r"treated arm has zero total mass"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt"), ) @@ -881,8 +905,11 @@ def test_fit_raises_on_zero_total_control_survey_mass(self): df["wt"] = np.where(df["treated"] == 0, 0.0, 1.0) with pytest.raises(ValueError, match=r"control arm has zero total mass"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt"), ) @@ -903,8 +930,11 @@ def test_fit_raises_on_zero_treated_mass_under_full_design(self): df["psu"] = df["unit"] with pytest.raises(ValueError, match=r"treated arm has zero total mass"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", strata="stratum", psu="psu"), ) @@ -954,14 +984,15 @@ def sparse_unit_weights(Y_pre_control, *args, **kwargs): w[0] = 1.0 return w - monkeypatch.setattr( - sdid_mod, "compute_sdid_unit_weights", sparse_unit_weights - ) + monkeypatch.setattr(sdid_mod, "compute_sdid_unit_weights", sparse_unit_weights) with pytest.raises(ValueError, match=r"unidentified|zero survey weight"): SyntheticDiD(variance_method="placebo", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt"), ) @@ -992,14 +1023,15 @@ def sparse_unit_weights(Y_pre_control, *args, **kwargs): w[0] = 1.0 return w - monkeypatch.setattr( - sdid_mod, "compute_sdid_unit_weights", sparse_unit_weights - ) + monkeypatch.setattr(sdid_mod, "compute_sdid_unit_weights", sparse_unit_weights) with pytest.raises(ValueError, match=r"unidentified|zero survey weight"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", strata="stratum", psu="psu"), ) @@ -1026,8 +1058,11 @@ def test_fit_raises_on_implicit_psu_fpc_below_unit_count_unstratified(self): df["fpc_pop"] = 5.0 with pytest.raises(ValueError, match=r"FPC.*less than the number of units"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", fpc="fpc_pop"), ) @@ -1047,11 +1082,16 @@ def test_fit_raises_on_implicit_psu_fpc_below_stratum_unit_count(self): df["fpc_pop"] = 3.0 with pytest.raises(ValueError, match=r"FPC.*less than the number of units in that stratum"): SyntheticDiD(variance_method="bootstrap", n_bootstrap=20, seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign( - weights="wt", strata="stratum", fpc="fpc_pop", + weights="wt", + strata="stratum", + fpc="fpc_pop", ), ) @@ -1092,28 +1132,43 @@ def test_bootstrap_scale_invariance_under_pweight_rescaling(self): df["wt_scaled"] = df["wt"] * 10.0 kwargs = dict( - outcome="outcome", treatment="treated", - unit="unit", time="period", + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) result_base = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=1, + variance_method="bootstrap", + n_bootstrap=50, + seed=1, ).fit(df, survey_design=SurveyDesign(weights="wt"), **kwargs) result_scaled = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=1, + variance_method="bootstrap", + n_bootstrap=50, + seed=1, ).fit(df, survey_design=SurveyDesign(weights="wt_scaled"), **kwargs) assert np.isfinite(result_base.se) and result_base.se > 0 np.testing.assert_allclose( - result_scaled.se, result_base.se, rtol=1e-13, atol=0, + result_scaled.se, + result_base.se, + rtol=1e-13, + atol=0, err_msg="bootstrap SE is not invariant to pweight global rescaling", ) np.testing.assert_allclose( - result_scaled.p_value, result_base.p_value, rtol=1e-12, atol=1e-14, + result_scaled.p_value, + result_base.p_value, + rtol=1e-12, + atol=1e-14, err_msg="bootstrap p-value is not invariant to pweight global rescaling", ) np.testing.assert_allclose( - result_scaled.conf_int, result_base.conf_int, rtol=1e-13, atol=0, + result_scaled.conf_int, + result_base.conf_int, + rtol=1e-13, + atol=0, err_msg="bootstrap CI is not invariant to pweight global rescaling", ) @@ -1127,18 +1182,23 @@ def test_bootstrap_single_psu_returns_nan(self): #351 fixed-weight Rao-Wu branch (commit 91082e5). """ from diff_diff.survey import SurveyDesign + df = _make_panel(n_control=20, n_treated=3, seed=42) df["wt"] = 1.0 df["psu_single"] = 0 # all units in the same PSU; no strata sdid = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=1) result = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], survey_design=SurveyDesign(weights="wt", psu="psu_single"), ) - assert np.isnan(result.se), \ - f"single-PSU unstratified bootstrap should return NaN SE, got {result.se}" + assert np.isnan( + result.se + ), f"single-PSU unstratified bootstrap should return NaN SE, got {result.se}" def test_bootstrap_fw_nonconvergence_warning_fires_under_rust(self, monkeypatch): """Aggregate FW non-convergence warning surfaces on the Rust backend. @@ -1169,9 +1229,7 @@ def _always_not_converged(Y, zeta, intercept, init, min_decrease, max_iter): weights, _ = real_rust(Y, zeta, intercept, init, min_decrease, max_iter) return weights, False - monkeypatch.setattr( - dd_utils, "_rust_sc_weight_fw_with_convergence", _always_not_converged - ) + monkeypatch.setattr(dd_utils, "_rust_sc_weight_fw_with_convergence", _always_not_converged) df = _make_panel(n_control=20, n_treated=3, seed=42) sdid = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=42) @@ -1179,13 +1237,17 @@ def _always_not_converged(Y, zeta, intercept, init, min_decrease, max_iter): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) fw_warnings = [ - w for w in caught + w + for w in caught if issubclass(w.category, UserWarning) and "Frank-Wolfe did not converge" in str(w.message) ] @@ -1211,8 +1273,11 @@ def test_jackknife_se_positive(self): df = _make_panel(n_control=20, n_treated=3, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) assert results.se > 0 @@ -1222,13 +1287,19 @@ def test_jackknife_deterministic(self): """Jackknife should produce identical results regardless of seed.""" df = _make_panel(n_control=15, n_treated=3, seed=42) results1 = SyntheticDiD(variance_method="jackknife", seed=1).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) results2 = SyntheticDiD(variance_method="jackknife", seed=999).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) # ATT should be identical (same weights, no randomness in point est) @@ -1241,8 +1312,11 @@ def test_jackknife_se_formula(self): df = _make_panel(n_control=15, n_treated=3, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) assert results.variance_effects is not None @@ -1258,8 +1332,11 @@ def test_jackknife_n_iterations(self): df = _make_panel(n_control=n_co, n_treated=n_tr, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) assert results.variance_effects is not None @@ -1272,8 +1349,11 @@ def test_jackknife_single_treated_nan(self): warnings.simplefilter("ignore") sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 7)), ) assert np.isnan(results.se) @@ -1287,8 +1367,11 @@ def test_jackknife_analytical_pvalue(self): df = _make_panel(n_control=20, n_treated=3, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) if np.isfinite(results.t_stat): @@ -1299,13 +1382,19 @@ def test_jackknife_same_att_as_placebo(self): """Jackknife should produce the same point estimate as placebo.""" df = _make_panel(n_control=15, n_treated=3, seed=42) res_jk = SyntheticDiD(variance_method="jackknife", seed=42).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) res_pl = SyntheticDiD(variance_method="placebo", seed=42, n_bootstrap=50).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) assert abs(res_jk.att - res_pl.att) < 1e-10 @@ -1321,8 +1410,11 @@ def test_jackknife_n_bootstrap_none_in_results(self): df = _make_panel(n_control=15, n_treated=3, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), ) assert results.n_bootstrap is None @@ -1338,8 +1430,11 @@ def test_jackknife_with_pweights(self): sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 8)), survey_design=SurveyDesign(weights="weight"), ) @@ -1368,8 +1463,11 @@ def test_jackknife_zero_effective_control_nan(self): warnings.simplefilter("ignore") sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 7)), survey_design=SurveyDesign(weights="weight"), ) @@ -1394,8 +1492,11 @@ def test_jackknife_zero_treated_weight_nan(self): warnings.simplefilter("ignore") sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(5, 7)), survey_design=SurveyDesign(weights="weight"), ) @@ -1430,67 +1531,189 @@ class TestJackknifeSERParity: # R's Y matrix (23 units x 8 periods), row-major Y_FLAT = [ - 12.459567808595292, 13.223481099962006, 13.658348196773856, - 13.844051055863837, 13.888854636247594, 14.997677893012806, - 14.494587375086788, 15.851128751231856, 10.527006629006900, - 11.317894498245712, 9.780141451338988, 10.635177418486473, - 11.007911133698329, 11.692547000930196, 11.532445341187122, - 10.646344091442769, 5.779122815714058, 5.265746845809725, - 4.828411925858962, 5.933107464969151, 6.926403492435262, - 7.566662873481445, 6.703831577045862, 7.090431451464497, - 6.703722507026075, 6.453676391630379, 7.301398891231049, - 7.726092498224848, 8.191225590595401, 7.669210641906834, - 8.526151391259425, 7.715169490073769, 8.005628186152748, - 7.523978158267692, 9.049143286687135, 9.434081283341134, - 9.450553333966674, 10.310163601090766, 9.867729569702721, - 9.846941461031360, 10.459939463684098, 11.887686682638062, - 11.249912950470762, 12.093459993478538, 12.226598684379407, - 11.973716581337246, 13.453499811673423, 13.287085704636093, - 10.317796666844943, 10.819165701226847, 10.824437736488752, - 9.582976251622744, 11.521962769964540, 11.495903971828724, - 12.072136575632017, 12.570433156881965, 12.435827624848123, - 13.750744970607428, 13.567397714461393, 14.218726703934166, - 14.459837938730677, 14.659912736018788, 14.077914185301429, - 14.854380461280002, 10.770274645112915, 11.275621916712160, - 12.137534572839927, 12.531125692916383, 12.678920118269170, - 12.304148175294246, 12.497145874675160, 14.103389828901550, - 10.560062989643855, 10.755394606294518, 10.518678427483797, - 11.721841324084256, 11.607272952190801, 11.924464521898100, - 12.782516039349641, 13.026729430318186, 12.546145790341205, - 13.409407032231695, 14.079787980063543, 13.128838312144593, - 13.553836458429620, 13.718363411441658, 13.854625752117343, - 14.924224028489123, 11.906891367097627, 12.128784222882244, - 11.404804355878456, 13.130649630134753, 12.173021974919472, - 12.859165585526416, 12.895280738363951, 13.345233593320895, - 10.435966548001499, 10.663839793569295, 11.030422432974012, - 11.033668451079661, 11.324277503659044, 11.045836529045589, - 11.985219205566086, 12.220060940064094, 14.722723885094736, - 15.772410109968900, 15.256969467031452, 15.568564129971197, - 16.666133193788099, 16.405462433247578, 17.202870693537243, - 17.289652559976691, 7.760317864391456, 8.460282811921017, - 9.462415007659978, 9.956467084312777, 9.726218110324272, - 10.272688229133685, 11.134101608790994, 11.592584658589104, - 7.747112683063268, 8.706521663648207, 8.170907672905205, - 8.679537720718859, 8.962718814069811, 8.861932954235140, - 9.383430460745986, 9.891050023644237, 9.728955313568255, - 9.231765881057163, 9.555677785583788, 10.420693590160205, - 9.844078095298698, 10.651913064308546, 10.196489890710358, - 11.855847076501993, 9.218785934915712, 9.133582433258733, - 10.048827580363175, 9.952567508276010, 10.385962432276619, - 11.596546220044132, 11.164945662130776, 11.016817405176500, - 10.145044557120791, 10.921420538928436, 11.642624728800259, - 10.730067509380019, 11.753738913724906, 11.868862794274008, - 12.574196556067037, 12.311524695461632, 10.800710206252880, - 12.817967597577915, 12.705627126180516, 12.497850142478354, - 12.148734571851643, 13.494742486942219, 13.714835068828613, - 13.770060323710533, 10.010857300549947, 10.787315152039971, - 11.050238955584605, 11.063282099053561, 10.834793458278272, - 17.153286194944865, 17.380010096861866, 16.984758489324143, - 6.913302966281331, 6.938279687001069, 7.537129527669741, - 7.063822443245238, 7.531238453797332, 13.853711102827464, - 13.812711128345372, 14.204067444347162, 13.694867606609098, - 12.929992273442151, 14.397345491024691, 15.116119455987304, - 15.860226513457558, 19.442026093187646, 19.855029109494353, + 12.459567808595292, + 13.223481099962006, + 13.658348196773856, + 13.844051055863837, + 13.888854636247594, + 14.997677893012806, + 14.494587375086788, + 15.851128751231856, + 10.527006629006900, + 11.317894498245712, + 9.780141451338988, + 10.635177418486473, + 11.007911133698329, + 11.692547000930196, + 11.532445341187122, + 10.646344091442769, + 5.779122815714058, + 5.265746845809725, + 4.828411925858962, + 5.933107464969151, + 6.926403492435262, + 7.566662873481445, + 6.703831577045862, + 7.090431451464497, + 6.703722507026075, + 6.453676391630379, + 7.301398891231049, + 7.726092498224848, + 8.191225590595401, + 7.669210641906834, + 8.526151391259425, + 7.715169490073769, + 8.005628186152748, + 7.523978158267692, + 9.049143286687135, + 9.434081283341134, + 9.450553333966674, + 10.310163601090766, + 9.867729569702721, + 9.846941461031360, + 10.459939463684098, + 11.887686682638062, + 11.249912950470762, + 12.093459993478538, + 12.226598684379407, + 11.973716581337246, + 13.453499811673423, + 13.287085704636093, + 10.317796666844943, + 10.819165701226847, + 10.824437736488752, + 9.582976251622744, + 11.521962769964540, + 11.495903971828724, + 12.072136575632017, + 12.570433156881965, + 12.435827624848123, + 13.750744970607428, + 13.567397714461393, + 14.218726703934166, + 14.459837938730677, + 14.659912736018788, + 14.077914185301429, + 14.854380461280002, + 10.770274645112915, + 11.275621916712160, + 12.137534572839927, + 12.531125692916383, + 12.678920118269170, + 12.304148175294246, + 12.497145874675160, + 14.103389828901550, + 10.560062989643855, + 10.755394606294518, + 10.518678427483797, + 11.721841324084256, + 11.607272952190801, + 11.924464521898100, + 12.782516039349641, + 13.026729430318186, + 12.546145790341205, + 13.409407032231695, + 14.079787980063543, + 13.128838312144593, + 13.553836458429620, + 13.718363411441658, + 13.854625752117343, + 14.924224028489123, + 11.906891367097627, + 12.128784222882244, + 11.404804355878456, + 13.130649630134753, + 12.173021974919472, + 12.859165585526416, + 12.895280738363951, + 13.345233593320895, + 10.435966548001499, + 10.663839793569295, + 11.030422432974012, + 11.033668451079661, + 11.324277503659044, + 11.045836529045589, + 11.985219205566086, + 12.220060940064094, + 14.722723885094736, + 15.772410109968900, + 15.256969467031452, + 15.568564129971197, + 16.666133193788099, + 16.405462433247578, + 17.202870693537243, + 17.289652559976691, + 7.760317864391456, + 8.460282811921017, + 9.462415007659978, + 9.956467084312777, + 9.726218110324272, + 10.272688229133685, + 11.134101608790994, + 11.592584658589104, + 7.747112683063268, + 8.706521663648207, + 8.170907672905205, + 8.679537720718859, + 8.962718814069811, + 8.861932954235140, + 9.383430460745986, + 9.891050023644237, + 9.728955313568255, + 9.231765881057163, + 9.555677785583788, + 10.420693590160205, + 9.844078095298698, + 10.651913064308546, + 10.196489890710358, + 11.855847076501993, + 9.218785934915712, + 9.133582433258733, + 10.048827580363175, + 9.952567508276010, + 10.385962432276619, + 11.596546220044132, + 11.164945662130776, + 11.016817405176500, + 10.145044557120791, + 10.921420538928436, + 11.642624728800259, + 10.730067509380019, + 11.753738913724906, + 11.868862794274008, + 12.574196556067037, + 12.311524695461632, + 10.800710206252880, + 12.817967597577915, + 12.705627126180516, + 12.497850142478354, + 12.148734571851643, + 13.494742486942219, + 13.714835068828613, + 13.770060323710533, + 10.010857300549947, + 10.787315152039971, + 11.050238955584605, + 11.063282099053561, + 10.834793458278272, + 17.153286194944865, + 17.380010096861866, + 16.984758489324143, + 6.913302966281331, + 6.938279687001069, + 7.537129527669741, + 7.063822443245238, + 7.531238453797332, + 13.853711102827464, + 13.812711128345372, + 14.204067444347162, + 13.694867606609098, + 12.929992273442151, + 14.397345491024691, + 15.116119455987304, + 15.860226513457558, + 19.442026093187646, + 19.855029109494353, 20.377546194927845, ] N0, N1, T0, T1 = 20, 3, 5, 3 @@ -1506,20 +1729,25 @@ def r_panel_df(self): rows = [] for i in range(N): for t in range(T): - rows.append({ - "unit": i, - "time": t, - "outcome": Y[i, t], - "treated": int(i >= self.N0), - }) + rows.append( + { + "unit": i, + "time": t, + "outcome": Y[i, t], + "treated": int(i >= self.N0), + } + ) return pd.DataFrame(rows) def test_att_matches_r(self, r_panel_df): """ATT should match R's synthdid_estimate to machine precision.""" sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - r_panel_df, outcome="outcome", treatment="treated", - unit="unit", time="time", + r_panel_df, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", post_periods=[5, 6, 7], ) assert abs(results.att - self.R_ATT) < 1e-10 @@ -1528,8 +1756,11 @@ def test_jackknife_se_matches_r(self, r_panel_df): """Jackknife SE should match R's vcov(method='jackknife').""" sdid = SyntheticDiD(variance_method="jackknife", seed=42) results = sdid.fit( - r_panel_df, outcome="outcome", treatment="treated", - unit="unit", time="time", + r_panel_df, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", post_periods=[5, 6, 7], ) assert abs(results.se - self.R_JACKKNIFE_SE) < 1e-10 @@ -1563,10 +1794,7 @@ def test_placebo_se_matches_r(self, r_panel_df): import json import pathlib - fixture_path = ( - pathlib.Path(__file__).parent - / "data" / "sdid_placebo_indices_r.json" - ) + fixture_path = pathlib.Path(__file__).parent / "data" / "sdid_placebo_indices_r.json" if not fixture_path.exists(): pytest.skip( f"Missing R-parity fixture {fixture_path}; regenerate via " @@ -1594,8 +1822,11 @@ def capture_then_call(*args, **kwargs): sdid._placebo_variance_se = capture_then_call # type: ignore[assignment] sdid.fit( - r_panel_df, outcome="outcome", treatment="treated", - unit="unit", time="time", + r_panel_df, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", post_periods=[5, 6, 7], ) sdid._placebo_variance_se = original_method # type: ignore[assignment] @@ -1608,18 +1839,16 @@ def capture_then_call(*args, **kwargs): kwargs = dict(captured["kwargs"]) kwargs["replications"] = replications kwargs["_placebo_indices"] = r_perms - se_n, placebo_effects_n = sdid._placebo_variance_se( - *captured["args"], **kwargs - ) + se_n, placebo_effects_n = sdid._placebo_variance_se(*captured["args"], **kwargs) Y_scale = sdid.results_.zeta_omega / kwargs["zeta_omega"] py_se = se_n * Y_scale py_taus = np.asarray(placebo_effects_n) * Y_scale # Match R within cross-library FW tolerance (Rust vs R BLAS # reductions differ at sub-ULP; 1e-8 absorbs that without # masking a real divergence). - assert abs(py_se - r_se) < 1e-8, ( - f"Python placebo SE {py_se} != R {r_se} (delta {py_se - r_se})" - ) + assert ( + abs(py_se - r_se) < 1e-8 + ), f"Python placebo SE {py_se} != R {r_se} (delta {py_se - r_se})" # Per-draw τ regression: equal-SE doesn't imply equal sample, and # the placebo τ vector is user-visible through ``variance_effects`` # and feeds the empirical placebo p-value (synthetic_did.py @@ -1627,11 +1856,14 @@ def capture_then_call(*args, **kwargs): # diverged at a single draw — but happened to leave sd() unchanged # — still trips the regression. r_taus = np.asarray(payload["R_PLACEBO_TAUS"], dtype=float) - assert r_taus.shape == py_taus.shape == (replications,), ( - f"shape mismatch: r {r_taus.shape}, py {py_taus.shape}" - ) + assert ( + r_taus.shape == py_taus.shape == (replications,) + ), f"shape mismatch: r {r_taus.shape}, py {py_taus.shape}" np.testing.assert_allclose( - py_taus, r_taus, atol=1e-8, rtol=1e-8, + py_taus, + r_taus, + atol=1e-8, + rtol=1e-8, err_msg="Per-draw placebo τ diverges from R despite SE match", ) @@ -1656,13 +1888,15 @@ def test_n_bootstrap_validation(self): def test_single_treated_unit(self, ci_params): """Estimation should work with a single treated unit.""" - df = _make_panel(n_control=10, n_treated=1, n_pre=5, n_post=2, - att=3.0, seed=42) + df = _make_panel(n_control=10, n_treated=1, n_pre=5, n_post=2, att=3.0, seed=42) n_boot = ci_params.bootstrap(30) sdid = SyntheticDiD(n_bootstrap=n_boot, seed=42) results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6], ) assert np.isfinite(results.att) @@ -1675,8 +1909,11 @@ def test_insufficient_controls_for_placebo(self): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6], ) user_warnings = [x for x in w if issubclass(x.category, UserWarning)] @@ -1693,8 +1930,11 @@ def test_se_zero_propagation(self): with warnings.catch_warnings(record=True): warnings.simplefilter("always") results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6], ) @@ -1733,12 +1973,17 @@ def mock_estimator(*args, **kwargs): sdid = SyntheticDiD(n_bootstrap=20, seed=42) - with patch('diff_diff.synthetic_did.compute_sdid_estimator', - side_effect=mock_estimator), \ - patch('diff_diff._backend.HAS_RUST_BACKEND', False): + with ( + patch("diff_diff.synthetic_did.compute_sdid_estimator", side_effect=mock_estimator), + patch("diff_diff._backend.HAS_RUST_BACKEND", False), + ): se, estimates = sdid._bootstrap_se( - Y_pre_c, Y_post_c, Y_pre_t, Y_post_t, - unit_weights, time_weights, + Y_pre_c, + Y_post_c, + Y_pre_t, + Y_post_t, + unit_weights, + time_weights, ) # All retained estimates must be finite (non-finite never leaks). @@ -1747,9 +1992,9 @@ def mock_estimator(*args, **kwargs): assert len(estimates) == 20 # Retry fired: estimator was called more than B times because every # third call returned inf and triggered another attempt. - assert call_count[0] > 20, ( - f"expected retry path to fire (call_count > 20); got {call_count[0]}" - ) + assert ( + call_count[0] > 20 + ), f"expected retry path to fire (call_count > 20); got {call_count[0]}" def test_nonfinite_tau_filtered_in_placebo(self): """Non-finite tau values are filtered in Python placebo path (matches Rust).""" @@ -1770,15 +2015,22 @@ def mock_estimator(*args, **kwargs): sdid = SyntheticDiD(seed=42) - with patch('diff_diff.synthetic_did.compute_sdid_estimator', - side_effect=mock_estimator), \ - patch('diff_diff.synthetic_did.compute_sdid_unit_weights', - return_value=np.ones(n_control - n_treated) / (n_control - n_treated)), \ - patch('diff_diff.synthetic_did.compute_time_weights', - return_value=np.ones(n_pre) / n_pre), \ - patch('diff_diff._backend.HAS_RUST_BACKEND', False): + with ( + patch("diff_diff.synthetic_did.compute_sdid_estimator", side_effect=mock_estimator), + patch( + "diff_diff.synthetic_did.compute_sdid_unit_weights", + return_value=np.ones(n_control - n_treated) / (n_control - n_treated), + ), + patch( + "diff_diff.synthetic_did.compute_time_weights", return_value=np.ones(n_pre) / n_pre + ), + patch("diff_diff._backend.HAS_RUST_BACKEND", False), + ): se, estimates = sdid._placebo_variance_se( - Y_pre_c, Y_post_c, Y_pre_t_mean, Y_post_t_mean, + Y_pre_c, + Y_post_c, + Y_pre_t_mean, + Y_post_t_mean, n_treated=n_treated, replications=20, ) @@ -1794,12 +2046,16 @@ def test_inf_se_produces_nan_inference(self): sdid = SyntheticDiD(variance_method="bootstrap", n_bootstrap=10, seed=42) with patch.object( - SyntheticDiD, '_bootstrap_se', + SyntheticDiD, + "_bootstrap_se", return_value=(np.inf, np.array([1.0, 2.0, 3.0])), ): results = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6], ) @@ -2027,8 +2283,7 @@ def test_placebo_reestimates_weights_not_fixed(self): # Need enough controls for placebo to work (n_control > n_treated) # and enough variation for weights to differ between re-estimation # and renormalization. - df = _make_panel(n_control=15, n_treated=2, n_pre=6, n_post=3, - att=5.0, seed=123) + df = _make_panel(n_control=15, n_treated=2, n_pre=6, n_post=3, att=5.0, seed=123) post_periods = list(range(6, 9)) # Fit SDID to get original weights and matrices @@ -2063,8 +2318,9 @@ def test_placebo_reestimates_weights_not_fixed(self): # Extract numpy arrays from result dicts (ordered by control unit) unit_weights_arr = np.array([results.unit_weights[u] for u in control_units]) - time_weights_arr = np.array([results.time_weights[t] - for t in sorted(results.time_weights.keys())]) + time_weights_arr = np.array( + [results.time_weights[t] for t in sorted(results.time_weights.keys())] + ) n_control = len(control_idx) n_treated_count = len(treated_idx) @@ -2087,9 +2343,12 @@ def test_placebo_reestimates_weights_not_fixed(self): try: tau = compute_sdid_estimator( - Y_pre_pc, Y_post_pc, - Y_pre_pt_mean, Y_post_pt_mean, - fixed_omega, fixed_lambda, + Y_pre_pc, + Y_post_pc, + Y_pre_pt_mean, + Y_post_pt_mean, + fixed_omega, + fixed_lambda, ) fixed_estimates.append(tau) except (ValueError, np.linalg.LinAlgError): @@ -2097,8 +2356,7 @@ def test_placebo_reestimates_weights_not_fixed(self): if len(fixed_estimates) >= 2: n_s = len(fixed_estimates) - fixed_se = (np.sqrt((n_s - 1) / n_s) - * np.std(fixed_estimates, ddof=1)) + fixed_se = np.sqrt((n_s - 1) / n_s) * np.std(fixed_estimates, ddof=1) # The two SEs should differ because re-estimation produces # different weights than renormalization assert actual_se != pytest.approx(fixed_se, rel=0.01), ( @@ -2119,18 +2377,24 @@ class TestTreatmentValidation: def test_varying_treatment_within_unit_raises(self): """Unit whose treatment switches over time should raise ValueError.""" np.random.seed(42) - data = pd.DataFrame({ - "unit": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3], - "time": [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4], - "outcome": np.random.randn(12), - # Unit 1: treatment turns on at time 3 (staggered) - "treated": [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], - }) + data = pd.DataFrame( + { + "unit": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3], + "time": [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4], + "outcome": np.random.randn(12), + # Unit 1: treatment turns on at time 3 (staggered) + "treated": [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], + } + ) sdid = SyntheticDiD() with pytest.raises(ValueError, match="Treatment indicator varies within"): sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[3, 4], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[3, 4], ) def test_constant_treatment_passes(self): @@ -2141,16 +2405,23 @@ def test_constant_treatment_passes(self): for u in range(n_units): is_treated = 1 if u < 3 else 0 for t in range(n_periods): - rows.append({ - "unit": u, "time": t, - "outcome": np.random.randn() + (2.0 if is_treated and t >= 5 else 0), - "treated": is_treated, - }) + rows.append( + { + "unit": u, + "time": t, + "outcome": np.random.randn() + (2.0 if is_treated and t >= 5 else 0), + "treated": is_treated, + } + ) data = pd.DataFrame(rows) sdid = SyntheticDiD() result = sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[5, 6, 7], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[5, 6, 7], ) assert result is not None @@ -2170,11 +2441,14 @@ def test_unbalanced_panel_raises(self): for u in range(6): is_treated = 1 if u < 2 else 0 for t in range(5): - rows.append({ - "unit": u, "time": t, - "outcome": np.random.randn(), - "treated": is_treated, - }) + rows.append( + { + "unit": u, + "time": t, + "outcome": np.random.randn(), + "treated": is_treated, + } + ) data = pd.DataFrame(rows) # Drop one observation to make panel unbalanced data = data[~((data["unit"] == 3) & (data["time"] == 2))].reset_index(drop=True) @@ -2182,8 +2456,12 @@ def test_unbalanced_panel_raises(self): sdid = SyntheticDiD() with pytest.raises(ValueError, match="Panel is not balanced"): sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[3, 4], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[3, 4], ) def test_balanced_panel_passes(self): @@ -2193,16 +2471,23 @@ def test_balanced_panel_passes(self): for u in range(8): is_treated = 1 if u < 2 else 0 for t in range(6): - rows.append({ - "unit": u, "time": t, - "outcome": np.random.randn() + (1.5 if is_treated and t >= 4 else 0), - "treated": is_treated, - }) + rows.append( + { + "unit": u, + "time": t, + "outcome": np.random.randn() + (1.5 if is_treated and t >= 4 else 0), + "treated": is_treated, + } + ) data = pd.DataFrame(rows) sdid = SyntheticDiD() result = sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[4, 5], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[4, 5], ) assert result is not None @@ -2224,23 +2509,30 @@ def test_poor_fit_emits_warning(self): # Large level difference: treated ~100, control ~10 level = 100.0 if is_treated else 10.0 for t in range(8): - rows.append({ - "unit": u, "time": t, - "outcome": level + np.random.randn() * 0.5, - "treated": is_treated, - }) + rows.append( + { + "unit": u, + "time": t, + "outcome": level + np.random.randn() * 0.5, + "treated": is_treated, + } + ) data = pd.DataFrame(rows) sdid = SyntheticDiD() with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[6, 7], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[6, 7], ) fit_warnings = [x for x in w if "Pre-treatment fit is poor" in str(x.message)] - assert len(fit_warnings) >= 1, ( - "Expected warning about poor pre-treatment fit but none was raised" - ) + assert ( + len(fit_warnings) >= 1 + ), "Expected warning about poor pre-treatment fit but none was raised" def test_good_fit_no_warning(self): """Parallel trends data with similar levels should not warn.""" @@ -2250,23 +2542,32 @@ def test_good_fit_no_warning(self): is_treated = 1 if u < 3 else 0 for t in range(8): # Same level, parallel trends, treatment effect only in post - rows.append({ - "unit": u, "time": t, - "outcome": t + np.random.randn() * 0.3 + (2.0 if is_treated and t >= 5 else 0), - "treated": is_treated, - }) + rows.append( + { + "unit": u, + "time": t, + "outcome": t + + np.random.randn() * 0.3 + + (2.0 if is_treated and t >= 5 else 0), + "treated": is_treated, + } + ) data = pd.DataFrame(rows) sdid = SyntheticDiD() with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") sdid.fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="time", post_periods=[5, 6, 7], + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="time", + post_periods=[5, 6, 7], ) fit_warnings = [x for x in w if "Pre-treatment fit is poor" in str(x.message)] - assert len(fit_warnings) == 0, ( - f"Unexpected pre-treatment fit warning: {fit_warnings[0].message}" - ) + assert ( + len(fit_warnings) == 0 + ), f"Unexpected pre-treatment fit warning: {fit_warnings[0].message}" class TestMinDecreaseFloor: @@ -2279,13 +2580,15 @@ def test_floor_equivalence(self): np.random.seed(42) Y = np.tile(np.array([1.0, 2.0, 3.0, 4.0]), (5, 1)) - w_floor = _sc_weight_fw(Y, zeta=0.0, intercept=True, - min_decrease=1e-5, max_iter=10000) - w_zero = _sc_weight_fw(Y, zeta=0.0, intercept=True, - min_decrease=0.0, max_iter=10000) + w_floor = _sc_weight_fw(Y, zeta=0.0, intercept=True, min_decrease=1e-5, max_iter=10000) + w_zero = _sc_weight_fw(Y, zeta=0.0, intercept=True, min_decrease=0.0, max_iter=10000) - np.testing.assert_allclose(w_floor, w_zero, atol=1e-10, - err_msg="Floor min_decrease should give same weights as 0.0") + np.testing.assert_allclose( + w_floor, + w_zero, + atol=1e-10, + err_msg="Floor min_decrease should give same weights as 0.0", + ) # Both should be valid simplex weights assert np.all(w_floor >= -1e-12), "Weights should be non-negative" assert abs(w_floor.sum() - 1.0) < 1e-10, "Weights should sum to 1" @@ -2306,76 +2609,70 @@ def test_placebo_se_matches_across_backends(self): """ import sys - df = _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, - att=5.0, seed=42) + df = _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, att=5.0, seed=42) post_periods = list(range(5, 8)) # Run with default backend (Rust-accelerated inner calls if available) - sdid_default = SyntheticDiD( - variance_method="placebo", n_bootstrap=50, seed=42 - ) - results_default = sdid_default.fit( - df, "outcome", "treated", "unit", "period", post_periods - ) + sdid_default = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=42) + results_default = sdid_default.fit(df, "outcome", "treated", "unit", "period", post_periods) # Run with pure Python backend utils_mod = sys.modules["diff_diff.utils"] with patch.object(utils_mod, "HAS_RUST_BACKEND", False): - sdid_py = SyntheticDiD( - variance_method="placebo", n_bootstrap=50, seed=42 - ) + sdid_py = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=42) results_py = sdid_py.fit( df.copy(), "outcome", "treated", "unit", "period", post_periods ) # ATT must be identical (same data, same algorithm) np.testing.assert_allclose( - results_default.att, results_py.att, rtol=1e-6, - err_msg="ATT should match between backends" + results_default.att, + results_py.att, + rtol=1e-6, + err_msg="ATT should match between backends", ) # SE must match within tolerance (faer vs NumPy in FW inner loop) np.testing.assert_allclose( - results_default.se, results_py.se, rtol=1e-4, - err_msg="Placebo SE should match between backends" + results_default.se, + results_py.se, + rtol=1e-4, + err_msg="Placebo SE should match between backends", ) def test_bootstrap_se_matches_across_backends(self): """SyntheticDiD bootstrap SE is identical regardless of backend.""" import sys - df = _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, - att=5.0, seed=42) + df = _make_panel(n_control=20, n_treated=3, n_pre=5, n_post=3, att=5.0, seed=42) post_periods = list(range(5, 8)) # Run with default backend - sdid_default = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=42 - ) - results_default = sdid_default.fit( - df, "outcome", "treated", "unit", "period", post_periods - ) + sdid_default = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=42) + results_default = sdid_default.fit(df, "outcome", "treated", "unit", "period", post_periods) # Run with pure Python backend utils_mod = sys.modules["diff_diff.utils"] with patch.object(utils_mod, "HAS_RUST_BACKEND", False): - sdid_py = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=50, seed=42 - ) + sdid_py = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=42) results_py = sdid_py.fit( df.copy(), "outcome", "treated", "unit", "period", post_periods ) # ATT must match np.testing.assert_allclose( - results_default.att, results_py.att, rtol=1e-6, - err_msg="ATT should match between backends" + results_default.att, + results_py.att, + rtol=1e-6, + err_msg="ATT should match between backends", ) # Bootstrap SE must match (same RNG, same loop, only inner FW differs) np.testing.assert_allclose( - results_default.se, results_py.se, rtol=1e-4, - err_msg="Bootstrap SE should match between backends" + results_default.se, + results_py.se, + rtol=1e-4, + err_msg="Bootstrap SE should match between backends", ) @@ -2384,8 +2681,8 @@ class TestEmptyPostGuard: def test_empty_post_raises(self): """compute_time_weights raises ValueError when Y_post_control has 0 rows.""" - Y_pre = np.random.randn(4, 5) # 4 pre-periods, 5 controls - Y_post = np.empty((0, 5)) # 0 post-periods + Y_pre = np.random.randn(4, 5) # 4 pre-periods, 5 controls + Y_post = np.empty((0, 5)) # 0 post-periods with pytest.raises(ValueError, match="Y_post_control has no rows"): compute_time_weights(Y_pre, Y_post, zeta_lambda=0.0) @@ -2403,9 +2700,14 @@ class TestFitSnapshot: def test_shapes_and_ordering(self): df = _make_panel(n_control=10, n_treated=2, n_pre=5, n_post=3, seed=42) sdid = SyntheticDiD(variance_method="jackknife", seed=42) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(5, 8))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(5, 8)), + ) snap = res._fit_snapshot assert snap is not None assert snap.Y_pre_control.shape == (5, 10) @@ -2422,8 +2724,7 @@ def test_shapes_and_ordering(self): def test_matrices_are_read_only(self): df = _make_panel(seed=7) sdid = SyntheticDiD(variance_method="jackknife", seed=7) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") snap = res._fit_snapshot assert not snap.Y_pre_control.flags.writeable assert not snap.Y_post_control.flags.writeable @@ -2439,19 +2740,15 @@ class TestTrajectories: def test_synthetic_pre_matches_weighted_sum(self): df = _make_panel(seed=11) sdid = SyntheticDiD(variance_method="jackknife", seed=11) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") - omega_eff = np.array( - [res.unit_weights[u] for u in res._fit_snapshot.control_unit_ids] - ) + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") + omega_eff = np.array([res.unit_weights[u] for u in res._fit_snapshot.control_unit_ids]) expected = res._fit_snapshot.Y_pre_control @ omega_eff assert np.allclose(res.synthetic_pre_trajectory, expected, atol=1e-12) def test_treated_trajectory_matches_mean(self): df = _make_panel(seed=13) sdid = SyntheticDiD(variance_method="jackknife", seed=13) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") assert np.allclose( res.treated_pre_trajectory, np.mean(res._fit_snapshot.Y_pre_treated, axis=1), @@ -2475,8 +2772,11 @@ def test_treated_trajectory_survey_weighted(self): df = df.assign(weight=df["unit"].map(w_by_unit)) sdid = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=83) res = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", survey_design=SurveyDesign(weights="weight", weight_type="pweight"), ) snap = res._fit_snapshot @@ -2486,19 +2786,14 @@ def test_treated_trajectory_survey_weighted(self): assert np.allclose(res.treated_pre_trajectory, expected_pre, atol=1e-12) assert np.allclose(res.treated_post_trajectory, expected_post, atol=1e-12) - omega_eff = np.array( - [res.unit_weights[u] for u in snap.control_unit_ids] - ) + omega_eff = np.array([res.unit_weights[u] for u in snap.control_unit_ids]) expected_synth_pre = snap.Y_pre_control @ omega_eff - assert np.allclose( - res.synthetic_pre_trajectory, expected_synth_pre, atol=1e-12 - ) + assert np.allclose(res.synthetic_pre_trajectory, expected_synth_pre, atol=1e-12) def test_public_trajectory_arrays_are_read_only(self): df = _make_panel(seed=87) sdid = SyntheticDiD(variance_method="jackknife", seed=87) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") for arr in ( res.synthetic_pre_trajectory, res.synthetic_post_trajectory, @@ -2511,14 +2806,9 @@ def test_public_trajectory_arrays_are_read_only(self): def test_pre_fit_rmse_recoverable(self): df = _make_panel(seed=17) sdid = SyntheticDiD(variance_method="jackknife", seed=17) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") rmse = float( - np.sqrt( - np.mean( - (res.treated_pre_trajectory - res.synthetic_pre_trajectory) ** 2 - ) - ) + np.sqrt(np.mean((res.treated_pre_trajectory - res.synthetic_pre_trajectory) ** 2)) ) assert abs(rmse - res.pre_treatment_fit) < 1e-10 @@ -2527,12 +2817,16 @@ class TestLooEffectsDf: """get_loo_effects_df joins jackknife pseudo-values to unit identities.""" def _fit_jackknife(self, seed=42, **kwargs): - df = _make_panel(n_control=10, n_treated=3, n_pre=5, n_post=3, - seed=seed, **kwargs) + df = _make_panel(n_control=10, n_treated=3, n_pre=5, n_post=3, seed=seed, **kwargs) sdid = SyntheticDiD(variance_method="jackknife", seed=seed) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(5, 8))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(5, 8)), + ) return res def test_columns_and_roles(self): @@ -2542,9 +2836,7 @@ def test_columns_and_roles(self): assert set(loo["role"]) == {"control", "treated"} assert (loo["role"] == "control").sum() == res.n_control assert (loo["role"] == "treated").sum() == res.n_treated - assert set(loo["unit"]) == set(res.unit_weights) | set( - res._fit_snapshot.treated_unit_ids - ) + assert set(loo["unit"]) == set(res.unit_weights) | set(res._fit_snapshot.treated_unit_ids) def test_delta_from_full_matches_math(self): res = self._fit_jackknife() @@ -2561,16 +2853,14 @@ def test_sorted_by_absolute_delta_descending(self): def test_placebo_raises_value_error(self): df = _make_panel(seed=21) sdid = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=21) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") with pytest.raises(ValueError, match="variance_method='jackknife'"): res.get_loo_effects_df() def test_bootstrap_raises_value_error(self): df = _make_panel(seed=23) sdid = SyntheticDiD(variance_method="bootstrap", n_bootstrap=50, seed=23) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") with pytest.raises(ValueError, match="variance_method='jackknife'"): res.get_loo_effects_df() @@ -2604,10 +2894,18 @@ def test_equal_weights_yield_effective_n_equal_to_count(self): uniform = {f"u{i}": 0.1 for i in range(10)} res = SyntheticDiDResults( - att=0.0, se=0.0, t_stat=0.0, p_value=1.0, conf_int=(0.0, 0.0), - n_obs=100, n_treated=1, n_control=10, - unit_weights=uniform, time_weights={}, - pre_periods=[], post_periods=[], + att=0.0, + se=0.0, + t_stat=0.0, + p_value=1.0, + conf_int=(0.0, 0.0), + n_obs=100, + n_treated=1, + n_control=10, + unit_weights=uniform, + time_weights={}, + pre_periods=[], + post_periods=[], ) c = res.get_weight_concentration() assert abs(c["effective_n"] - 10.0) < 1e-10 @@ -2619,10 +2917,18 @@ def test_single_unit_yields_effective_n_one(self): from diff_diff.results import SyntheticDiDResults res = SyntheticDiDResults( - att=0.0, se=0.0, t_stat=0.0, p_value=1.0, conf_int=(0.0, 0.0), - n_obs=10, n_treated=1, n_control=1, - unit_weights={"only": 1.0}, time_weights={}, - pre_periods=[], post_periods=[], + att=0.0, + se=0.0, + t_stat=0.0, + p_value=1.0, + conf_int=(0.0, 0.0), + n_obs=10, + n_treated=1, + n_control=1, + unit_weights={"only": 1.0}, + time_weights={}, + pre_periods=[], + post_periods=[], ) c = res.get_weight_concentration() assert abs(c["effective_n"] - 1.0) < 1e-10 @@ -2632,10 +2938,18 @@ def test_top_k_clamped_to_n(self): from diff_diff.results import SyntheticDiDResults res = SyntheticDiDResults( - att=0.0, se=0.0, t_stat=0.0, p_value=1.0, conf_int=(0.0, 0.0), - n_obs=30, n_treated=1, n_control=3, - unit_weights={"a": 0.5, "b": 0.3, "c": 0.2}, time_weights={}, - pre_periods=[], post_periods=[], + att=0.0, + se=0.0, + t_stat=0.0, + p_value=1.0, + conf_int=(0.0, 0.0), + n_obs=30, + n_treated=1, + n_control=3, + unit_weights={"a": 0.5, "b": 0.3, "c": 0.2}, + time_weights={}, + pre_periods=[], + post_periods=[], ) c = res.get_weight_concentration(top_k=10) assert c["top_k"] == 3 @@ -2645,10 +2959,18 @@ def test_negative_top_k_raises(self): from diff_diff.results import SyntheticDiDResults res = SyntheticDiDResults( - att=0.0, se=0.0, t_stat=0.0, p_value=1.0, conf_int=(0.0, 0.0), - n_obs=30, n_treated=1, n_control=3, - unit_weights={"a": 0.5, "b": 0.3, "c": 0.2}, time_weights={}, - pre_periods=[], post_periods=[], + att=0.0, + se=0.0, + t_stat=0.0, + p_value=1.0, + conf_int=(0.0, 0.0), + n_obs=30, + n_treated=1, + n_control=3, + unit_weights={"a": 0.5, "b": 0.3, "c": 0.2}, + time_weights={}, + pre_periods=[], + post_periods=[], ) with pytest.raises(ValueError, match="top_k must be non-negative"): res.get_weight_concentration(top_k=-1) @@ -2656,7 +2978,6 @@ def test_negative_top_k_raises(self): def test_uses_composed_weights_under_survey(self): """Metrics come from self.unit_weights which stores composed ω_eff for survey fits.""" - import pandas as pd from diff_diff.survey import SurveyDesign df = _make_panel(n_control=8, n_treated=2, n_pre=4, n_post=2, seed=29) @@ -2666,8 +2987,11 @@ def test_uses_composed_weights_under_survey(self): df = df.assign(weight=df["unit"].map(w_by_unit)) sdid = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=29) res = sdid.fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", survey_design=SurveyDesign(weights="weight", weight_type="pweight"), ) c = res.get_weight_concentration() @@ -2675,7 +2999,7 @@ def test_uses_composed_weights_under_survey(self): import numpy as np weights = np.array(list(res.unit_weights.values())) - expected_effective_n = 1.0 / np.sum(weights ** 2) + expected_effective_n = 1.0 / np.sum(weights**2) assert abs(c["effective_n"] - expected_effective_n) < 1e-10 @@ -2683,29 +3007,40 @@ class TestInTimePlacebo: """in_time_placebo re-estimates on shifted fake dates.""" def test_default_sweep_shape(self): - df = _make_panel(n_control=10, n_treated=2, n_pre=8, n_post=3, - att=0.0, seed=31) + df = _make_panel(n_control=10, n_treated=2, n_pre=8, n_post=3, att=0.0, seed=31) sdid = SyntheticDiD(variance_method="jackknife", seed=31) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(8, 11))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(8, 11)), + ) placebo = res.in_time_placebo() # Feasible positions: i in [2, n_pre - 1] = [2, 7] -> 6 rows assert len(placebo) == 6 assert list(placebo.columns) == [ - "fake_treatment_period", "att", "pre_fit_rmse", - "n_pre_fake", "n_post_fake", + "fake_treatment_period", + "att", + "pre_fit_rmse", + "n_pre_fake", + "n_post_fake", ] def test_no_effect_dgp_gives_small_placebo_atts(self): """On a clean no-effect DGP, placebo ATTs should be small vs the DGP's noise level.""" - df = _make_panel(n_control=20, n_treated=3, n_pre=8, n_post=3, - att=0.0, seed=33) + df = _make_panel(n_control=20, n_treated=3, n_pre=8, n_post=3, att=0.0, seed=33) sdid = SyntheticDiD(variance_method="jackknife", seed=33) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(8, 11))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(8, 11)), + ) placebo = res.in_time_placebo() median_abs = float(placebo["att"].abs().median()) # Noise sd in _make_panel is 0.5; ATTs should be well under that at @@ -2715,9 +3050,14 @@ def test_no_effect_dgp_gives_small_placebo_atts(self): def test_explicit_list_overrides_sweep(self): df = _make_panel(n_control=10, n_treated=2, n_pre=8, n_post=3, seed=37) sdid = SyntheticDiD(variance_method="jackknife", seed=37) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(8, 11))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(8, 11)), + ) placebo = res.in_time_placebo(fake_treatment_periods=[4, 6]) assert len(placebo) == 2 assert set(placebo["fake_treatment_period"]) == {4, 6} @@ -2725,18 +3065,28 @@ def test_explicit_list_overrides_sweep(self): def test_post_period_raises(self): df = _make_panel(n_control=10, n_treated=2, n_pre=6, n_post=3, seed=41) sdid = SyntheticDiD(variance_method="jackknife", seed=41) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(6, 9))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(6, 9)), + ) with pytest.raises(ValueError, match="post_periods"): res.in_time_placebo(fake_treatment_periods=[6]) def test_missing_period_raises(self): df = _make_panel(n_control=10, n_treated=2, n_pre=6, n_post=3, seed=43) sdid = SyntheticDiD(variance_method="jackknife", seed=43) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(6, 9))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(6, 9)), + ) with pytest.raises(ValueError, match="not found in pre_periods"): res.in_time_placebo(fake_treatment_periods=[999]) @@ -2745,22 +3095,35 @@ def test_empty_default_sweep_preserves_schema(self): still carry the documented columns.""" df = _make_panel(n_control=10, n_treated=2, n_pre=2, n_post=3, seed=50) sdid = SyntheticDiD(variance_method="jackknife", seed=50) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(2, 5))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(2, 5)), + ) placebo = res.in_time_placebo() assert len(placebo) == 0 assert list(placebo.columns) == [ - "fake_treatment_period", "att", "pre_fit_rmse", - "n_pre_fake", "n_post_fake", + "fake_treatment_period", + "att", + "pre_fit_rmse", + "n_pre_fake", + "n_post_fake", ] def test_zeta_override_changes_result(self): df = _make_panel(n_control=10, n_treated=2, n_pre=8, n_post=3, seed=47) sdid = SyntheticDiD(variance_method="jackknife", seed=47) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(8, 11))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(8, 11)), + ) default_ = res.in_time_placebo(fake_treatment_periods=[5]) override_ = res.in_time_placebo( fake_treatment_periods=[5], @@ -2778,12 +3141,16 @@ class TestSensitivityToZetaOmega: """sensitivity_to_zeta_omega sweeps regularization values.""" def test_multiplier_one_reproduces_original_att(self): - df = _make_panel(n_control=12, n_treated=2, n_pre=6, n_post=3, - att=2.0, seed=53) + df = _make_panel(n_control=12, n_treated=2, n_pre=6, n_post=3, att=2.0, seed=53) sdid = SyntheticDiD(variance_method="jackknife", seed=53) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(6, 9))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(6, 9)), + ) sens = res.sensitivity_to_zeta_omega() # Row where zeta_omega == self.zeta_omega (multiplier 1.0) match = sens.loc[np.isclose(sens["zeta_omega"], res.zeta_omega)] @@ -2793,8 +3160,7 @@ def test_multiplier_one_reproduces_original_att(self): def test_grid_length_matches_default_multipliers(self): df = _make_panel(seed=59) sdid = SyntheticDiD(variance_method="jackknife", seed=59) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega() assert len(sens) == 5 @@ -2803,8 +3169,7 @@ def test_default_grid_values_match_documented_multipliers(self): so the contract cannot drift unnoticed.""" df = _make_panel(seed=60) sdid = SyntheticDiD(variance_method="jackknife", seed=60) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega() expected = np.array([0.25, 0.5, 1.0, 2.0, 4.0]) * res.zeta_omega assert np.allclose(sens["zeta_omega"].to_numpy(), expected, atol=1e-12) @@ -2812,8 +3177,7 @@ def test_default_grid_values_match_documented_multipliers(self): def test_explicit_grid_overrides_multipliers(self): df = _make_panel(seed=61) sdid = SyntheticDiD(variance_method="jackknife", seed=61) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega(zeta_grid=[0.1, 1.0, 10.0]) assert len(sens) == 3 assert np.allclose(sens["zeta_omega"], [0.1, 1.0, 10.0]) @@ -2823,9 +3187,14 @@ def test_effective_n_grows_with_zeta(self): should be monotone non-decreasing across the default grid.""" df = _make_panel(n_control=15, n_treated=2, n_pre=8, n_post=3, seed=67) sdid = SyntheticDiD(variance_method="jackknife", seed=67) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(8, 11))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(8, 11)), + ) sens = res.sensitivity_to_zeta_omega() eff_n = sens["effective_n"].to_numpy() # Allow tiny wobble from solver tolerance @@ -2834,36 +3203,42 @@ def test_effective_n_grows_with_zeta(self): def test_columns_shape(self): df = _make_panel(seed=71) sdid = SyntheticDiD(variance_method="jackknife", seed=71) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega() assert list(sens.columns) == [ - "zeta_omega", "att", "pre_fit_rmse", - "max_unit_weight", "effective_n", + "zeta_omega", + "att", + "pre_fit_rmse", + "max_unit_weight", + "effective_n", ] def test_empty_zeta_grid_preserves_schema(self): df = _make_panel(seed=91) sdid = SyntheticDiD(variance_method="jackknife", seed=91) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega(zeta_grid=[]) assert sens.shape == (0, 5) assert list(sens.columns) == [ - "zeta_omega", "att", "pre_fit_rmse", - "max_unit_weight", "effective_n", + "zeta_omega", + "att", + "pre_fit_rmse", + "max_unit_weight", + "effective_n", ] def test_empty_multipliers_preserves_schema(self): df = _make_panel(seed=93) sdid = SyntheticDiD(variance_method="jackknife", seed=93) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") sens = res.sensitivity_to_zeta_omega(multipliers=()) assert sens.shape == (0, 5) assert list(sens.columns) == [ - "zeta_omega", "att", "pre_fit_rmse", - "max_unit_weight", "effective_n", + "zeta_omega", + "att", + "pre_fit_rmse", + "max_unit_weight", + "effective_n", ] @@ -2878,8 +3253,7 @@ def test_snippets_reference_existing_methods(self): df = _make_panel(seed=73) sdid = SyntheticDiD(variance_method="jackknife", seed=73) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") steps, _ = _handle_synthetic(res) expected_methods = [ "get_weight_concentration", @@ -2889,9 +3263,7 @@ def test_snippets_reference_existing_methods(self): ] joined_code = "\n".join(step["code"] for step in steps) for method in expected_methods: - assert method in joined_code, ( - f"Expected {method}() referenced in practitioner guidance" - ) + assert method in joined_code, f"Expected {method}() referenced in practitioner guidance" assert hasattr(SyntheticDiDResults, method), ( f"Practitioner guidance references {method}() but it's not " "on SyntheticDiDResults" @@ -2900,12 +3272,12 @@ def test_snippets_reference_existing_methods(self): def test_snippets_parse_as_python(self): """Each snippet in _handle_synthetic should be syntactically valid.""" import ast + from diff_diff.practitioner import _handle_synthetic df = _make_panel(seed=77) sdid = SyntheticDiD(variance_method="jackknife", seed=77) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + res = sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") steps, _ = _handle_synthetic(res) for step in steps: ast.parse(step["code"]) @@ -2920,19 +3292,23 @@ def test_jackknife_loo_snippet_handles_unavailable_loo(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") sdid = SyntheticDiD(variance_method="jackknife", seed=97) - res = sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period", - post_periods=list(range(5, 8))) + res = sdid.fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + post_periods=list(range(5, 8)), + ) assert res.variance_method == "jackknife" assert res._loo_unit_ids is None # LOO intentionally unavailable steps, _ = _handle_synthetic(res) - loo_snippet = next( - s["code"] for s in steps if "get_loo_effects_df" in s["code"] - ) + loo_snippet = next(s["code"] for s in steps if "get_loo_effects_df" in s["code"]) # Executing the snippet against this result must not raise. - import io import contextlib + import io + captured = io.StringIO() with contextlib.redirect_stdout(captured): exec(loo_snippet, {"results": res}) @@ -2946,8 +3322,7 @@ class TestSyntheticDiDResultsPickle: def _fit(self, seed=101): df = _make_panel(seed=seed) sdid = SyntheticDiD(variance_method="jackknife", seed=seed) - return sdid.fit(df, outcome="outcome", treatment="treated", - unit="unit", time="period") + return sdid.fit(df, outcome="outcome", treatment="treated", unit="unit", time="period") def test_snapshot_dropped_on_pickle(self): import pickle @@ -2960,9 +3335,7 @@ def test_snapshot_dropped_on_pickle(self): # Public fields survive assert restored.att == res.att assert restored.se == res.se - assert np.allclose( - restored.synthetic_pre_trajectory, res.synthetic_pre_trajectory - ) + assert np.allclose(restored.synthetic_pre_trajectory, res.synthetic_pre_trajectory) def test_legacy_pickle_state_maps_placebo_effects(self): """A result pickled before the placebo_effects → variance_effects @@ -3071,7 +3444,7 @@ class TestScaleEquivariance: # protection moves to ``test_placebo_se_matches_r``, which # bypasses the platform-divergent fit-time path through a test # seam. Pinned value is the Linux/OpenBLAS capture. - "placebo": (4.603349837478791, 0.2938403592163006, 0.004975124378109453, 200), + "placebo": (4.603349837478791, 0.2938403592163006, 0.004975124378109453, 200), # bootstrap = paper-faithful refit with R-default warm-start: FW is # initialized with ``sum_normalize(unit_weights[boot_control_idx])`` # for ω and with the fit-time ``time_weights`` for λ on each draw, @@ -3082,8 +3455,8 @@ class TestScaleEquivariance: # init; strict-convexity of the FW objective means the converged # answer is unique, so the warm-start matches R more faithfully on # problems where the pre-sparsify budget is tight. - "bootstrap": (4.6033498374787865, 0.21427381053829253, 2.2215821875845446e-102, 200), - "jackknife": (4.603349837478791, 0.19908075946622925, 2.716551077849484e-118, 23), + "bootstrap": (4.6033498374787865, 0.21427381053829253, 2.2215821875845446e-102, 200), + "jackknife": (4.603349837478791, 0.19908075946622925, 2.716551077849484e-118, 23), } # (a, b) pairs. Includes extreme scales where pre-fix SDID loses @@ -3191,19 +3564,24 @@ def test_detects_true_effect_at_extreme_scale(self, variance_method): y = baseline_level + unit_fe + t * 3e5 + rng.normal(0, 5e5) if is_treated and t >= n_pre: y += true_att - rows.append({ - "unit": unit, "period": t, - "treated": int(is_treated), "outcome": y, - }) + rows.append( + { + "unit": unit, + "period": t, + "treated": int(is_treated), + "outcome": y, + } + ) data = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - r = SyntheticDiD( - variance_method=variance_method, n_bootstrap=200, seed=7 - ).fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="period", + r = SyntheticDiD(variance_method=variance_method, n_bootstrap=200, seed=7).fit( + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=list(range(n_pre, n_pre + n_post)), ) @@ -3248,14 +3626,17 @@ def test_bootstrap_p_value_matches_analytical(self, ci_params): r = SyntheticDiD( variance_method="bootstrap", n_bootstrap=ci_params.bootstrap(200), seed=1 ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) _, expected_p, _ = safe_inference(r.att, r.se, alpha=0.05) - assert abs(r.p_value - expected_p) < 1e-12, ( - f"bootstrap p_value={r.p_value} != analytical {expected_p}" - ) + assert ( + abs(r.p_value - expected_p) < 1e-12 + ), f"bootstrap p_value={r.p_value} != analytical {expected_p}" def test_placebo_p_value_uses_empirical_formula(self, ci_params): """Placebo p-value must equal max(mean(|draws| >= |att|), 1/(r+1)).""" @@ -3267,16 +3648,19 @@ def test_placebo_p_value_uses_empirical_formula(self, ci_params): r = SyntheticDiD( variance_method="placebo", n_bootstrap=ci_params.bootstrap(200), seed=1 ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) variance_effects = np.asarray(r.variance_effects) empirical_p = float(np.mean(np.abs(variance_effects) >= np.abs(r.att))) expected_p = max(empirical_p, 1.0 / (len(variance_effects) + 1)) - assert abs(r.p_value - expected_p) < 1e-12, ( - f"placebo p_value={r.p_value} != empirical {expected_p}" - ) + assert ( + abs(r.p_value - expected_p) < 1e-12 + ), f"placebo p_value={r.p_value} != empirical {expected_p}" def test_bootstrap_p_value_detects_large_effect(self): """Bootstrap p-value must reject decisively when z is large. @@ -3288,18 +3672,17 @@ def test_bootstrap_p_value_detects_large_effect(self): df = _make_panel(att=5.0, seed=123) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - r = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=200, seed=1 - ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + r = SyntheticDiD(variance_method="bootstrap", n_bootstrap=200, seed=1).fit( + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) z = abs(r.att / r.se) assert z > 6, f"setup error: z={z} too small to test rejection" - assert r.p_value < 1e-6, ( - f"bootstrap p_value={r.p_value} too large at z={z}" - ) + assert r.p_value < 1e-6, f"bootstrap p_value={r.p_value} too large at z={z}" @pytest.mark.slow def test_bootstrap_p_value_null_dispersion(self): @@ -3331,10 +3714,15 @@ def test_bootstrap_p_value_null_dispersion(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) r = SyntheticDiD( - variance_method="bootstrap", n_bootstrap=200, seed=seed, + variance_method="bootstrap", + n_bootstrap=200, + seed=seed, ).fit( - df, outcome="outcome", treatment="treated", - unit="unit", time="period", + df, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) if np.isfinite(r.p_value): @@ -3376,8 +3764,11 @@ def _rescale(df, a, b): @staticmethod def _fit(data, seed=1): return SyntheticDiD(variance_method="jackknife", seed=seed).fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="period", + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) @@ -3400,15 +3791,13 @@ def test_in_time_placebo_scale_equivariance(self): placebo = r.in_time_placebo(fake_treatment_periods=fake_periods) assert list(placebo["fake_treatment_period"]) == fake_periods - for row0, row in zip(placebo0.to_dict("records"), - placebo.to_dict("records")): + for row0, row in zip(placebo0.to_dict("records"), placebo.to_dict("records")): if np.isnan(row0["att"]): assert np.isnan(row["att"]), f"att at (a={a}, b={b})" assert np.isnan(row["pre_fit_rmse"]) continue assert row["att"] / a == pytest.approx(row0["att"], rel=1e-6), ( - f"att at (a={a}, b={b}), " - f"fake_period={row0['fake_treatment_period']}" + f"att at (a={a}, b={b}), " f"fake_period={row0['fake_treatment_period']}" ) assert row["pre_fit_rmse"] / abs(a) == pytest.approx( row0["pre_fit_rmse"], rel=1e-6 @@ -3433,24 +3822,15 @@ def test_sensitivity_to_zeta_omega_scale_equivariance(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) r = self._fit(scaled) - sens = r.sensitivity_to_zeta_omega( - zeta_grid=[a * z for z in zeta_grid] - ) - for row0, row in zip(sens0.to_dict("records"), - sens.to_dict("records")): - assert row["att"] / a == pytest.approx(row0["att"], rel=1e-6), ( - f"att at (a={a}, b={b}), zeta={row0['zeta_omega']}" - ) - assert row["pre_fit_rmse"] / abs(a) == pytest.approx( - row0["pre_fit_rmse"], rel=1e-6 - ) + sens = r.sensitivity_to_zeta_omega(zeta_grid=[a * z for z in zeta_grid]) + for row0, row in zip(sens0.to_dict("records"), sens.to_dict("records")): + assert row["att"] / a == pytest.approx( + row0["att"], rel=1e-6 + ), f"att at (a={a}, b={b}), zeta={row0['zeta_omega']}" + assert row["pre_fit_rmse"] / abs(a) == pytest.approx(row0["pre_fit_rmse"], rel=1e-6) # omega-derived diagnostics are scale-invariant. - assert row["max_unit_weight"] == pytest.approx( - row0["max_unit_weight"], rel=1e-6 - ) - assert row["effective_n"] == pytest.approx( - row0["effective_n"], rel=1e-6 - ) + assert row["max_unit_weight"] == pytest.approx(row0["max_unit_weight"], rel=1e-6) + assert row["effective_n"] == pytest.approx(row0["effective_n"], rel=1e-6) def test_in_time_placebo_detectable_at_extreme_scale(self): """Pre-fix regression: at Y~1e9 the placebo re-fit corrupted ATTs via @@ -3465,8 +3845,9 @@ def test_in_time_placebo_detectable_at_extreme_scale(self): unit_fe = rng.normal(0, 2e6) for t in range(n_pre + n_post): y = baseline_level + unit_fe + t * 3e5 + rng.normal(0, 5e5) - rows.append({"unit": unit, "period": t, - "treated": int(unit >= n_control), "outcome": y}) + rows.append( + {"unit": unit, "period": t, "treated": int(unit >= n_control), "outcome": y} + ) data = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) @@ -3496,8 +3877,11 @@ def test_legacy_snapshot_defaults_are_noop(self): with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) r = SyntheticDiD(variance_method="jackknife", seed=1).fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="period", + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) @@ -3555,8 +3939,11 @@ class TestHeterogeneousAndRampingScale: @staticmethod def _fit(data, seed=1): return SyntheticDiD(variance_method="jackknife", seed=seed).fit( - data, outcome="outcome", treatment="treated", - unit="unit", time="period", + data, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", post_periods=[5, 6, 7], ) @@ -3575,8 +3962,7 @@ def test_cross_unit_heterogeneous_scale(self): y = unit_level * (1 + 0.02 * t) + rng.normal(0, unit_level * 0.01) if is_treated and t >= n_pre: y += 0.05 * unit_level - rows.append({"unit": unit, "period": t, - "treated": int(is_treated), "outcome": y}) + rows.append({"unit": unit, "period": t, "treated": int(is_treated), "outcome": y}) data = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) @@ -3606,8 +3992,7 @@ def test_cross_period_ramping_trend(self): y = trend + unit_fe * trend * 0.01 + rng.normal(0, trend * 0.005) if is_treated and t >= n_pre: y += 0.01 * trend - rows.append({"unit": unit, "period": t, - "treated": int(is_treated), "outcome": y}) + rows.append({"unit": unit, "period": t, "treated": int(is_treated), "outcome": y}) data = pd.DataFrame(rows) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) @@ -3649,8 +4034,7 @@ def test_coverage_artifacts_present(self): import pathlib artifact = ( - pathlib.Path(__file__).parent.parent - / "benchmarks" / "data" / "sdid_coverage.json" + pathlib.Path(__file__).parent.parent / "benchmarks" / "data" / "sdid_coverage.json" ) if not artifact.exists(): pytest.skip( @@ -3664,8 +4048,15 @@ def test_coverage_artifacts_present(self): assert key in payload, f"missing top-level key: {key}" meta = payload["metadata"] - for key in ("n_seeds", "n_bootstrap", "library_version", "backend", - "generated_at", "methods", "alphas"): + for key in ( + "n_seeds", + "n_bootstrap", + "library_version", + "backend", + "generated_at", + "methods", + "alphas", + ): assert key in meta, f"missing metadata key: {key}" assert meta["n_seeds"] >= 100, ( f"n_seeds={meta['n_seeds']} too small; the REGISTRY calibration " @@ -3682,19 +4073,20 @@ def test_coverage_artifacts_present(self): assert dgp in payload["per_dgp"], f"missing DGP block: {dgp}" per_method = payload["per_dgp"][dgp] for method in ("placebo", "bootstrap", "jackknife"): - assert method in per_method, ( - f"missing method block {method!r} under DGP {dgp!r}" - ) + assert method in per_method, f"missing method block {method!r} under DGP {dgp!r}" block = per_method[method] - for field in ("rejection_rate", "mean_se", "true_sd_tau_hat", - "se_over_truesd", "n_successful_fits"): - assert field in block, ( - f"missing field {field!r} in {dgp}/{method}" - ) + for field in ( + "rejection_rate", + "mean_se", + "true_sd_tau_hat", + "se_over_truesd", + "n_successful_fits", + ): + assert field in block, f"missing field {field!r} in {dgp}/{method}" for alpha_key in ("0.01", "0.05", "0.10"): - assert alpha_key in block["rejection_rate"], ( - f"missing alpha {alpha_key} in {dgp}/{method} rejection_rate" - ) + assert ( + alpha_key in block["rejection_rate"] + ), f"missing alpha {alpha_key} in {dgp}/{method} rejection_rate" # Post-PR #365: stratified_survey runs bootstrap (validation # gate) + jackknife (anti-conservative but reported for @@ -3736,4 +4128,3 @@ def test_coverage_artifacts_present(self): "the PSU-level LOO + stratum aggregation path is broken if " "this drops to 0." ) - diff --git a/tests/test_methodology_triple_diff.py b/tests/test_methodology_triple_diff.py index d9fcc4df7..98aa218d2 100644 --- a/tests/test_methodology_triple_diff.py +++ b/tests/test_methodology_triple_diff.py @@ -1421,7 +1421,7 @@ def test_no_overlap_warning_for_reg(self): ddd = TripleDifference(estimation_method="reg") with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - result = ddd.fit( + ddd.fit( data, outcome="outcome", group="group", @@ -1444,9 +1444,7 @@ def _failing_lr(*args, **kwargs): monkeypatch.setattr(td_module, "solve_logit", _failing_lr) - ddd = TripleDifference( - estimation_method=method, pscore_fallback="unconditional" - ) + ddd = TripleDifference(estimation_method=method, pscore_fallback="unconditional") with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") result = ddd.fit( @@ -1601,12 +1599,20 @@ def test_constant_covariate_finite_se_matches_drop_one(self, method): with warnings.catch_warnings(): warnings.simplefilter("ignore") drop_one = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age"], ) with_const = TripleDifference(estimation_method=method).fit( - data_const, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "xc"], + data_const, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "xc"], ) assert np.isfinite(with_const.se) @@ -1620,8 +1626,12 @@ def test_constant_covariate_emits_single_rank_guard_warning(self): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") TripleDifference(estimation_method="reg").fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "xc"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "xc"], ) rank_guard = [w for w in caught if "rank-guarded inverse" in str(w.message)] # The per-fit aggregate warning fires exactly once, not per comparison. @@ -1632,8 +1642,12 @@ def test_well_conditioned_covariates_no_rank_guard_warning(self): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") res = TripleDifference(estimation_method="dr").fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "education"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "education"], ) assert not any("rank-guarded inverse" in str(w.message) for w in caught) assert np.isfinite(res.se) @@ -1652,12 +1666,22 @@ def test_survey_weighted_constant_covariate_finite_se(self, method): with warnings.catch_warnings(): warnings.simplefilter("ignore") drop_one = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age"], survey_design=sd, + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age"], + survey_design=sd, ) with_const = TripleDifference(estimation_method=method).fit( - data_const, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "xc"], survey_design=sd, + data_const, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "xc"], + survey_design=sd, ) assert np.isfinite(with_const.se) assert with_const.se < 1.0 @@ -1675,11 +1699,13 @@ def test_error_mode_raises_before_rank_guard(self, method): data = generate_ddd_data(n_per_cell=200, seed=42, add_covariates=True) data["agec"] = 2.0 * data["age"] # exactly collinear with age with pytest.raises(ValueError, match="(?i)rank-deficient"): - TripleDifference( - estimation_method=method, rank_deficient_action="error" - ).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "agec"], + TripleDifference(estimation_method=method, rank_deficient_action="error").fit( + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "agec"], ) @pytest.mark.parametrize("method", ["reg", "dr"]) @@ -1706,36 +1732,49 @@ def test_cell_aliasing_rank_guard(self, method, cell): with warnings.catch_warnings(): warnings.simplefilter("ignore") drop_one = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age"], ) with_deg = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "aged"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "aged"], ) assert np.isfinite(with_deg.se) and with_deg.se > 0 np.testing.assert_allclose(with_deg.se, drop_one.se, rtol=5e-2) else: # treated-pre cell: compare to the near-collinear (full-rank) limit - mask = ( - (data["group"] == 1) & (data["partition"] == 1) & (data["time"] == 0) - ) + mask = (data["group"] == 1) & (data["partition"] == 1) & (data["time"] == 0) base = np.where(mask, 2.0 * data["age"], rng.normal(size=len(data))) data["aged_exact"] = base near = base.copy() - near[mask.to_numpy()] = ( - 2.0 * data.loc[mask, "age"].to_numpy() - + 1e-6 * rng.standard_normal(int(mask.sum())) - ) + near[mask.to_numpy()] = 2.0 * data.loc[ + mask, "age" + ].to_numpy() + 1e-6 * rng.standard_normal(int(mask.sum())) data["aged_near"] = near with warnings.catch_warnings(): warnings.simplefilter("ignore") r_guard = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "aged_exact"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "aged_exact"], ) r_full = TripleDifference(estimation_method=method).fit( - data, outcome="outcome", group="group", partition="partition", - time="time", covariates=["age", "aged_near"], + data, + outcome="outcome", + group="group", + partition="partition", + time="time", + covariates=["age", "aged_near"], ) assert np.isfinite(r_guard.se) and r_guard.se > 0 # not 1e17 garbage np.testing.assert_allclose(r_guard.se, r_full.se, rtol=1e-3) diff --git a/tests/test_openai_review.py b/tests/test_openai_review.py index 17cb95126..dc3a86ae4 100644 --- a/tests/test_openai_review.py +++ b/tests/test_openai_review.py @@ -11,7 +11,6 @@ import pathlib import re import subprocess -import sys import pytest @@ -24,10 +23,7 @@ def _find_script() -> "pathlib.Path | None": """Find openai_review.py relative to the repo root.""" # Method 1: relative to this test file (works in local checkout) candidate = ( - pathlib.Path(__file__).resolve().parent.parent - / ".claude" - / "scripts" - / "openai_review.py" + pathlib.Path(__file__).resolve().parent.parent / ".claude" / "scripts" / "openai_review.py" ) if candidate.exists(): return candidate @@ -132,11 +128,7 @@ def test_utility_files_no_sections(self, review_mod): assert review_mod._needed_sections(text) == set() def test_mixed_files(self, review_mod): - text = ( - "M\tdiff_diff/bacon.py\n" - "M\tdiff_diff/linalg.py\n" - "M\ttests/test_bacon.py" - ) + text = "M\tdiff_diff/bacon.py\n" "M\tdiff_diff/linalg.py\n" "M\ttests/test_bacon.py" sections = review_mod._needed_sections(text) assert sections == {"BaconDecomposition"} @@ -159,9 +151,7 @@ class TestExtractRegistrySections: ) def test_extract_single_section(self, review_mod): - result = review_mod.extract_registry_sections( - self.SAMPLE_REGISTRY, {"BaconDecomposition"} - ) + result = review_mod.extract_registry_sections(self.SAMPLE_REGISTRY, {"BaconDecomposition"}) assert "Bacon content line 1" in result assert "SA content" not in result @@ -182,9 +172,7 @@ def test_empty_section_names(self, review_mod): assert review_mod.extract_registry_sections(self.SAMPLE_REGISTRY, set()) == "" def test_nonexistent_section(self, review_mod): - result = review_mod.extract_registry_sections( - self.SAMPLE_REGISTRY, {"NonExistent"} - ) + result = review_mod.extract_registry_sections(self.SAMPLE_REGISTRY, {"NonExistent"}) assert result == "" @@ -237,7 +225,7 @@ def test_strips_shell_grep_directive_from_real_prompt(self, review_mod): assert 'Command to check: `grep -n "pattern" diff_diff/*.py`' in source # After local adaptation: directive is gone, no-shell-access note is in. adapted = review_mod._adapt_review_criteria(source) - assert 'Command to check: `grep' not in adapted + assert "Command to check: `grep" not in adapted assert "no shell access" in adapted @@ -397,20 +385,26 @@ def test_fresh_review_no_delta_sections(self, review_mod): assert "Full Branch Diff (Reference Only)" not in result assert "Changes Under Review" in result - def test_findings_table_escapes_pipe_chars(self, review_mod): """Summary containing | should be escaped in the findings table.""" findings = [ { - "id": "R1-P1-1", "severity": "P1", "section": "Code Quality", + "id": "R1-P1-1", + "severity": "P1", + "section": "Code Quality", "summary": "Return type str | None is wrong", - "location": "foo.py:L10", "status": "open", + "location": "foo.py:L10", + "status": "open", } ] result = review_mod.compile_prompt( - criteria_text="C.", registry_content="R.", diff_text="D.", - changed_files_text="M\tf.py", branch_info="b", - previous_review="Prev.", delta_diff_text="delta", + criteria_text="C.", + registry_content="R.", + diff_text="D.", + changed_files_text="M\tf.py", + branch_info="b", + previous_review="Prev.", + delta_diff_text="delta", structured_findings=findings, ) # The pipe in "str | None" should be escaped as "str \| None" @@ -548,9 +542,7 @@ def test_role_attribute(self, review_mod, repo_root): assert 'role="import-context"' in result def test_handles_missing_file(self, review_mod, repo_root, capsys): - result = review_mod.read_source_files( - ["/nonexistent/path.py"], repo_root - ) + result = review_mod.read_source_files(["/nonexistent/path.py"], repo_root) assert result == "" captured = capsys.readouterr() assert "Warning" in captured.err @@ -587,12 +579,8 @@ def test_submodule_imports_not_truncated(self, review_mod, repo_root): pytest.skip("diff_diff/visualization/_staggered.py not found") imports = review_mod.parse_imports(path) # Should include full submodule paths like diff_diff.visualization._common - has_submodule = any( - m.count(".") >= 2 for m in imports # at least 3 components - ) - assert has_submodule, ( - f"Expected submodule imports (3+ components) but got: {imports}" - ) + has_submodule = any(m.count(".") >= 2 for m in imports) # at least 3 components + assert has_submodule, f"Expected submodule imports (3+ components) but got: {imports}" def test_relative_import_aliases_expanded(self, review_mod, repo_root): """from . import _event_study should resolve to diff_diff.visualization._event_study.""" @@ -602,9 +590,7 @@ def test_relative_import_aliases_expanded(self, review_mod, repo_root): imports = review_mod.parse_imports(path) # Should include individual submodule names, not just the package submodules = [m for m in imports if m.startswith("diff_diff.visualization._")] - assert len(submodules) > 0, ( - f"Expected visualization submodule imports but got: {imports}" - ) + assert len(submodules) > 0, f"Expected visualization submodule imports but got: {imports}" def test_handles_syntax_error(self, review_mod, tmp_path, capsys): test_file = tmp_path / "bad.py" @@ -655,9 +641,9 @@ def test_visualization_init_includes_submodules(self, review_mod, repo_root): result = review_mod.expand_import_graph([path], repo_root) filenames = [os.path.basename(p) for p in result] # Should include visualization submodules like _event_study.py, _staggered.py - assert any(f.startswith("_") and f.endswith(".py") for f in filenames), ( - f"Expected visualization submodule files but got: {filenames}" - ) + assert any( + f.startswith("_") and f.endswith(".py") for f in filenames + ), f"Expected visualization submodule files but got: {filenames}" def test_empty_input(self, review_mod, repo_root): assert review_mod.expand_import_graph([], repo_root) == [] @@ -809,8 +795,10 @@ def test_non_dict_findings_filtered(self, review_mod, tmp_path): """Non-dict elements in findings list are filtered out, not crash.""" state_file = tmp_path / "review-state.json" good_finding = { - "id": "R1-P1-1", "severity": "P1", - "summary": "Test finding", "status": "open", + "id": "R1-P1-1", + "severity": "P1", + "summary": "Test finding", + "status": "open", } state = { "schema_version": 1, @@ -905,10 +893,7 @@ def test_empty_review(self, review_mod): assert not uncertain def test_finding_ids_follow_format(self, review_mod): - review_text = ( - "**P0** Critical bug in `foo.py:L1`\n" - "**P1** Minor issue in the code\n" - ) + review_text = "**P0** Critical bug in `foo.py:L1`\n" "**P1** Minor issue in the code\n" findings, _ = review_mod.parse_review_findings(review_text, 2) for f in findings: assert f["id"].startswith("R2-") @@ -999,7 +984,9 @@ def test_finding_with_skip_marker_in_summary_still_parsed(self, review_mod): def test_finding_with_looks_good_in_summary(self, review_mod): """Finding mentioning 'Looks good' in summary should not be skipped.""" - review_text = "**P1** Assessment says Looks good but edge case is unhandled in `bar.py:L5`\n" + review_text = ( + "**P1** Assessment says Looks good but edge case is unhandled in `bar.py:L5`\n" + ) findings, _ = review_mod.parse_review_findings(review_text, 1) assert len(findings) == 1 assert findings[0]["severity"] == "P1" @@ -1116,24 +1103,39 @@ def test_ignores_instructional_text(self, review_mod): class TestMergeFindings: def test_matching_finding_stays_open(self, review_mod): previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard", "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard", + "status": "open", + } ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard", "status": "open"} + { + "id": "R2-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) - open_at_loc = [ - f for f in merged - if f["location"] == "foo.py:L10" and f["status"] == "open" - ] + open_at_loc = [f for f in merged if f["location"] == "foo.py:L10" and f["status"] == "open"] assert len(open_at_loc) >= 1 def test_absent_finding_marked_addressed(self, review_mod): previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard", "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard", + "status": "open", + } ] current = [] # Finding was addressed merged = review_mod.merge_findings(previous, current) @@ -1144,8 +1146,14 @@ def test_absent_finding_marked_addressed(self, review_mod): def test_new_finding_added_as_open(self, review_mod): previous = [] current = [ - {"id": "R2-P0-1", "severity": "P0", "location": "bar.py:L5", - "section": "Methodology", "summary": "Missing check", "status": "open"} + { + "id": "R2-P0-1", + "severity": "P0", + "location": "bar.py:L5", + "section": "Methodology", + "summary": "Missing check", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) assert len(merged) == 1 @@ -1155,14 +1163,24 @@ def test_new_finding_added_as_open(self, review_mod): def test_matching_with_shifted_line_numbers(self, review_mod): """Same finding at different line ranges should still match via summary.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "foo.py:L10-L12", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R2-P1-1", + "severity": "P1", + "location": "foo.py:L10-L12", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1174,14 +1192,24 @@ def test_matching_with_shifted_line_numbers(self, review_mod): def test_matching_with_missing_location(self, review_mod): """Finding with no location should still match on summary fingerprint.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R2-P1-1", + "severity": "P1", + "location": "", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1193,17 +1221,32 @@ def test_matching_with_missing_location(self, review_mod): def test_multiple_findings_same_key(self, review_mod): """Multiple previous findings with same key should not overwrite each other.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"}, - {"id": "R1-P1-2", "severity": "P1", "location": "foo.py:L20", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"}, + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + }, + { + "id": "R1-P1-2", + "severity": "P1", + "location": "foo.py:L20", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + }, ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"}, + { + "id": "R2-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + }, ] merged = review_mod.merge_findings(previous, current) # One should match, one should be addressed @@ -1215,17 +1258,32 @@ def test_multiple_findings_same_key(self, review_mod): def test_duplicate_no_location_findings_one_to_one(self, review_mod): """Two prior no-location findings should not both match one current finding.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "", - "section": "Code Quality", "summary": "Missing NaN guard", - "status": "open"}, - {"id": "R1-P1-2", "severity": "P1", "location": "", - "section": "Methodology", "summary": "Missing NaN guard", - "status": "open"}, + { + "id": "R1-P1-1", + "severity": "P1", + "location": "", + "section": "Code Quality", + "summary": "Missing NaN guard", + "status": "open", + }, + { + "id": "R1-P1-2", + "severity": "P1", + "location": "", + "section": "Methodology", + "summary": "Missing NaN guard", + "status": "open", + }, ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard", - "status": "open"}, + { + "id": "R2-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard", + "status": "open", + }, ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1237,14 +1295,24 @@ def test_duplicate_no_location_findings_one_to_one(self, review_mod): def test_previous_missing_location_current_has_location(self, review_mod): """Previous finding with no location, current has one → should match.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "staggered.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in staggered", - "status": "open"} + { + "id": "R2-P1-1", + "severity": "P1", + "location": "staggered.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in staggered", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1256,12 +1324,24 @@ def test_previous_missing_location_current_has_location(self, review_mod): def test_same_basename_different_dirs_no_cross_match(self, review_mod): """__init__.py in different dirs with same summary should NOT cross-match.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "diff_diff/__init__.py:L10", - "section": "Code Quality", "summary": "Missing type export", "status": "open"} + { + "id": "R1-P1-1", + "severity": "P1", + "location": "diff_diff/__init__.py:L10", + "section": "Code Quality", + "summary": "Missing type export", + "status": "open", + } ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "diff_diff/visualization/__init__.py:L5", - "section": "Code Quality", "summary": "Missing type export", "status": "open"} + { + "id": "R2-P1-1", + "severity": "P1", + "location": "diff_diff/visualization/__init__.py:L5", + "section": "Code Quality", + "summary": "Missing type export", + "status": "open", + } ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1274,20 +1354,40 @@ def test_long_summaries_dont_collide(self, review_mod): """Two findings with same first 50 chars but different suffixes should NOT collapse.""" prefix = "a" * 50 previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": prefix + " first issue details", - "status": "open"}, - {"id": "R1-P1-2", "severity": "P1", "location": "foo.py:L20", - "section": "Code Quality", "summary": prefix + " second different issue", - "status": "open"}, + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": prefix + " first issue details", + "status": "open", + }, + { + "id": "R1-P1-2", + "severity": "P1", + "location": "foo.py:L20", + "section": "Code Quality", + "summary": prefix + " second different issue", + "status": "open", + }, ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": prefix + " first issue details", - "status": "open"}, - {"id": "R2-P1-2", "severity": "P1", "location": "foo.py:L20", - "section": "Code Quality", "summary": prefix + " second different issue", - "status": "open"}, + { + "id": "R2-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": prefix + " first issue details", + "status": "open", + }, + { + "id": "R2-P1-2", + "severity": "P1", + "location": "foo.py:L20", + "section": "Code Quality", + "summary": prefix + " second different issue", + "status": "open", + }, ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1299,14 +1399,24 @@ def test_long_summaries_dont_collide(self, review_mod): def test_same_summary_different_files_no_cross_match(self, review_mod): """Two findings with same summary but different files should NOT cross-match.""" previous = [ - {"id": "R1-P1-1", "severity": "P1", "location": "foo.py:L10", - "section": "Code Quality", "summary": "Missing NaN guard in estimator", - "status": "open"}, + { + "id": "R1-P1-1", + "severity": "P1", + "location": "foo.py:L10", + "section": "Code Quality", + "summary": "Missing NaN guard in estimator", + "status": "open", + }, ] current = [ - {"id": "R2-P1-1", "severity": "P1", "location": "bar.py:L20", - "section": "Code Quality", "summary": "Missing NaN guard in estimator", - "status": "open"}, + { + "id": "R2-P1-1", + "severity": "P1", + "location": "bar.py:L20", + "section": "Code Quality", + "summary": "Missing NaN guard in estimator", + "status": "open", + }, ] merged = review_mod.merge_findings(previous, current) open_findings = [f for f in merged if f["status"] == "open"] @@ -1378,6 +1488,7 @@ def test_stores_and_retrieves_branch_and_base(self, review_mod, tmp_path): ) # Read back and verify fields are present import json + with open(path) as f: data = json.load(f) assert data["branch"] == "feature/test" @@ -1458,10 +1569,12 @@ class TestValidateReviewState: def test_valid_state_returns_true(self, review_mod, tmp_path): path = str(tmp_path / "review-state.json") review_mod.write_review_state( - path=path, commit_sha="abc123", base_ref="main", - branch="feature/test", review_round=1, - findings=[{"id": "R1-P1-1", "severity": "P1", - "summary": "Test", "status": "open"}], + path=path, + commit_sha="abc123", + base_ref="main", + branch="feature/test", + review_round=1, + findings=[{"id": "R1-P1-1", "severity": "P1", "summary": "Test", "status": "open"}], ) findings, rnd, commit, valid = review_mod.validate_review_state( path, "feature/test", "main" @@ -1473,26 +1586,24 @@ def test_valid_state_returns_true(self, review_mod, tmp_path): def test_branch_mismatch_returns_false(self, review_mod, tmp_path): path = str(tmp_path / "review-state.json") review_mod.write_review_state( - path=path, commit_sha="abc123", base_ref="main", - branch="feature/old", review_round=1, findings=[], - ) - _, _, _, valid = review_mod.validate_review_state( - path, "feature/new", "main" + path=path, + commit_sha="abc123", + base_ref="main", + branch="feature/old", + review_round=1, + findings=[], ) + _, _, _, valid = review_mod.validate_review_state(path, "feature/new", "main") assert not valid def test_schema_mismatch_returns_false(self, review_mod, tmp_path): state_file = tmp_path / "review-state.json" state_file.write_text(json.dumps({"schema_version": 999})) - _, _, _, valid = review_mod.validate_review_state( - str(state_file), "b", "main" - ) + _, _, _, valid = review_mod.validate_review_state(str(state_file), "b", "main") assert not valid def test_missing_file_returns_false(self, review_mod): - _, _, _, valid = review_mod.validate_review_state( - "/nonexistent.json", "b", "main" - ) + _, _, _, valid = review_mod.validate_review_state("/nonexistent.json", "b", "main") assert not valid def test_malformed_finding_returns_false(self, review_mod, tmp_path): @@ -1509,9 +1620,7 @@ def test_malformed_finding_returns_false(self, review_mod, tmp_path): ], } state_file.write_text(json.dumps(state)) - _, _, _, valid = review_mod.validate_review_state( - str(state_file), "feature/test", "main" - ) + _, _, _, valid = review_mod.validate_review_state(str(state_file), "feature/test", "main") assert not valid # fail closed on malformed finding @@ -1641,23 +1750,17 @@ class TestSanitizePreviousReview: """Hostile prior-review content must not be able to close the wrapper tag.""" def test_strips_lowercase_closing_tag(self, review_mod): - result = review_mod._sanitize_previous_review( - "hi there" - ) + result = review_mod._sanitize_previous_review("hi there") assert "" not in result assert "</previous-review-output>" in result def test_strips_uppercase_closing_tag(self, review_mod): - result = review_mod._sanitize_previous_review( - "hi there" - ) + result = review_mod._sanitize_previous_review("hi there") assert "" not in result assert "</previous-review-output>" in result def test_strips_mixed_case_with_whitespace(self, review_mod): - result = review_mod._sanitize_previous_review( - "hi there" - ) + result = review_mod._sanitize_previous_review("hi there") assert "' in text + assert r"" in text assert "" in text def test_workflow_wraps_pr_body_with_untrusted_attr(self): @@ -1733,7 +1836,7 @@ def test_workflow_wraps_pr_body_with_untrusted_attr(self): pytest.skip("workflow not found") text = wf.read_text() # Shell uses backslash-escaped quotes inside the YAML literal block. - assert r'' in text + assert r"" in text assert "" in text def test_workflow_sanitizes_pr_title_closing_tag(self): @@ -1767,7 +1870,7 @@ def test_workflow_wraps_notebook_prose_with_untrusted_attr(self): pytest.skip("workflow not found") text = wf.read_text() # Shell uses backslash-escaped quotes inside the YAML literal block. - assert r'' in text + assert r"" in text assert "" in text def test_workflow_sanitizes_notebook_prose_closing_tag(self): @@ -1833,16 +1936,11 @@ def test_workflow_bootstrap_branch_has_parity_with_steady_state(self): "did the elif transition get rewritten?" ) assert end_anchor in text, ( - f"end anchor {end_anchor!r} missing from workflow — " - "did the Codex step get renamed?" + f"end anchor {end_anchor!r} missing from workflow — " "did the Codex step get renamed?" ) - steady_state = text[ - text.index(steady_anchor) : text.index(bootstrap_anchor) - ] - bootstrap = text[ - text.index(bootstrap_anchor) : text.index(end_anchor) - ] + steady_state = text[text.index(steady_anchor) : text.index(bootstrap_anchor)] + bootstrap = text[text.index(bootstrap_anchor) : text.index(end_anchor)] # Each branch must apply close-tag sanitization independently. sanitize_re = r"" @@ -2000,12 +2098,12 @@ def test_workflow_steady_state_has_aggregate_budget_cap(self): ) assert 'NB_OMITTED+=("$nb")' in text, ( "Omitted paths must be appended via array push " - "(`NB_OMITTED+=(\"$nb\")`) with explicit double-quoting to " + '(`NB_OMITTED+=("$nb")`) with explicit double-quoting to ' "preserve literal path content." ) assert '"${NB_OMITTED[@]}"' in text, ( "Truncation marker must iterate NB_OMITTED via quoted array " - "expansion (`for omitted in \"${NB_OMITTED[@]}\"; do`) to " + 'expansion (`for omitted in "${NB_OMITTED[@]}"; do`) to ' "survive paths with whitespace." ) @@ -2436,9 +2534,9 @@ def test_both_gates_enumerate_same_triggers(self, workflow_text): 'context.payload.action !== "opened"', # JS ] for signal in rerun_signals: - assert signal in workflow_text, ( - f"Expected rerun-set signal {signal!r} not found in workflow YAML" - ) + assert ( + signal in workflow_text + ), f"Expected rerun-set signal {signal!r} not found in workflow YAML" class TestWorkflowCodexActionContract: @@ -2494,8 +2592,7 @@ def _step_block(workflow_text, step_name): elsewhere in the file (e.g. a comment). Mirrors ``TestWorkflowDoesNotExecutePRHeadCode._extract_step_block``.""" pattern = re.compile( - rf"^ - name:\s*{re.escape(step_name)}\s*\n" - r"((?:[ ]{8,}.*\n|[ ]*\n)*)", + rf"^ - name:\s*{re.escape(step_name)}\s*\n" r"((?:[ ]{8,}.*\n|[ ]*\n)*)", re.MULTILINE, ) m = pattern.search(workflow_text) @@ -2644,9 +2741,7 @@ def test_prev_review_fetch_filter_is_prefix_of_markers(self, workflow_text): every run is framed as a fresh review.""" fetch = self._require_block(workflow_text, self.FETCH_PREV_STEP) fetch_filter = "