diff --git a/CHANGELOG.md b/CHANGELOG.md index f1802804..be2c23e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -53,6 +53,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 supersedes it. ### Changed +- **CallawaySantAnna aggregation SE assembly rewritten to O(n_units)** (also inherited by + `StaggeredTripleDifference` aggregation). The combined influence-function assembly behind + simple/event-study/group aggregated SEs — previously ~56-85% of analytical fit time at + scale — replaces its per-aggregation-target full-DataFrame cohort scans, per-unit Python + loops, and dense `(n_units × n_gt)` weight-influence-function matrices with per-fit cohort + tables and a closed-form WIF (algebraically identical; details in REGISTRY). Measured + (medians of 3, Apple M4 Max, Rust backend): analytical fits at 100k units × 20 periods + (2M rows) 1.32→0.21s no-covariate (6.3x), 2.49→1.08s 5-covariate DR (2.3x); long panels + 40p×10 cohorts 9.09→3.91s and 60p×15 cohorts 9.86→4.34s (2.3x); bootstrap-999 fits + 1.4-1.6x (the remaining draw-loop matmul is unchanged); repeated cross-sections + 4.17→1.50s no-covariate (2.8x) with peak memory 6.4→2.1 GB (-67%; the dense WIF matrices + were observation-scale in RCS mode). Point estimates are bit-identical; aggregated SEs + move only at floating-point reassociation level (measured ≤5e-16 relative, drift-bound + tested at 1e-9 against a frozen copy of the prior implementation). For scale context, + a full `fit(aggregate="all")` at 2M rows now runs 3-11x faster than R `did` 2.5.1 + (equal work, analytical or bootstrap-999, single- or multi-core R; R's `pl`/`cores` + parallelism is BLAS-bound on this path so both R arms time identically). - **ImputationDiD/TwoStageDiD demeaning modernized onto the shared MAP engine.** The private per-estimator pandas `_iterative_demean` loops (rebuilt a `pd.Series.groupby().transform()` hash table every alternating-projection iteration) diff --git a/TODO.md b/TODO.md index 75106fca..12c6528b 100644 --- a/TODO.md +++ b/TODO.md @@ -32,6 +32,7 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo | `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; remaining: (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. Also deferred from the covariate work: `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate — documented `NotImplementedError`), and `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF). | `continuous_did.py` | CGBS-2024 | Heavy | Low | | `ImputationDiD` LOO conservative-variance refinement (BJS 2024 Supp. Appendix A.9) — a finite-sample improvement to the auxiliary-model residuals reducing overfit of `tau_tilde_g` to `epsilon`. Asymptotic Theorem-3 variance is implemented and matches R `didimputation` (which also omits LOO by default). | `imputation.py` | imputation-validation | Mid | Low | | `TwoWayFixedEffects(vcov_type in {hc2, hc2_bm})` with replicate-weight designs raises `NotImplementedError` (`twfe.py:~233`). The replicate path re-demeans per replicate, which doesn't compose with the full-dummy HC2/HC2-BM build — a correct impl needs per-replicate full-dummy refit. Workaround: `hc1` for replicate-weight CR1. | `twfe.py::fit` | follow-up | Heavy | Low | +| `CallawaySantAnna` reg-method aggregated SEs sit 3-20% from the R golden fixtures (`two_period` simple/group/dynamic, `dynamic_effects` dynamic) while the ATTs match at 1e-11 and the dr-method aggregated SEs match at ≤7e-6. Pre-existing (byte-identical deltas on the pre-fast-path tree); golden aggregated-SE assertions were therefore enabled for dr scenarios only (`test_golden_simple_aggregation_se`, `test_golden_event_study_aggregation_se`). Investigate the reg-path per-cell IF vs R `DRDID::reg_did_panel`'s. | `staggered.py::_outcome_regression`, `tests/test_csdid_ported.py` | CS-scaling | Mid | Medium | | TWFE's HC2/HC2-BM inline full-dummy build (`twfe.py:280-315`) duplicates the dummy-construction logic in `DifferenceInDifferences(fixed_effects=...)` (`estimators.py:478-486`). Extract a shared helper, or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch (with TWFE-specific cluster-default threading), to reduce drift risk on FE naming / survey behavior / result-surface conventions. Substantive refactor — touches both estimators. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit` | follow-up | Heavy | Low | ### Performance @@ -53,6 +54,9 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low | | Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low | | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low | +| `_cluster_robust_se_from_per_gt_if` scatters per-cell IFs with `np.add.at` (unbuffered ufunc, 5-20x slower than fancy `+=`; same slow-scatter class the aggregation fast path fixed). Index arrays come from `np.where` over boolean masks (duplicate-free), so per-cell `psi[idx] += vals` is exact. Runs once per (g,t) cell when `cluster=` is set — noticeable at many-cell fits. | `staggered.py::_cluster_robust_se_from_per_gt_if` | CS-scaling | Quick | Low | +| Per-cell `treated_units`/`control_units` label arrays (`all_units[positions]`, ~O(n_control) alloc per (g,t) cell) are consumed only by the precomputed-None fallback of the combined-IF assembly, which no in-package caller reaches — build them lazily (or drop from the IF-info dict) to cut per-cell allocation at high cell counts. | `staggered.py::_compute_att_gt_fast` | CS-scaling | Mid | Low | +| `_compute_aggregated_se` is dead code (zero in-package callers; superseded by `_compute_aggregated_se_with_wif`) — remove it, or fold its docstring into the WIF variant. Its `np.add.at` scatter also predates the fancy-`+=` convention. | `staggered_aggregation.py::_compute_aggregated_se` | CS-scaling | Quick | Low | ### Testing / docs diff --git a/benchmarks/R/benchmark_did.R b/benchmarks/R/benchmark_did.R index 3d85cc94..210d98b2 100644 --- a/benchmarks/R/benchmark_did.R +++ b/benchmarks/R/benchmark_did.R @@ -2,7 +2,17 @@ # Benchmark: Callaway-Sant'Anna Estimator (R `did` package) # # Usage: -# Rscript benchmark_did.R --data path/to/data.csv --output path/to/results.json +# Rscript benchmark_did.R --data path/to/data.csv --output path/to/results.json \ +# [--method dr|ipw|reg] [--control-group nevertreated|notyettreated] \ +# [--xformla "~ x1 + x2"] [--bstrap true|false] [--biters N] \ +# [--cband true|false] [--pl true|false] [--cores N] \ +# [--faster-mode true|false] +# +# Defaults reproduce the historical behavior (analytical SEs, no covariates, +# single core), so existing callers (benchmarks/run_benchmarks.py) are +# unaffected. The optional flags exist for the R-yardstick arms: bootstrap +# inference (bstrap/biters/cband applied at BOTH att_gt and aggte), covariate +# formulas, and did's parallel processing (pl/cores). library(did) library(jsonlite) @@ -11,35 +21,73 @@ library(data.table) # Parse command line arguments args <- commandArgs(trailingOnly = TRUE) +parse_bool <- function(x, flag) { + v <- tolower(x) + if (v %in% c("true", "t", "1", "yes")) return(TRUE) + if (v %in% c("false", "f", "0", "no")) return(FALSE) + stop(sprintf("Invalid boolean for %s: '%s' (use true/false)", flag, x)) +} + parse_args <- function(args) { result <- list( data = NULL, output = NULL, method = "dr", - control_group = "nevertreated" + control_group = "nevertreated", + xformla = NULL, + bstrap = FALSE, + biters = 1000L, + cband = FALSE, + pl = FALSE, + cores = 1L, + faster_mode = NULL # NULL -> use did's own default for this version ) i <- 1 while (i <= length(args)) { - if (args[i] == "--data") { - result$data <- args[i + 1] - i <- i + 2 - } else if (args[i] == "--output") { - result$output <- args[i + 1] - i <- i + 2 - } else if (args[i] == "--method") { - result$method <- args[i + 1] - i <- i + 2 - } else if (args[i] == "--control-group") { - result$control_group <- args[i + 1] - i <- i + 2 + flag <- args[i] + if (i + 1 > length(args) && flag != "") { + stop(sprintf("Missing value for flag: %s", flag)) + } + val <- args[i + 1] + if (flag == "--data") { + result$data <- val + } else if (flag == "--output") { + result$output <- val + } else if (flag == "--method") { + result$method <- val + } else if (flag == "--control-group") { + result$control_group <- val + } else if (flag == "--xformla") { + result$xformla <- as.formula(val) + } else if (flag == "--bstrap") { + result$bstrap <- parse_bool(val, flag) + } else if (flag == "--biters") { + result$biters <- as.integer(val) + } else if (flag == "--cband") { + result$cband <- parse_bool(val, flag) + } else if (flag == "--pl") { + result$pl <- parse_bool(val, flag) + } else if (flag == "--cores") { + result$cores <- as.integer(val) + } else if (flag == "--faster-mode") { + result$faster_mode <- parse_bool(val, flag) } else { - i <- i + 1 + # Unknown flags used to be silently skipped, which turned typos into + # silent default runs. Fail loudly instead. + stop(sprintf("Unknown flag: %s", flag)) } + i <- i + 2 } if (is.null(result$data) || is.null(result$output)) { - stop("Usage: Rscript benchmark_did.R --data --output [--method dr|ipw|reg] [--control-group nevertreated|notyettreated]") + stop("Usage: Rscript benchmark_did.R --data --output [--method dr|ipw|reg] [--control-group nevertreated|notyettreated] [--xformla '~ x1 + x2'] [--bstrap true|false] [--biters N] [--cband true|false] [--pl true|false] [--cores N] [--faster-mode true|false]") + } + if (is.na(result$biters) || result$biters <= 0) { + stop(sprintf("--biters must be a positive integer, got '%s'", result$biters)) + } + if (is.na(result$cores) || result$cores <= 0) { + stop(sprintf("--cores must be a positive integer, got '%s'", result$cores)) } return(result) @@ -66,28 +114,50 @@ message(sprintf("Never-treated units (first_treat=Inf): %d", sum(is.infinite(dat message("Running Callaway-Sant'Anna estimation...") start_time <- Sys.time() -out <- att_gt( +att_gt_args <- list( yname = "outcome", tname = "time", idname = "unit", gname = "first_treat", - xformla = NULL, + xformla = config$xformla, data = data, est_method = config$method, control_group = config$control_group, - bstrap = FALSE, # Use analytical SEs for speed - cband = FALSE + bstrap = config$bstrap, + biters = config$biters, + cband = config$cband, + pl = config$pl, + cores = config$cores ) +# faster_mode exists in recent did releases only; pass it only when both +# requested and supported, so the script still runs on older installs. +# effective_faster_mode records what actually reached att_gt (vs the +# requested flag) so benchmark artifacts are never mislabeled. +effective_faster_mode <- "did-default" +if (!is.null(config$faster_mode)) { + if ("faster_mode" %in% names(formals(att_gt))) { + att_gt_args$faster_mode <- config$faster_mode + effective_faster_mode <- config$faster_mode + } else { + message("faster_mode not supported by this did version; ignoring flag") + effective_faster_mode <- "unsupported-ignored" + } +} +out <- do.call(att_gt, att_gt_args) estimation_time <- as.numeric(difftime(Sys.time(), start_time, units = "secs")) -# Aggregate results +# Aggregate results (same inference mode as att_gt so the timing arm does +# equal work end-to-end) message("Aggregating results...") agg_start <- Sys.time() -agg_simple <- aggte(out, type = "simple", bstrap = FALSE, cband = FALSE) -agg_dynamic <- aggte(out, type = "dynamic", bstrap = FALSE, cband = FALSE) -agg_group <- aggte(out, type = "group", bstrap = FALSE, cband = FALSE) +agg_simple <- aggte(out, type = "simple", bstrap = config$bstrap, + biters = config$biters, cband = FALSE) +agg_dynamic <- aggte(out, type = "dynamic", bstrap = config$bstrap, + biters = config$biters, cband = config$cband) +agg_group <- aggte(out, type = "group", bstrap = config$bstrap, + biters = config$biters, cband = FALSE) aggregation_time <- as.numeric(difftime(Sys.time(), agg_start, units = "secs")) total_time <- estimation_time + aggregation_time @@ -131,13 +201,22 @@ results <- list( total_seconds = total_time ), - # Metadata + # Metadata (records the full inference/parallelism config so every + # yardstick number is reproducible from its own artifact) metadata = list( r_version = R.version.string, did_version = as.character(packageVersion("did")), n_units = length(unique(data$unit)), n_periods = length(unique(data$time)), - n_obs = nrow(data) + n_obs = nrow(data), + xformla = if (is.null(config$xformla)) NULL else deparse(config$xformla), + bstrap = config$bstrap, + biters = if (config$bstrap) config$biters else NULL, + cband = config$cband, + pl = config$pl, + cores = config$cores, + faster_mode = effective_faster_mode, + blas = tryCatch(sessionInfo()$BLAS, error = function(e) NULL) ) ) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index cb256354..d739d8cc 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1008,6 +1008,10 @@ def _compute_att_gt_fast( # Package influence function info with index arrays (positions into # precomputed['all_units']) for O(1) downstream lookups instead of # O(n) Python dict lookups. + # INVARIANT: treated_idx/control_idx are np.where over boolean masks, + # so each is duplicate-free - the aggregation fast path relies on this + # to scatter with fancy `psi[idx] += vals` (see + # staggered_aggregation._combined_if_fast). n_t = int(n_treated) all_units = precomputed["all_units"] treated_positions = np.where(treated_valid)[0] @@ -1576,6 +1580,9 @@ def _compute_all_att_gt_covariate_reg( "skip_reason": None, } + # INVARIANT: np.where over boolean masks -> duplicate-free + # index arrays (fancy-+= scatter contract, see + # staggered_aggregation._combined_if_fast). all_units = precomputed["all_units"] treated_positions = np.where(treated_valid)[0] control_positions = np.where(control_valid)[0] @@ -3511,7 +3518,11 @@ def _compute_att_gt_rc( ) # Build influence function info - # For RCS, treated_idx/control_idx combine obs from BOTH periods + # For RCS, treated_idx/control_idx combine obs from BOTH periods. + # INVARIANT: the two period masks are disjoint (obs_time == t vs + # == base period), so the concatenated index arrays stay + # duplicate-free (fancy-+= scatter contract, see + # staggered_aggregation._combined_if_fast). treated_idx = np.concatenate([np.where(treated_t)[0], np.where(treated_s)[0]]) control_idx = np.concatenate([np.where(control_t)[0], np.where(control_s)[0]]) diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 567a64bd..b52a7c98 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -212,6 +212,178 @@ def _compute_aggregated_se( variance = np.sum(psi_overall**2) return np.sqrt(variance) + @staticmethod + def _get_agg_cache(precomputed: "PrecomputedData") -> Dict[str, Any]: + """ + Per-fit cohort tables for the combined-IF fast path, lazily memoized + on the precomputed dict. + + The cache is validated by ARRAY IDENTITY, not dict residency: + StaggeredTripleDifference aggregates through a shallow copy of + precomputed with a replaced (eligibility-zeroed) ``unit_cohorts``, + so a cache keyed to the dict could serve stale tables across the + copy. ``cohorts_ref``/``sw_ref`` pin the exact arrays the tables + were built from; any mismatch rebuilds into a FRESH dict (never + mutated in place - the shallow copy shares the cache reference). + """ + unit_cohorts = precomputed["unit_cohorts"] + survey_w = precomputed.get("survey_weights") + cache = precomputed.get("_agg_cache") + if ( + cache is not None + and cache["cohorts_ref"] is unit_cohorts + and cache["sw_ref"] is survey_w + ): + return cache + + cohort_values, cohort_codes = np.unique(unit_cohorts, return_inverse=True) + if survey_w is not None: + # Survey-weighted cohort masses. np.bincount accumulation order + # differs from the historical per-group mask-sums at the ~1 ULP + # level (documented drift budget; REGISTRY CallawaySantAnna SE + # notes). + cohort_masses = np.bincount( + cohort_codes, weights=survey_w, minlength=len(cohort_values) + ) + total_weight = float(np.sum(survey_w)) + else: + cohort_masses = np.bincount(cohort_codes, minlength=len(cohort_values)).astype( + np.float64 + ) + total_weight = float(len(unit_cohorts)) + + cache = { + "cohorts_ref": unit_cohorts, + "sw_ref": survey_w, + "cohort_values": cohort_values, + "cohort_codes": cohort_codes, + "cohort_masses": cohort_masses, + "total_weight": total_weight, + } + precomputed["_agg_cache"] = cache + return cache + + def _combined_if_fast( + self, + gt_pairs: List[Tuple[Any, Any]], + weights: np.ndarray, + effects: np.ndarray, + groups_for_gt: np.ndarray, + influence_func_info: Dict, + precomputed: "PrecomputedData", + n_units: int, + ) -> Optional[Tuple[np.ndarray, None]]: + """ + O(n_units) combined-IF assembly over per-fit cohort tables. + + Replaces the general path's per-group full-DataFrame scans, per-unit + Python loops, and dense (n_units x n_gt) WIF matrices with cohort- + indexed lookups. The WIF uses the closed form (algebraically + identical to the dense ``wif_matrix @ effects``, floating-point + accumulation order differs - not bit-for-bit): + + wif_i = w_i * (E(c_i)/S - K(c_i) * d / S**2) + + where c_i is unit i's cohort, E(c) sums ``effects`` over keeper + (g,t) pairs with g == c, K(c) counts them, S = sum of keeper pg, + d = pg_keepers @ effects, and w_i is the survey weight (1 when + unweighted). Units whose cohort is not among the keepers get + exactly 0 (the old dense form realizes the same value through + cancelling terms). + + Returns None when the cohort lookup cannot be resolved exactly + (non-numeric cohort dtypes, or a keeper group missing from the + cohort table) - the caller then falls back to the general path. + """ + cache = self._get_agg_cache(precomputed) + cohort_values = cache["cohort_values"] + cohort_codes = cache["cohort_codes"] + cohort_masses = cache["cohort_masses"] + total_weight = cache["total_weight"] + survey_w = cache["sw_ref"] + + groups_arr = np.asarray(groups_for_gt) + if not ( + np.issubdtype(groups_arr.dtype, np.number) + and np.issubdtype(cohort_values.dtype, np.number) + ): + return None + + # Unique keeper groups + exact positions in the cohort table. + unique_groups = np.unique(groups_arr) + pos = np.searchsorted(cohort_values, unique_groups) + if np.any(pos >= len(cohort_values)) or np.any( + cohort_values[np.minimum(pos, len(cohort_values) - 1)] != unique_groups + ): + return None # keeper group absent from cohort table + + # pg per keeper (same values as the general path's group_sizes / + # total_weight; survey masses differ only in accumulation order). + pg_by_group = cohort_masses[pos] / total_weight + kpos = np.searchsorted(unique_groups, groups_arr) + pg_keepers = pg_by_group[kpos] + sum_pg_keepers = np.sum(pg_keepers) + + # Guard against zero weights (no keepers = no variance). Must stay + # BEFORE the psi_standard scatter - the general path returns zeros + # without ever accumulating the standard IF. + if sum_pg_keepers == 0: + return np.zeros(n_units), None + + # Standard aggregated influence (without wif). Index arrays are + # unique within each cell by construction at every producer + # (np.where on disjoint masks), so fancy += is exact. + psi_standard = np.zeros(n_units) + for j, (g, t) in enumerate(gt_pairs): + if (g, t) not in influence_func_info: + continue + info = influence_func_info[(g, t)] + w = weights[j] + treated_idx = info["treated_idx"] + if len(treated_idx) > 0: + psi_standard[treated_idx] += w * info["treated_inf"] + control_idx = info["control_idx"] + if len(control_idx) > 0: + psi_standard[control_idx] += w * info["control_inf"] + + # Closed-form WIF over per-cohort tables. + n_ug = len(unique_groups) + E_keepers = np.bincount(kpos, weights=effects, minlength=n_ug) + K_keepers = np.bincount(kpos, minlength=n_ug).astype(np.float64) + E_full = np.zeros(len(cohort_values)) + K_full = np.zeros(len(cohort_values)) + E_full[pos] = E_keepers + K_full[pos] = K_keepers + + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + d = pg_keepers @ effects + wif_contrib = E_full[cohort_codes] / sum_pg_keepers - K_full[cohort_codes] * ( + d / sum_pg_keepers**2 + ) + if survey_w is not None: + wif_contrib = wif_contrib * survey_w + + # Check for non-finite values from edge cases (same fail-closed + # contract as the general path: warn + all-NaN vector, before the + # 1/total_weight scaling). + if not np.all(np.isfinite(wif_contrib)): + import warnings + + n_nonfinite = np.sum(~np.isfinite(wif_contrib)) + warnings.warn( + f"Non-finite values ({n_nonfinite}/{len(wif_contrib)}) in weight influence " + "function computation. This may occur with very small samples or extreme " + "weights. Returning NaN for SE to signal invalid inference.", + RuntimeWarning, + stacklevel=2, + ) + return np.full(n_units, np.nan), None + + # Scale by 1/total_weight to match R's getSE formula + psi_wif = wif_contrib / total_weight + + return psi_standard + psi_wif, None + def _compute_combined_influence_function( self, gt_pairs: List[Tuple[Any, Any]], @@ -248,6 +420,37 @@ def _compute_combined_influence_function( # Detect RCS mode via explicit flag. In RCS, obs indices ARE array positions. _is_rcs = precomputed is not None and not precomputed.get("is_panel", True) + # Fast-path dispatch: all in-package callers thread the SAME + # precomputed structures they index psi by, so the cohort tables can + # be looked up in O(n_units) instead of re-scanning the DataFrame per + # group and looping units in Python. Order matters: the RCS check + # must precede the panel identity check (for RCS both + # global_unit_to_idx and precomputed["unit_to_idx"] are None, so the + # identity guard alone would spuriously pass). Anything not exactly + # matched (direct callers with foreign index maps, size mismatches, + # non-numeric cohorts) falls through to the general path below, + # which is preserved unchanged. + if precomputed is not None and n_global_units is not None: + _fast_ok = False + if _is_rcs: + _fast_ok = n_global_units == len(precomputed["unit_cohorts"]) + elif global_unit_to_idx is not None and global_unit_to_idx is precomputed.get( + "unit_to_idx" + ): + _fast_ok = n_global_units == len(precomputed["unit_cohorts"]) + if _fast_ok: + fast = self._combined_if_fast( + gt_pairs, + weights, + effects, + groups_for_gt, + influence_func_info, + precomputed, + n_global_units, + ) + if fast is not None: + return fast + # Build unit index mapping (local or global) if _is_rcs and n_global_units is not None: # RCS: direct indexing — obs indices are the array positions diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index 757d06d8..18e57ab9 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -482,6 +482,9 @@ def fit( inf_gmm_rescaled = inf_gmm * (n_units / size_gt) inf_gmm_scaled = inf_gmm_rescaled / n_units + # INVARIANT: np.where over boolean masks -> duplicate-free + # index arrays (fancy-+= scatter contract, see + # staggered_aggregation._combined_if_fast). treated_idx = np.where(treated_mask)[0] treated_inf = inf_gmm_scaled[treated_idx] nonzero_mask = (inf_gmm_scaled != 0) & ~treated_mask diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index f0ab536c..29ae7e3b 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -505,6 +505,22 @@ Aggregations: - Bootstrap: Multiplier bootstrap with Rademacher, Mammen, or Webb weights. Bootstrap perturbs the combined influence function (standard IF + WIF) directly, not just fixed-weight re-aggregation. This correctly propagates weight estimation uncertainty. +- **Note:** The combined-IF assembly (`_compute_combined_influence_function`) routes + in-package callers through an O(n_units) fast path (v3.7): per-fit cohort tables + (unit counts or survey-weighted masses via `np.unique` + `np.bincount`, cached on the + precomputed structures with array-identity validation) and a closed-form WIF + `wif_i = w_i * (E(c_i)/S - K(c_i) * d / S**2)` — algebraically identical to the prior + dense `(n_units x n_gt)` `wif_matrix @ effects` (E(c)/K(c) = per-cohort effect + sums/keeper counts, S = sum of keeper pg, d = pg @ effects, w_i = survey weight or 1). + Point estimates are unaffected (bit-identical); aggregated SEs agree with the dense + form only to floating-point reassociation (measured ≤ 5e-16 relative; drift-bound + frozen-copy tests pin ≤ 1e-9, `tests/test_staggered_aggregation.py`), not + bit-for-bit — same accumulation-order posture as the shared demeaning engine + (see "Absorbed Fixed Effects with Survey Weights"). Units whose cohort is not + among the keeper (g,t) groups get an exact 0 WIF contribution (the dense form + realized the same value through cancelling terms). The pre-rewrite general path + is preserved verbatim as the fallback for direct callers with foreign index maps, + non-numeric cohort dtypes, or absent precomputed structures. - Block structure preserves within-unit correlation - Simultaneous confidence bands (`cband=True`, default): Uses sup-t bootstrap to compute a uniform critical value across event times, controlling family-wise error rate. Matches diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 048b63aa..ad8eb7fb 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -129,6 +129,60 @@ its non-convergence warning as evidence of the `max_iter=100` contract issue. --- +## CallawaySantAnna aggregation IF-assembly rewrite (v3.7, 2026-07) + +Profiling at 2M+ rows showed `staggered_aggregation._compute_combined_influence_function` +was 56-85% of every analytical CS fit: called once per aggregation target (38 targets at +20 periods, 128 at 60 periods; plus 1 + N_e bootstrap-prep calls), each call re-ran +per-group full-DataFrame scans (`df[df["first_treat"] == g][unit].nunique()` — cost scales +with row count × frame width × cohorts × targets: 105 scans/fit at 5 cohorts, 975 at 15), +per-unit Python dict-lookup loops, and dense `(n_units × n_gt)` WIF matrices (~76MB × 3 +arrays per call at 100k units). The rewrite: per-fit cohort tables (`np.unique` + +`np.bincount`, identity-validated cache) + closed-form WIF + fancy-index scatter; general +path preserved as fallback. Point estimates bit-identical; aggregated SEs ≤5e-16 relative. + +Measured (medians of 3, M4 Max, Rust backend, otherwise-idle machine; BEFORE = pristine +`origin/main` worktree, AFTER = frozen branch): + +| Scenario | Before | After | Speedup | Peak RSS | +|---|---:|---:|---:|---| +| no-cov analytical, 100k units × 20p (2M rows) | 1.32s | 0.21s | **6.3x** | 960→582 MB | +| no-cov analytical, 250k units (5M rows) | 3.57s | 0.60s | 5.9x | 2.1→1.3 GB | +| 5-cov reg / ipw / dr (2M rows) | 1.96/2.15/2.49s | 0.56/0.73/1.08s | 3.5/2.9/2.3x | ~-20% | +| 10/20/40-cov dr (2M rows) | 3.3/4.9/9.3s | 1.5/2.6/5.7s | 2.1/1.9/1.6x | 4.1→3.6 GB @40 | +| 40 periods × 10 cohorts (4M rows) | 9.09s | 3.91s | 2.3x | 2.6→2.1 GB | +| 60 periods × 15 cohorts (3M rows) | 9.86s | 4.34s | 2.3x | 2.5→1.9 GB | +| bootstrap-999 legs (2-4M rows) | 7.2-33.6s | 4.9-24.2s | 1.4-1.6x | ~-10-25% | +| RCS no-cov / dr / dr+boot (2M obs) | 4.2/7.3/44.0s | 1.5/4.4/39.5s | 2.8/1.7/1.1x | **6.4→2.1 / 7.0→3.2 GB** | + +The RCS memory drop is the headline tail-risk fix: in RCS mode the WIF matrices were +observation-scale. Remaining CS bottlenecks after this rewrite (honest forecast, next +targets): `_precompute_structures` pandas groupby+pivot (now #1 on analytical no-cov +fits), the bootstrap draw-loop matmul (Phase 2 candidate; 999 × 100k × 390 cells), +and per-cell DR/IPW nuisance solvers at high covariate counts (IRLS `lstsq`, rank-guard +QR — Phase 3 candidate). + +**R `did` yardstick** (did 2.5.1 — the 2.5.0 performance release with `faster_mode` and +the tiled-contraction bootstrap — R 4.5.2 + OpenBLAS 0.3.33; equal work: `att_gt` + +`aggte` simple+dynamic+group vs `fit(aggregate="all")`; identical `biters`/`cband`; +medians of 3; run configs recorded in each JSON artifact's metadata): + +| Scenario | diff-diff | R 1-core | R 14-core (`pl=TRUE, cores=14`) | +|---|---:|---:|---:| +| no-cov analytical (2M rows) | 0.21s | 2.29s | 2.29s | +| 5-cov dr analytical | 1.09s | 3.67s | 3.74s | +| no-cov bootstrap 999 | 5.11s | 16.33s | 16.14s | +| 5-cov dr bootstrap 999 | 6.13s | 18.39s | 17.58s | +| 40p × 10c, 5-cov (4M rows) | 3.88s | 13.76s | 13.72s | + +3-11x vs R at equal work under both R threading modes. The flat R multi-core column is +not a rigged comparison: did 2.5's bootstrap is a tiled BLAS contraction that already +multithreads through OpenBLAS, so fork-level `pl`/`cores` parallelism adds nothing on +this path — both arms are effectively multithreaded R. ATTs match R exactly at displayed +precision; analytical SEs to 5 significant digits; bootstrap SEs within ~1% (Monte Carlo). + +--- + ## Memory scaling at the millions-of-units tail (v3.x, June 2026) At the scale where the dense working arrays - not compute - are the binding constraint diff --git a/tests/test_csdid_ported.py b/tests/test_csdid_ported.py index 52cf9d83..36fc2c3f 100644 --- a/tests/test_csdid_ported.py +++ b/tests/test_csdid_ported.py @@ -1254,3 +1254,82 @@ def test_golden_zero_pretreatment(self, golden_values): assert abs(py_att - r_att) < 0.05, ( f"ZeroPre {bp} ATT(g={g},t={t}): " f"Py={py_att:.4f}, R={r_att:.4f}" ) + + def test_golden_simple_aggregation_se(self, golden_values): + """Simple (overall) aggregated ATT and SE match R ``aggte(type="simple")``. + + Enables the previously-unasserted ``simple`` block of the + ``default_no_covariates_dr`` golden scenario (analytical SEs, + ``bstrap=FALSE``). Observed agreement is ~1e-12; the 2e-3 tolerance + matches the per-cell DR SE precedent above. + """ + if "default_no_covariates_dr" not in golden_values: + pytest.skip("Scenario not in golden values") + scenario = golden_values["default_no_covariates_dr"] + data = _golden_to_df(scenario["data"]) + + cs = CallawaySantAnna(estimation_method="dr") + results = cs.fit( + data, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="simple", + ) + r_simple = scenario["results"]["simple"] + assert ( + abs(results.att - r_simple["att"]) < 2e-3 + ), f"Simple ATT: Py={results.att:.6f}, R={r_simple['att']:.6f}" + assert ( + abs(results.se - r_simple["se"]) < 2e-3 + ), f"Simple SE: Py={results.se:.6f}, R={r_simple['se']:.6f}" + + def test_golden_event_study_aggregation_se(self, golden_values): + """Event-study aggregated ATT and SE match R ``aggte(type="dynamic")`` + under BOTH base-period modes. + + Enables the previously-unasserted ``varying_dynamic`` / + ``universal_dynamic`` blocks of the ``varying_vs_universal`` golden + scenario (analytical SEs). R reports NA at the universal reference + period e=-1 (identically-zero estimate); NA entries are skipped - + this pins parity on the estimated horizons (observed agreement + ~7e-6). + + NOT enabled (pre-existing, unrelated to the IF-assembly fast path - + verified byte-identical deltas on the pre-rewrite tree): the + ``two_period`` and ``dynamic_effects`` aggregated SEs, where the reg + method's aggregated SEs sit 3-20% from these R fixtures while the + ATTs match at 1e-11. Tracked in TODO.md. + """ + if "varying_vs_universal" not in golden_values: + pytest.skip("Scenario not in golden values") + scenario = golden_values["varying_vs_universal"] + data = _golden_to_df(scenario["data"]) + + for bp in ["varying", "universal"]: + cs = CallawaySantAnna(estimation_method="dr", base_period=bp) + results = cs.fit( + data, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + covariates=["X"], + aggregate="event_study", + ) + blk = scenario["results"][f"{bp}_dynamic"] + n_checked = 0 + for e, r_att, r_se in zip(blk["egt"], blk["att"], blk["se"]): + if r_att == "NA" or r_se == "NA": + continue + es = results.event_study_effects.get(int(e)) + assert es is not None, f"{bp}: missing event time e={e}" + assert ( + abs(es["effect"] - float(r_att)) < 2e-3 + ), f"{bp} dyn ATT(e={e}): Py={es['effect']:.6f}, R={r_att}" + assert ( + abs(es["se"] - float(r_se)) < 2e-3 + ), f"{bp} dyn SE(e={e}): Py={es['se']:.6f}, R={r_se}" + n_checked += 1 + assert n_checked >= 8, f"{bp}: too few comparable horizons" diff --git a/tests/test_staggered_aggregation.py b/tests/test_staggered_aggregation.py new file mode 100644 index 00000000..a5c96230 --- /dev/null +++ b/tests/test_staggered_aggregation.py @@ -0,0 +1,630 @@ +""" +Frozen-copy parity and contract tests for the CS aggregation IF fast path. + +The O(n_units) rewrite of ``_compute_combined_influence_function`` (per-fit +cohort tables + closed-form WIF) replaces the dense (n_units x n_gt) WIF +matrices, the per-group full-DataFrame scans, and the per-unit Python loops +of the pre-rewrite implementation. ``_frozen_combined_influence_function`` +below is a byte-copy of that pre-rewrite implementation (self removed), kept +as a drift-bound guard: the fast path must agree with it at rtol=0, +atol=1e-9. The closed form is algebraically identical to the dense +``wif_matrix @ effects``; only floating-point accumulation order differs +(bincount cohort masses vs per-group mask-sums, table-gather vs matmul), so +observed drift is ~1e-16 - the 1e-9 bound mirrors the frozen-copy convention +in tests/test_utils.py. + +One intentional value difference: for units whose cohort is not among the +keeper (g,t) groups, the old dense form realizes an exact algebraic zero +through cancelling terms (~1e-16 residue in floating point) while the closed +form returns exactly 0.0 - parity therefore holds at atol, not bitwise. +""" + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import CallawaySantAnna, StaggeredTripleDifference, SurveyDesign +from diff_diff.staggered_aggregation import CallawaySantAnnaAggregationMixin + + +def _frozen_combined_influence_function( + gt_pairs, + weights, + effects, + groups_for_gt, + influence_func_info, + df, + unit, + precomputed=None, + global_unit_to_idx=None, + n_global_units=None, +): + """Byte-copy of the pre-rewrite ``_compute_combined_influence_function`` + (dense-WIF implementation, ``self`` removed - the body never used it).""" + if not influence_func_info: + if n_global_units is not None: + return np.zeros(n_global_units), None + return np.zeros(0), None + + # Detect RCS mode via explicit flag. In RCS, obs indices ARE array positions. + _is_rcs = precomputed is not None and not precomputed.get("is_panel", True) + + # Build unit index mapping (local or global) + if _is_rcs and n_global_units is not None: + # RCS: direct indexing — obs indices are the array positions + n_units = n_global_units + all_units = None + elif global_unit_to_idx is not None and n_global_units is not None: + n_units = n_global_units + all_units = None # caller already has the unit list + else: + all_units_set = set() + for g, t in gt_pairs: + if (g, t) in influence_func_info: + info = influence_func_info[(g, t)] + all_units_set.update(info["treated_units"]) + all_units_set.update(info["control_units"]) + + if not all_units_set: + return np.zeros(0), [] + + all_units = sorted(all_units_set) + n_units = len(all_units) + + # Get unique groups and their information + unique_groups = sorted(set(groups_for_gt)) + unique_groups_set = set(unique_groups) + group_to_idx = {g: i for i, g in enumerate(unique_groups)} + + # Check for survey weights in precomputed data + survey_w = precomputed.get("survey_weights") if precomputed is not None else None + + # Compute group-level probabilities matching R's formula + group_sizes = {} + if survey_w is not None: + precomputed_cohorts = precomputed["unit_cohorts"] + for g in unique_groups: + mask_g = precomputed_cohorts == g + group_sizes[g] = float(np.sum(survey_w[mask_g])) + total_weight = float(np.sum(survey_w)) + elif _is_rcs: + precomputed_cohorts = precomputed["unit_cohorts"] + for g in unique_groups: + group_sizes[g] = int(np.sum(precomputed_cohorts == g)) + total_weight = float(n_units) + else: + for g in unique_groups: + treated_in_g = df[df["first_treat"] == g][unit].nunique() + group_sizes[g] = treated_in_g + total_weight = float(n_units) + + pg_by_group = np.array([group_sizes[g] / total_weight for g in unique_groups]) + pg_keepers = np.array([pg_by_group[group_to_idx[g]] for g in groups_for_gt]) + sum_pg_keepers = np.sum(pg_keepers) + + if sum_pg_keepers == 0: + return np.zeros(n_units), all_units + + psi_standard = np.zeros(n_units) + + for j, (g, t) in enumerate(gt_pairs): + if (g, t) not in influence_func_info: + continue + + info = influence_func_info[(g, t)] + w = weights[j] + + treated_idx = info["treated_idx"] + if len(treated_idx) > 0: + np.add.at(psi_standard, treated_idx, w * info["treated_inf"]) + + control_idx = info["control_idx"] + if len(control_idx) > 0: + np.add.at(psi_standard, control_idx, w * info["control_inf"]) + + unit_groups_array = np.full(n_units, -1, dtype=np.float64) + + if _is_rcs: + precomputed_cohorts = precomputed["unit_cohorts"] + for g in unique_groups: + mask_g = precomputed_cohorts == g + unit_groups_array[mask_g] = g + elif global_unit_to_idx is not None: + idx_uid_pairs = [(idx, uid) for uid, idx in global_unit_to_idx.items()] + + if precomputed is not None: + precomputed_cohorts = precomputed["unit_cohorts"] + precomputed_unit_to_idx = precomputed["unit_to_idx"] + for idx, uid in idx_uid_pairs: + if uid in precomputed_unit_to_idx: + cohort = precomputed_cohorts[precomputed_unit_to_idx[uid]] + if cohort in unique_groups_set: + unit_groups_array[idx] = cohort + else: + for idx, uid in idx_uid_pairs: + unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] + if unit_first_treat in unique_groups_set: + unit_groups_array[idx] = unit_first_treat + else: + idx_uid_pairs = list(enumerate(all_units)) + for idx, uid in idx_uid_pairs: + unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] + if unit_first_treat in unique_groups_set: + unit_groups_array[idx] = unit_first_treat + + groups_for_gt_array = np.array(groups_for_gt) + indicator_matrix = ( + unit_groups_array[:, np.newaxis] == groups_for_gt_array[np.newaxis, :] + ).astype(np.float64) + + if survey_w is not None: + if _is_rcs and precomputed is not None: + unit_sw = survey_w + elif global_unit_to_idx is not None and precomputed is not None: + unit_sw = np.zeros(n_units) + precomputed_unit_to_idx_local = precomputed["unit_to_idx"] + idx_uid_pairs_sw = [(idx, uid) for uid, idx in global_unit_to_idx.items()] + for idx, uid in idx_uid_pairs_sw: + if uid in precomputed_unit_to_idx_local: + pc_idx = precomputed_unit_to_idx_local[uid] + unit_sw[idx] = survey_w[pc_idx] + else: + unit_sw = np.ones(n_units) + + weighted_indicator = indicator_matrix * unit_sw[:, np.newaxis] + indicator_diff = weighted_indicator - pg_keepers + indicator_sum_w = np.sum(indicator_diff, axis=1) + + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + if1_matrix = indicator_diff / sum_pg_keepers + if2_matrix = np.outer(indicator_sum_w, pg_keepers) / (sum_pg_keepers**2) + wif_matrix = if1_matrix - if2_matrix + wif_contrib = wif_matrix @ effects + else: + indicator_sum = np.sum(indicator_matrix - pg_keepers, axis=1) + + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + if1_matrix = (indicator_matrix - pg_keepers) / sum_pg_keepers + if2_matrix = np.outer(indicator_sum, pg_keepers) / (sum_pg_keepers**2) + wif_matrix = if1_matrix - if2_matrix + wif_contrib = wif_matrix @ effects + + if not np.all(np.isfinite(wif_contrib)): + return np.full(n_units, np.nan), all_units + + psi_wif = wif_contrib / total_weight + psi_total = psi_standard + psi_wif + + return psi_total, all_units + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_panel( + n_units=300, + n_periods=10, + n_cohorts=3, + n_cov=0, + never_frac=0.3, + drop_frac=0.0, + survey=False, + inf_coded=False, + seed=42, +): + """Small staggered panel with covariate-selected cohorts.""" + rng = np.random.default_rng(seed) + x = rng.normal(size=(n_units, max(n_cov, 1))) + treat_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + logits = 0.4 * x[:, 0][:, None] + rng.normal(0, 1, size=(n_units, n_cohorts)) + never = rng.random(n_units) < never_frac + g = np.where(never, 0, treat_periods[logits.argmax(axis=1)]) + + unit = np.repeat(np.arange(n_units), n_periods) + t = np.tile(np.arange(1, n_periods + 1), n_units) + g_long = np.repeat(g, n_periods) + treated = (g_long > 0) & (t >= g_long) + y = ( + np.repeat(rng.normal(0, 1, n_units), n_periods) + + 0.05 * t + + 2.0 * treated + + rng.normal(0, 1, n_units * n_periods) + ) + df = pd.DataFrame({"unit": unit, "time": t, "outcome": y, "first_treat": g_long}) + if inf_coded: + # staggered.py accepts np.inf for never-treated; exercises the + # input-normalization path in front of the cohort tables. + df["first_treat"] = df["first_treat"].astype(float) + df.loc[df["first_treat"] == 0, "first_treat"] = np.inf + for j in range(n_cov): + df[f"x{j + 1}"] = np.repeat(x[:, j], n_periods) + if survey: + df["pw"] = np.repeat(np.exp(rng.normal(0, 0.4, n_units)), n_periods) + if drop_frac > 0: + df = df[rng.random(len(df)) >= drop_frac].reset_index(drop=True) + return df + + +def _make_rcs(n_obs=20_000, n_periods=10, n_cohorts=3, n_cov=2, seed=42): + """Repeated cross-sections: unique unit id per row (panel=False).""" + rng = np.random.default_rng(seed) + x = rng.normal(size=(n_obs, max(n_cov, 1))) + treat_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + logits = 0.4 * x[:, 0][:, None] + rng.normal(0, 1, size=(n_obs, n_cohorts)) + never = rng.random(n_obs) < 0.3 + g = np.where(never, 0, treat_periods[logits.argmax(axis=1)]) + t = rng.integers(1, n_periods + 1, size=n_obs) + treated = (g > 0) & (t >= g) + y = 0.05 * t + 2.0 * treated + rng.normal(0, 1, n_obs) + df = pd.DataFrame({"unit": np.arange(n_obs), "time": t, "outcome": y, "first_treat": g}) + for j in range(n_cov): + df[f"x{j + 1}"] = x[:, j] + return df + + +def _make_ddd(n_units=400, n_periods=8, n_cohorts=2, seed=42): + """Staggered triple-diff panel: unit-level eligibility partition.""" + rng = np.random.default_rng(seed) + treat_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + never = rng.random(n_units) < 0.4 + g = np.where(never, 0, treat_periods[rng.integers(0, n_cohorts, size=n_units)]) + elig = (rng.random(n_units) < 0.5).astype(int) + + unit = np.repeat(np.arange(n_units), n_periods) + t = np.tile(np.arange(1, n_periods + 1), n_units) + g_long = np.repeat(g, n_periods) + e_long = np.repeat(elig, n_periods) + treated = (g_long > 0) & (t >= g_long) & (e_long == 1) + y = ( + np.repeat(rng.normal(0, 1, n_units), n_periods) + + 0.05 * t + + 0.3 * e_long + + 2.0 * treated + + rng.normal(0, 1, n_units * n_periods) + ) + return pd.DataFrame( + { + "unit": unit, + "period": t, + "outcome": y, + "first_treat": g_long, + "eligibility": e_long, + } + ) + + +class _ParityRecorder: + """Monkeypatch wrapper: on every combined-IF call, compare the live + implementation against the frozen pre-rewrite copy on identical inputs.""" + + def __init__(self, monkeypatch, atol=1e-9): + self.n_calls = 0 + self.atol = atol + orig = CallawaySantAnnaAggregationMixin._compute_combined_influence_function + + recorder = self + + def wrapper(mixin_self, *args, **kwargs): + new_psi, new_units = orig(mixin_self, *args, **kwargs) + old_psi, _ = _frozen_combined_influence_function(*args, **kwargs) + assert new_psi.shape == old_psi.shape + np.testing.assert_allclose( + new_psi, + old_psi, + rtol=0, + atol=recorder.atol, + equal_nan=True, + err_msg="fast-path psi diverged from frozen pre-rewrite copy", + ) + recorder.n_calls += 1 + return new_psi, new_units + + monkeypatch.setattr( + CallawaySantAnnaAggregationMixin, + "_compute_combined_influence_function", + wrapper, + ) + + +class TestCombinedIFFrozenParity: + """Old-vs-new psi_total parity across data families and aggregation + targets. aggregate='all' exercises the multi-cohort overall target, + every per-horizon keeper subset (incl. single-keeper horizons), and + single-cohort group targets in one fit.""" + + def _fit_cs(self, df, covs=None, survey=False, rcs=False, boot=0): + cs = CallawaySantAnna( + estimation_method="dr", + n_bootstrap=boot, + panel=not rcs, + seed=7, + ) + cs.fit( + df, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=covs, + aggregate="all", + survey_design=SurveyDesign(weights="pw") if survey else None, + ) + return cs + + def test_plain_panel(self, monkeypatch): + rec = _ParityRecorder(monkeypatch) + self._fit_cs(_make_panel(n_cov=2), covs=["x1", "x2"]) + assert rec.n_calls > 0 + + def test_survey_pweights(self, monkeypatch): + rec = _ParityRecorder(monkeypatch) + self._fit_cs(_make_panel(n_cov=2, survey=True), covs=["x1", "x2"], survey=True) + assert rec.n_calls > 0 + + def test_rcs(self, monkeypatch): + rec = _ParityRecorder(monkeypatch) + with pytest.warns(UserWarning, match="stationary"): + self._fit_cs(_make_rcs(), covs=["x1", "x2"], rcs=True) + assert rec.n_calls > 0 + + def test_unbalanced_panel(self, monkeypatch): + rec = _ParityRecorder(monkeypatch) + self._fit_cs(_make_panel(n_cov=2, drop_frac=0.2), covs=["x1", "x2"]) + assert rec.n_calls > 0 + + def test_inf_coded_never_treated(self, monkeypatch): + rec = _ParityRecorder(monkeypatch) + self._fit_cs(_make_panel(inf_coded=True)) + assert rec.n_calls > 0 + + def test_bootstrap_prep_path(self, monkeypatch): + """Bootstrap call sites (overall + per-horizon prep) go through the + same parity wrapper.""" + rec = _ParityRecorder(monkeypatch) + with pytest.warns(UserWarning, match="n_bootstrap"): + self._fit_cs(_make_panel(), boot=19) + assert rec.n_calls > 0 + + def test_triple_diff_zeroed_cohorts(self, monkeypatch): + """StaggeredTripleDifference aggregates through a shallow copy of + precomputed with an eligibility-zeroed unit_cohorts - the fourth + consumer of the mixin.""" + rec = _ParityRecorder(monkeypatch) + std = StaggeredTripleDifference() + std.fit( + _make_ddd(), + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + aggregate="all", + ) + assert rec.n_calls > 0 + + +class TestWifZeroContract: + """Units whose cohort is not among the keeper groups get a wif + contribution of exactly 0.0 on the fast path (the old dense form left an + ~eps residue from cancelling terms).""" + + def test_non_keeper_cohort_exact_zero(self): + n_units = 12 + # cohorts: 0 = never treated (not a keeper), 3 and 5 = keepers + unit_cohorts = np.array([0, 0, 0, 0, 3, 3, 3, 5, 5, 5, 0, 3]) + precomputed = { + "unit_cohorts": unit_cohorts, + "is_panel": True, + "unit_to_idx": {i: i for i in range(n_units)}, + } + gt_pairs = [(3, 4), (5, 6)] + # Zero-valued IFs isolate psi_wif: psi_total == wif / total_weight + influence_func_info = { + gt: { + "treated_idx": np.where(unit_cohorts == gt[0])[0], + "control_idx": np.where(unit_cohorts == 0)[0], + "treated_inf": np.zeros(np.sum(unit_cohorts == gt[0])), + "control_inf": np.zeros(np.sum(unit_cohorts == 0)), + } + for gt in gt_pairs + } + mixin = CallawaySantAnna() + psi, _ = mixin._compute_combined_influence_function( + gt_pairs, + weights=np.array([0.5, 0.5]), + effects=np.array([1.7, -0.9]), + groups_for_gt=np.array([3, 5]), + influence_func_info=influence_func_info, + df=None, + unit="unit", + precomputed=precomputed, + global_unit_to_idx=precomputed["unit_to_idx"], + n_global_units=n_units, + ) + non_keeper = unit_cohorts == 0 + assert np.all(psi[non_keeper] == 0.0) + assert np.all(psi[~non_keeper] != 0.0) + + +class TestAggCache: + """The per-fit cohort-table cache is validated by array identity and + rebuilt into a fresh dict - the StaggeredTripleDifference shallow-copy + pattern must never see stale tables.""" + + def _precomputed(self, cohorts, sw=None): + return { + "unit_cohorts": cohorts, + "survey_weights": sw, + "is_panel": True, + "unit_to_idx": {i: i for i in range(len(cohorts))}, + } + + def test_cache_hit_on_same_arrays(self): + p = self._precomputed(np.array([0, 3, 3, 5])) + c1 = CallawaySantAnnaAggregationMixin._get_agg_cache(p) + c2 = CallawaySantAnnaAggregationMixin._get_agg_cache(p) + assert c1 is c2 + assert c1["cohorts_ref"] is p["unit_cohorts"] + + def test_shallow_copy_with_swapped_cohorts_rebuilds(self): + p = self._precomputed(np.array([0, 3, 3, 5])) + c1 = CallawaySantAnnaAggregationMixin._get_agg_cache(p) + # StaggeredTripleDifference pattern: dict() copy + replaced cohorts + p2 = dict(p) + p2["unit_cohorts"] = np.array([0, 0, 3, 5]) # eligibility-zeroed + c2 = CallawaySantAnnaAggregationMixin._get_agg_cache(p2) + assert c2 is not c1 # fresh dict, not in-place mutation + assert c2["cohorts_ref"] is p2["unit_cohorts"] + # original cache untouched and still valid for the original dict + assert p["_agg_cache"] is c1 + assert CallawaySantAnnaAggregationMixin._get_agg_cache(p) is c1 + # masses reflect the respective cohort arrays + assert not np.array_equal(c1["cohort_masses"], c2["cohort_masses"]) + + def test_survey_weight_identity_checked(self): + cohorts = np.array([0, 3, 3, 5]) + p = self._precomputed(cohorts, sw=np.array([1.0, 2.0, 1.0, 0.5])) + c1 = CallawaySantAnnaAggregationMixin._get_agg_cache(p) + p["survey_weights"] = np.array([1.0, 1.0, 1.0, 1.0]) + c2 = CallawaySantAnnaAggregationMixin._get_agg_cache(p) + assert c2 is not c1 + assert c2["sw_ref"] is p["survey_weights"] + + +class TestFastPathTaken: + """A broken dispatch guard would silently fall back to the general path + and revert the entire optimization while every value-level test stays + green - so pin that the fast path actually runs (and resolves) for the + canonical panel fit.""" + + def test_fast_path_taken_on_plain_panel_fit(self, monkeypatch): + calls = {"resolved": 0, "fallback": 0} + seen = {} + orig = CallawaySantAnnaAggregationMixin._combined_if_fast + + def spy( + self, + gt_pairs, + weights, + effects, + groups_for_gt, + influence_func_info, + precomputed, + n_units, + ): + out = orig( + self, + gt_pairs, + weights, + effects, + groups_for_gt, + influence_func_info, + precomputed, + n_units, + ) + calls["resolved" if out is not None else "fallback"] += 1 + seen["precomputed"] = precomputed + return out + + monkeypatch.setattr(CallawaySantAnnaAggregationMixin, "_combined_if_fast", spy) + cs = CallawaySantAnna(estimation_method="dr", seed=7) + cs.fit( + _make_panel(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="all", + ) + # every aggregation target routed through the fast path, none fell + # back to the general (DataFrame-scanning) branch + assert calls["resolved"] > 0 + assert calls["fallback"] == 0 + # the per-fit cache exists and pins the exact arrays it was built on + cache = seen["precomputed"]["_agg_cache"] + assert cache["cohorts_ref"] is seen["precomputed"]["unit_cohorts"] + + def _spy(self, monkeypatch): + """Count fast-path dispatch outcomes (resolved vs fell-back).""" + calls = {"resolved": 0, "fallback": 0} + orig = CallawaySantAnnaAggregationMixin._combined_if_fast + + def spy(mixin_self, *args): + out = orig(mixin_self, *args) + calls["resolved" if out is not None else "fallback"] += 1 + return out + + monkeypatch.setattr(CallawaySantAnnaAggregationMixin, "_combined_if_fast", spy) + return calls + + def test_rcs_fast_path_taken(self, monkeypatch): + calls = self._spy(monkeypatch) + cs = CallawaySantAnna(estimation_method="dr", panel=False, seed=7) + with pytest.warns(UserWarning, match="stationary"): + cs.fit( + _make_rcs(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + aggregate="all", + ) + assert calls["resolved"] > 0 + assert calls["fallback"] == 0 + + def test_survey_fast_path_taken(self, monkeypatch): + """Survey-weighted fits must dispatch through the fast path too - a + silent fallback would pass every parity test while reverting the + optimization for the whole survey family.""" + calls = self._spy(monkeypatch) + cs = CallawaySantAnna(estimation_method="dr", seed=7) + cs.fit( + _make_panel(n_cov=2, survey=True), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + aggregate="all", + survey_design=SurveyDesign(weights="pw"), + ) + assert calls["resolved"] > 0 + assert calls["fallback"] == 0 + + def test_bootstrap_prep_fast_path_taken(self, monkeypatch): + """Bootstrap prep (overall + per-horizon combined-IF calls) must + dispatch through the fast path.""" + calls = self._spy(monkeypatch) + cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=19, seed=7) + with pytest.warns(UserWarning, match="n_bootstrap"): + cs.fit( + _make_panel(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="all", + ) + assert calls["resolved"] > 0 + assert calls["fallback"] == 0 + + def test_triple_diff_fast_path_taken(self, monkeypatch): + """StaggeredTripleDifference's shallow-copied precomputed (zeroed + cohorts, no canonical_size key) must still resolve the fast path via + the .get defaults.""" + calls = self._spy(monkeypatch) + std = StaggeredTripleDifference() + std.fit( + _make_ddd(), + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + aggregate="all", + ) + assert calls["resolved"] > 0 + assert calls["fallback"] == 0