diff --git a/CHANGELOG.md b/CHANGELOG.md index 31705d37..48d9a1e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -228,6 +228,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 the legacy pseudoinverse path as described under `omega_ridge` above; at `omega_ridge=0` results are bit-identical to the previous release. The conditional-path scalability warning now fires once per fit at n > 50,000 (previously per cell at n > 5,000). +- **`EfficientDiD` conditional-path kernel-covariance tables hoisted across `(g,t)` cells** + (follow-up to the fused tiled pass above). Every Omega* term is a kernel covariance keyed only + by outcome columns, so the tables the tiled pass rebuilt per cell dedup to one table of + distinct product columns per comparison group per unit-tile (~26x fewer kernel GEMM columns + on a PT-All fit); kernel weight matrices are built and freed one group at a time, so the tile + memory budget is governed by the largest single group instead of the sum (proportionally + fatter tiles at large n); each cell's Omega* is gathered from the group tables through + compact per-cell column slices in the legacy per-entry operation order (value-exact, locked + by test). Measured (3-rep medians, 20 periods, 5 cohorts, 5 covariates, `aggregate="all"`): + fit 7.7s → 7.0s at 2k units, 57.2s → 35.7s at 10k (1.6x), 358s → 129s at 30k (2.8x), and + the 100k-unit fit drops from ~86 min to **17.8 min (~4.8x)** at the same ~4 GB peak RSS; + the kernel-covariance stage itself is 40x faster (21.2s → 0.5s at 10k). Results move only at + floating-point reassociation level (post-treatment cells ~1e-12 relative, overall ATT ~1e-13); + the no-covariates path is byte-identical and `omega_ridge=0` still routes the entire legacy + path. A `pt_assumption="post"` covariate fit builds no tables at all (all cells are + just-identified), keeping that path regression-free. - **CallawaySantAnna multiplier bootstrap rewritten as a fused, column-tiled scatter-GEMM** (EfficientDiD's bootstrap routes through the same kernel). The former loop sliced the `(block × n_units)` weight matrix twice per (g,t) cell per weight block — at a 40-period, diff --git a/TODO.md b/TODO.md index 71119124..48410b7a 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| `EfficientDiD` conditional path at very large n: the fused tiled GEMM rewrite (v3.7; 43x at 2k units, 100k now completes in ~86 min) computes kernel-covariance tables per cell per tile. Cross-cell table hoisting — Term 5 tables are (g,t)-independent and Term 2 tables depend only on t — would cut the intrinsic per-cell GEMM factor by ~10x at 100k-unit scale if practitioners need routine 100k covariate fits. Follow-up lever from the v3.7 flop analysis; not needed below ~50k units. | `efficient_did_covariates.py::compute_conditional_cells_tiled` | CS-scaling | Mid | Low | +| `EfficientDiD` conditional path: post-hoisting (2026-07, 100k fit ~18 min) the top stage at every scale is `_ridge_solve_weights` — batched `np.linalg.solve` LU over `(n × H × H)` stacks (17.2s of a 35.7s fit at 10k units, ~170s at 100k; linear in n). The matrices are SPD after the ridge, so a batched Cholesky (potrf ~1/2 the LU flops) in the Rust backend — which could also thread across the stack — is the natural lever; numpy has no batched triangular solve, so pure-Python gains are limited. | `efficient_did_covariates.py::_ridge_solve_weights`, `rust/src/` | CS-scaling | Mid | Low | | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | diff --git a/diff_diff/efficient_did_covariates.py b/diff_diff/efficient_did_covariates.py index 0158db26..42a35bf5 100644 --- a/diff_diff/efficient_did_covariates.py +++ b/diff_diff/efficient_did_covariates.py @@ -1104,6 +1104,11 @@ def compute_omega_star_conditional( # can monkeypatch it to force multi-tile execution. _TARGET_OMEGA_TILE_BYTES = 256 * 1024 * 1024 +# Column-chunk width for the product GEMMs inside _build_group_kcov_table. +# Bounds the (n_group x chunk) product temporary independently of how many +# distinct kernel-covariance columns a fit needs. +_KCOV_PRODUCT_CHUNK = 512 + def _kcov_batch(W: np.ndarray, A: np.ndarray, B: np.ndarray) -> np.ndarray: """Batched kernel-weighted local covariances via GEMM. @@ -1135,6 +1140,248 @@ def _kcov_batch(W: np.ndarray, A: np.ndarray, B: np.ndarray) -> np.ndarray: return W @ (A0 * B0) - (W @ A0) * (W @ B0) +def _collect_kcov_specs( + cell_specs: List[Dict], + period_to_col: Dict[float, int], + never_treated_val: float, +) -> Tuple[Dict[float, Dict], Dict[Tuple[float, float], Dict]]: + """Cross-cell kernel-covariance key registry + per-cell gather specs. + + Every Omega* entry is a sum of ``s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 + | group)`` terms, keyed only by wide-outcome column indices - not by the + (g, t) cell - so the H(H+1)/2-pair tables that + ``compute_conditional_cells_tiled`` used to rebuild per cell dedup to one + table of distinct product columns per comparison group per tile. + + Difference keys are ORDERED ``(minuend_col, subtrahend_col)`` and are + never sort-canonicalized: ``Y_u - Y_v = -(Y_v - Y_u)``, so sorting would + sign-flip covariances for cells whose pre-periods straddle t (possible + under PT-All). Only the symmetric product-pair swap is canonicalized + (``KCov(A, B) = KCov(B, A)``). + + Per-cell gather specs cover the UPPER TRIANGLE only; assembly builds a + ``(tile, H(H+1)/2)`` strip and writes it to both triangles, reproducing + the legacy loop's exact ``omega[k, j] = omega[j, k]`` mirror semantics + (computing the lower triangle independently would reassociate the + term-3/4 subtractions). Gathers go through compact per-cell column + slices of the group tables (cache-resident, ~3x faster than gathering + from the full table). Masked-out entries (no same-cohort term 3/4, no + matching-group term 5) resolve to an all-zero column: ``x - 0.0 == x`` + and ``x + 0.0 == x`` value-exactly in IEEE, matching the legacy + conditional skips. + + Returns + ------- + registry : dict mapping group key to + ``{"diff_keys": (nD, 2) int64 array, "prod_i", "prod_j": (nP,) int64 + arrays (indices into diff_keys), "zero_col": nP}``. + Only groups needed by some H >= 2 cell appear; H == 1 cells skip + Omega* entirely and contribute no keys (so e.g. a pure + pt_assumption="post" fit builds no tables and no kernel matrices). + cell_asm : dict mapping (g, t) of each H >= 2 cell to triu-entry specs + (entry order = ``np.triu_indices(H)``): + ``H``, ``iu``/``ju`` (triu index vectors), ``t1_col`` (group-g table + column), ``t2_cols``/``t2_loc`` (compact never-treated-table columns + + per-entry local map), ``t34_cols``/``t34_row_loc``/``t34_col_loc`` + (compact group-g columns incl. its zero column, or ``t34_cols = + None`` when the cell has no same-cohort pair), ``t5_parts`` (list of + (gkey, compact column array) in slot order), ``t5_loc`` (per-entry + local map into the slot-concatenated layout whose last column is + zero), ``m5`` (total concat width excl. the zero column). + """ + diff_index: Dict[float, Dict[Tuple[int, int], int]] = {} + prod_index: Dict[float, Dict[Tuple[int, int], int]] = {} + + def _prod_col(gkey: float, d1: Tuple[int, int], d2: Tuple[int, int]) -> int: + dmap = diff_index.setdefault(gkey, {}) + pmap = prod_index.setdefault(gkey, {}) + i1 = dmap.setdefault(d1, len(dmap)) + i2 = dmap.setdefault(d2, len(dmap)) + pkey = (i1, i2) if i1 <= i2 else (i2, i1) + return pmap.setdefault(pkey, len(pmap)) + + cell_asm: Dict[Tuple[float, float], Dict] = {} + for spec in cell_specs: + pairs = spec["pairs"] + H = len(pairs) + if H < 2: + continue + g = spec["g"] + t_col = spec["t_col"] + y1_col = spec["y1_col"] + j_c = [period_to_col[tp] for _, tp in pairs] + j_gp = [gp for gp, _ in pairs] + j_gkey = [never_treated_val if np.isinf(gp) else gp for gp in j_gp] + + t1_col = _prod_col(g, (t_col, y1_col), (t_col, y1_col)) + has_t34 = any(gp == g for gp in j_gp) + iu, ju = np.triu_indices(H) + n_tri = iu.size + t2_ids = np.empty(n_tri, dtype=np.int64) + t34_row_ids = np.full(n_tri, -1, dtype=np.int64) + t34_col_ids = np.full(n_tri, -1, dtype=np.int64) + cell_gkeys = sorted(set(j_gkey)) + slot_of = {gk: s for s, gk in enumerate(cell_gkeys)} + t5_slot = np.full(n_tri, -1, dtype=np.int64) + t5_ids = np.full(n_tri, -1, dtype=np.int64) + + for e in range(n_tri): + j = int(iu[e]) + k = int(ju[e]) + cj, ck = j_c[j], j_c[k] + t2_ids[e] = _prod_col(never_treated_val, (t_col, cj), (t_col, ck)) + if has_t34: + if j_gp[j] == g: + t34_row_ids[e] = _prod_col(g, (t_col, y1_col), (cj, y1_col)) + if j_gp[k] == g: + t34_col_ids[e] = _prod_col(g, (t_col, y1_col), (ck, y1_col)) + if j_gp[j] == j_gp[k]: + gk = j_gkey[j] + t5_slot[e] = slot_of[gk] + t5_ids[e] = _prod_col(gk, (cj, y1_col), (ck, y1_col)) + + cell_asm[(g, spec["t"])] = { + "H": H, + "iu": iu, + "ju": ju, + "t1_col": t1_col, + "has_t34": has_t34, + "t2_ids": t2_ids, + "t34_row_ids": t34_row_ids, + "t34_col_ids": t34_col_ids, + "gkeys": cell_gkeys, + "t5_slot": t5_slot, + "t5_ids": t5_ids, + } + + registry: Dict[float, Dict] = {} + for gkey, dmap in diff_index.items(): + pmap = prod_index[gkey] + diff_keys = np.array(sorted(dmap, key=lambda k: dmap[k]), dtype=np.int64) + prods = sorted(pmap, key=lambda k: pmap[k]) + registry[gkey] = { + "diff_keys": diff_keys, + "prod_i": np.array([p[0] for p in prods], dtype=np.int64), + "prod_j": np.array([p[1] for p in prods], dtype=np.int64), + "zero_col": len(prods), + } + + # Finalize: resolve -1 mask sentinels to each table's appended all-zero + # column, then convert global column ids to compact (unique columns + + # per-entry local map) form for cache-resident gathers at assembly time. + for (g, _t), asm in cell_asm.items(): + asm["t2_cols"], asm["t2_loc"] = np.unique(asm.pop("t2_ids"), return_inverse=True) + rows = asm.pop("t34_row_ids") + cols = asm.pop("t34_col_ids") + if asm.pop("has_t34"): + zc = registry[g]["zero_col"] + rows[rows == -1] = zc + cols[cols == -1] = zc + both, inv = np.unique(np.concatenate([rows, cols]), return_inverse=True) + asm["t34_cols"] = both + asm["t34_row_loc"] = inv[: rows.size] + asm["t34_col_loc"] = inv[rows.size :] + else: + asm["t34_cols"] = None + t5_slot = asm.pop("t5_slot") + t5_ids = asm.pop("t5_ids") + t5_loc = np.empty(t5_slot.size, dtype=np.intp) + parts = [] + offset = 0 + for s, gk in enumerate(asm.pop("gkeys")): + sel = t5_slot == s + u, inv = np.unique(t5_ids[sel], return_inverse=True) + parts.append((gk, u)) + t5_loc[sel] = offset + inv + offset += u.size + t5_loc[t5_slot == -1] = offset + asm["t5_parts"] = parts + asm["t5_loc"] = t5_loc + asm["m5"] = offset + + return registry, cell_asm + + +def _assemble_omega_tile( + asm: Dict, + tables: Dict[float, np.ndarray], + g: float, + never_treated_val: float, + n_rows: int, +) -> np.ndarray: + """Gather one cell's ``(n_rows, H, H)`` Omega* block from group tables. + + A ``(n_rows, H(H+1)/2)`` upper-triangle strip is accumulated in-place in + the legacy per-entry operation order (term 1 + term 2, then the term-3/4 + subtractions, then term 5), gathered through compact per-cell column + slices of the group tables, then written to both triangles (the legacy + ``omega[k, j] = omega[j, k]`` mirror). Masked entries hit an all-zero + column, matching the legacy conditional skips value-exactly; term 1 is + folded into the compact term-2 slice before the strip gather + (commutativity-exact, saves a full-width pass). + """ + table_g = tables[g] + t2c = tables[never_treated_val][:, asm["t2_cols"]] + t2c += table_g[:, asm["t1_col"]][:, np.newaxis] + tri = t2c[:, asm["t2_loc"]] + if asm["t34_cols"] is not None: + t34c = table_g[:, asm["t34_cols"]] + tri -= t34c[:, asm["t34_row_loc"]] + tri -= t34c[:, asm["t34_col_loc"]] + m5 = asm["m5"] + a5 = np.empty((n_rows, m5 + 1)) + off = 0 + for gkey, cols_gk in asm["t5_parts"]: + a5[:, off : off + cols_gk.size] = tables[gkey][:, cols_gk] + off += cols_gk.size + a5[:, m5] = 0.0 + tri += a5[:, asm["t5_loc"]] + + H = asm["H"] + omega_tile = np.empty((n_rows, H, H)) + omega_tile[:, asm["iu"], asm["ju"]] = tri + omega_tile[:, asm["ju"], asm["iu"]] = tri + return omega_tile + + +def _build_group_kcov_table( + W: np.ndarray, + y_group: np.ndarray, + group_registry: Dict, + s_tile: np.ndarray, +) -> np.ndarray: + """One group's kernel-covariance table for a unit tile. + + Same construction as :func:`_kcov_batch` (globally pre-centered + differenced columns through the GEMM identity + ``KCov = W @ (A0 * B0) - (W @ A0) * (W @ B0)``), with the per-difference + ``W @ D0`` GEMM shared across all product columns and the product columns + chunked to ``_KCOV_PRODUCT_CHUNK`` so the (n_group x chunk) temporary is + bounded. The finished table is pre-scaled by the group's own s vector + (every Omega* term's s factor is the table-group's s) and gets an + appended all-zero column for masked gather entries. + + Returns shape ``(n_rows, nP + 1)``. + """ + diff_keys = group_registry["diff_keys"] + prod_i = group_registry["prod_i"] + prod_j = group_registry["prod_j"] + n_p = prod_i.size + + D = y_group[:, diff_keys[:, 0]] - y_group[:, diff_keys[:, 1]] + D -= D.mean(axis=0) + WD = W @ D + + table = np.empty((W.shape[0], n_p + 1)) + for lo in range(0, n_p, _KCOV_PRODUCT_CHUNK): + hi = min(lo + _KCOV_PRODUCT_CHUNK, n_p) + table[:, lo:hi] = W @ (D[:, prod_i[lo:hi]] * D[:, prod_j[lo:hi]]) + table[:, :n_p] -= WD[:, prod_i] * WD[:, prod_j] + table[:, :n_p] *= s_tile[:, np.newaxis] + table[:, n_p] = 0.0 + return table + + def compute_conditional_cells_tiled( cell_specs: List[Dict], outcome_wide: np.ndarray, @@ -1157,14 +1404,20 @@ def compute_conditional_cells_tiled( Replaces the legacy per-cell chain (dense ``compute_omega_star_conditional`` H^2 loop -> per-unit SVD/pinv weights -> EIF) with, per unit-tile: - 1. ONE row-normalized kernel weight matrix per comparison group - (``_kernel_weights_matrix``), reused across ALL cells - the matrices - depend only on covariates/bandwidth/group, not on (g, t). - 2. Per cell, the five Eq 3.12 terms as kernel-covariance TABLES via - :func:`_kcov_batch` - Term 2/5 covariances depend only on the - ``(t_pre_j, t_pre_k)`` column combination, so the H(H+1)/2 pair loop - dedups to ~T^2 GEMM columns - then a gather-assembly of the - ``(tile, H, H)`` Omega* block. + 1. ONE kernel-covariance table per comparison group, HOISTED across all + cells (:func:`_collect_kcov_specs` / :func:`_build_group_kcov_table`): + every Eq 3.12 term is ``s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | + group)`` keyed only by wide-outcome columns, so the per-cell + H(H+1)/2-pair tables dedup to one table of distinct product columns + per group per tile (~26x fewer GEMM columns on a PT-All fit). The + group's row-normalized kernel weight matrix + (``_kernel_weights_matrix``) is built and FREED one group at a time, + so the tile budget is governed by the largest single group rather + than the sum of all groups. + 2. Per cell, a gather-assembly of the ``(tile, H, H)`` Omega* block from + the group tables via precomputed index arrays (upper triangle + mirrored at the index level, preserving the legacy loop's exact + per-entry operation sequence). 3. Ridge-regularized batched weights (:func:`_ridge_solve_weights`; requires ``omega_ridge > 0`` - the legacy ``omega_ridge=0`` path never reaches this function). @@ -1211,28 +1464,22 @@ def compute_conditional_cells_tiled( tile_bytes = _TARGET_OMEGA_TILE_BYTES h_max = max(len(spec["pairs"]) for spec in cell_specs) - # Per-tile-row footprint: (H, H) omega block + kernel-matrix rows across - # all groups (sum of group sizes <= n_units, counted twice for the cdist - # temp) + small per-cell vectors. - row_bytes = 8 * (h_max * h_max + 2 * n_units + 4 * h_max) - tile_units = int(max(1, min(n_units, tile_bytes // max(1, row_bytes)))) - - # Group-level (non-tiled) data, shared across cells and tiles - group_keys: set = set() - for spec in cell_specs: - group_keys.add(spec["g"]) - group_keys.add(never_treated_val) - for gp, _ in spec["pairs"]: - group_keys.add(never_treated_val if np.isinf(gp) else gp) + + # Cross-cell kcov key registry + per-cell gather specs (fit-level key + # bookkeeping; the tables they index are rebuilt per tile). Groups whose + # keys come only from H == 1 cells get no entry, no table, and no kernel + # matrix (e.g. a pure pt_assumption="post" fit builds none at all). + registry, cell_asm = _collect_kcov_specs(cell_specs, period_to_col, never_treated_val) def _mask_for(gkey: float) -> np.ndarray: return never_treated_mask if np.isinf(gkey) else cohort_masks[gkey] + # Group-level (non-tiled) data for the groups that need tables. y_grp: Dict[float, np.ndarray] = {} x_grp: Dict[float, np.ndarray] = {} w_grp: Dict[float, Optional[np.ndarray]] = {} s_all: Dict[float, np.ndarray] = {} - for gkey in group_keys: + for gkey in registry: mask = _mask_for(gkey) y_grp[gkey] = outcome_wide[mask] x_grp[gkey] = covariate_matrix[mask] @@ -1242,6 +1489,33 @@ def _mask_for(gkey: float) -> np.ndarray: np.full(n_units, 1.0 / max(cohort_fractions.get(gkey, 1e-10), 1e-10)), ) + # Per-tile-row footprint: the (H, H) omega block plus the triu strip and + # one gather temporary (together <= 2 H^2), the largest single group's + # kernel-matrix row counted twice (cdist temp + W - groups are built and + # freed ONE AT A TIME, so the budget is governed by the largest group, + # not the sum), the held group tables (product columns + WD columns + + # zero column per group), the widest cell's compact slices, the + # (tile x chunk) product-GEMM temporary, and small per-cell vectors. + # The group-side product temp (n_g x chunk) is tile-independent - carve + # it out of the budget before dividing. + n_g_max = max((y.shape[0] for y in y_grp.values()), default=0) + sum_tables = sum(r["zero_col"] + len(r["diff_keys"]) + 1 for r in registry.values()) + max_compact = max( + ( + asm["t2_cols"].size + + (asm["t34_cols"].size if asm["t34_cols"] is not None else 0) + + asm["m5"] + + 1 + for asm in cell_asm.values() + ), + default=0, + ) + row_bytes = 8 * ( + 2 * h_max * h_max + 2 * n_g_max + sum_tables + max_compact + _KCOV_PRODUCT_CHUNK + 4 * h_max + ) + budget = tile_bytes - 8 * _KCOV_PRODUCT_CHUNK * n_g_max + tile_units = int(max(1, min(n_units, budget // max(1, row_bytes)))) + scores: Dict[Tuple[float, float], np.ndarray] = { (spec["g"], spec["t"]): np.empty(n_units) for spec in cell_specs } @@ -1249,19 +1523,22 @@ def _mask_for(gkey: float) -> np.ndarray: for lo in range(0, n_units, tile_units): hi = min(lo + tile_units, n_units) - # Kernel weight matrices for this tile, one per group, lazily built - # and shared across ALL cells. - w_tile_cache: Dict[float, np.ndarray] = {} - - def _w_tile(gkey: float) -> np.ndarray: - if gkey not in w_tile_cache: - w_tile_cache[gkey] = _kernel_weights_matrix( + # One kcov table per group for this tile, shared across ALL cells. + # The kernel weight matrix is an argument temporary: built for one + # group, released when its table returns. + tables: Dict[float, np.ndarray] = {} + for gkey in registry: + tables[gkey] = _build_group_kcov_table( + _kernel_weights_matrix( covariate_matrix[lo:hi], x_grp[gkey], bandwidth, group_weights=w_grp[gkey], - ) - return w_tile_cache[gkey] + ), + y_grp[gkey], + registry[gkey], + s_all[gkey][lo:hi], + ) # Row-sliced views of the nuisance caches / masks for gen_out reuse masks_tile = {k: v[lo:hi] for k, v in cohort_masks.items()} @@ -1274,7 +1551,6 @@ def _w_tile(gkey: float) -> np.ndarray: g = spec["g"] t = spec["t"] pairs = spec["pairs"] - t_col = spec["t_col"] y1_col = spec["y1_col"] H = len(pairs) @@ -1297,71 +1573,9 @@ def _w_tile(gkey: float) -> np.ndarray: scores[(g, t)][lo:hi] = gen_out_tile[:, 0] continue - pcols = sorted({period_to_col[tp] for _, tp in pairs}) - pindex = {c: i for i, c in enumerate(pcols)} - n_p = len(pcols) - j_p = [pindex[period_to_col[tp]] for _, tp in pairs] - j_gp = [gp for gp, _ in pairs] - - iu, ju = np.triu_indices(n_p) - combo_idx = np.zeros((n_p, n_p), dtype=int) - combo_idx[iu, ju] = np.arange(len(iu)) - combo_idx[ju, iu] = combo_idx[iu, ju] - - s_g_tile = s_all[g][lo:hi] - s_inf_tile = s_all[never_treated_val][lo:hi] - wg = _w_tile(g) - winf = _w_tile(never_treated_val) - - y_g = y_grp[g] - y_inf = y_grp[never_treated_val] - a_g = y_g[:, t_col] - y_g[:, y1_col] - - # Term 1: s_g * KCov(a, a | g) - shared by every (j, k) - term1 = s_g_tile * _kcov_batch(wg, a_g[:, None], a_g[:, None])[:, 0] - - # Term 2 table: s_inf * KCov(Y_t - Y_pj, Y_t - Y_pk | inf) - u_inf = y_inf[:, [t_col]] - y_inf[:, pcols] - t2 = s_inf_tile[:, None] * _kcov_batch(winf, u_inf[:, iu], u_inf[:, ju]) - - # Term 3/4 table: s_g * KCov(a, Y_p - Y_1 | g), only if any - # same-cohort pair exists - t34 = None - if any(gp == g for gp in j_gp): - b_g = y_g[:, pcols] - y_g[:, [y1_col]] - t34 = s_g_tile[:, None] * _kcov_batch(wg, np.repeat(a_g[:, None], n_p, axis=1), b_g) - - # Term 5 tables per distinct comparison group - t5: Dict[float, np.ndarray] = {} - for gp in set(j_gp): - gkey = never_treated_val if np.isinf(gp) else gp - if gkey in t5: - continue - y_c = y_grp[gkey] - v_c = y_c[:, pcols] - y_c[:, [y1_col]] - t5[gkey] = s_all[gkey][lo:hi][:, None] * _kcov_batch( - _w_tile(gkey), v_c[:, iu], v_c[:, ju] - ) - - omega_tile = np.empty((hi - lo, H, H)) - for j in range(H): - pj = j_p[j] - gpj = j_gp[j] - for k in range(j, H): - pk = j_p[k] - gpk = j_gp[k] - val = term1 + t2[:, combo_idx[pj, pk]] - if gpj == g: - val = val - t34[:, pj] - if gpk == g: - val = val - t34[:, pk] - if gpj == gpk: - gkey = never_treated_val if np.isinf(gpj) else gpj - val = val + t5[gkey][:, combo_idx[pj, pk]] - omega_tile[:, j, k] = val - if j != k: - omega_tile[:, k, j] = val - + omega_tile = _assemble_omega_tile( + cell_asm[(g, t)], tables, g, never_treated_val, hi - lo + ) w_units = _ridge_solve_weights(omega_tile, omega_ridge) scores[(g, t)][lo:hi] = np.sum(w_units * gen_out_tile, axis=1) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 55f8e378..f17a8de1 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1254,7 +1254,7 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. - **PT-Post regime (just-identified)**: Under PT-Post, EDiD automatically reduces to standard single-baseline estimator (Corollary 3.2). No downside to using EDiD -- it subsumes standard estimators - **Duplicate rows**: Duplicate `(unit, time)` entries are rejected with `ValueError`. The estimator requires exactly one observation per unit-period - **Note:** PT-All index set includes g'=∞ (never-treated) as a candidate comparison group and excludes period_1 for all g'. When g'=∞, the second and third Eq 3.9 terms telescope so all (∞, t_pre) moments produce the same 2x2 DiD value; these redundant moments are damped by the default Omega* ridge (see the Omega* ridge Note below; at `omega_ridge=0`, by the legacy pseudoinverse). When t_pre = period_1, the third term degenerates to E[Y_1 - Y_1 | G=g'] = 0 for any g', adding no information. Valid pairs require only t_pre < g' (pre-treatment for comparison group), not t_pre < g. Same-group pairs (g'=g) are valid and contribute overidentifying moments (Equation 3.9) — except the degenerate PRE-treatment self-pair (g'=g, t_pre=t), which the default ridge path excludes (see the ridge Note). -- **Note:** Omega* ridge regularization (`omega_ridge`, default 1e-6; v3.7) — a documented refinement of the Omega* inversion in Eq 3.5/3.13/4.3, in the space the paper leaves open (it assumes Omega* is invertible and does not prescribe finite-sample handling of singularity). Under PT-All the overidentified moment set makes the SAMPLE Omega* numerically singular by construction (the telescoping (∞, t_pre) moments above plus near-duplicate cross-cohort moments): measured cond 1e17–1e22 for 100% of units on realistic panels, with a spectrum of one exact-null direction (the degenerate self-pair, below) plus a cluster of statistically-null directions at relative eigenvalue 1e-5–1e-8, far below the ~1e-2 sampling noise of the covariance entries — and no clean spectral gap. The prior pseudoinverse fallback therefore sat on the rcond-cutoff cliff: ANY floating-point-level change (BLAS reordering, platform change, a 1-ulp data perturbation) redrew per-cell weights and moved per-cell ATT(g,t) at up to ~1e-2 relative (measured 1.2e-4 per-cell rel for a 1-ulp outcome perturbation on the pre-v3.7 code; overall ATT stable at ~1e-9 because the redraw averages out across units and cells). The ridge solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` (trace-scaled: scale-equivariant, O(H), no SVD) — a smooth regularization with sensitivity bounded by 1/omega_ridge, not a cutoff. **Why the deviation is safe:** every moment individually identifies ATT(g,t), so any fixed weights summing to 1 keep the estimator consistent; the ridge trades a numerically ill-defined efficiency optimum for a stable one, and the plug-in EIF treatment of estimated weights (Remark 4.2) is unaffected. **Calibration evidence (2026-07):** default 1e-6 chosen by 1-ulp stability (per-cell rel <= 3e-9 vs 1.2e-4 legacy; candidates 1e-4/1e-6/1e-8 all pass the <=1e-6 target) with the HRS Table 6 anchors as a hard gate (all anchors within the published-value tolerance at every candidate; worst deviation 0.0257 SE, unchanged from legacy; shift <= 0.0001 SE); Monte Carlo on covariate-confounded DGPs shows bias/RMSE/SE-calibration/coverage statistically identical to legacy (ridge marginally better point metrics). **One-time value shift:** per-cell and event-study values on the covariate path move within the pre-existing indeterminacy band when upgrading (worst observed post-treatment cell shift ~0.6 of its own SE at n=500, shrinking with n: overall-ATT shift 1.6e-2 → 4e-3 → 1.1e-3 rel at n=500/1k/2k); the no-covariates path is essentially unchanged (~1e-7). `omega_ridge=0` restores the ENTIRE legacy code path bit-for-bit (both the omega constructions and the inv/pinv weights, including per-cell condition-number warnings and the legacy O(n^2 H^2) runtime). **Degenerate self-pair:** for PRE-treatment cells (t < g), the pair (g'=g, t_pre=t) telescopes to the identically-zero moment 0=0 (the exact-null Omega* direction). The legacy pseudoinverse truncated it, spreading weight over noisy moments (spurious pre-treatment placebos of ~5% of the effect size, pure noise amplification); a naive ridge would instead load all weight on the zero-variance moment, making placebos deterministically zero and silently disabling the pre-trend diagnostic. The default ridge path therefore drops this zero-information pair (fit-level filter; post-treatment cells never contain it), restoring honest data-driven placebos; `omega_ridge=0` keeps the legacy pair set. **Warnings/diagnostics:** the no-covariates path still computes per-cell condition numbers (cheap at (H,H)) for `results.omega_condition_numbers` and consolidates cells with cond > 1e12 into ONE fit-level warning (count + max cond) instead of the legacy per-cell pseudoinverse warnings; the covariate path intentionally computes NO per-unit condition numbers — they would cost exactly the per-unit SVDs the v3.7 rewrite removes, near-singularity there is structural (~always true under PT-All, so a warning would be always-on noise rather than signal), and the ridge handles it by design (the legacy covariate path likewise had no per-unit diagnostics). The scalability warning threshold moved from n > 5000 (legacy per-cell warning) to n > 50,000 (one fit-level warning; the kernel stage is still intrinsically O(n^2) but with a ~100x lower constant and tile-bounded memory). +- **Note:** Omega* ridge regularization (`omega_ridge`, default 1e-6; v3.7) — a documented refinement of the Omega* inversion in Eq 3.5/3.13/4.3, in the space the paper leaves open (it assumes Omega* is invertible and does not prescribe finite-sample handling of singularity). Under PT-All the overidentified moment set makes the SAMPLE Omega* numerically singular by construction (the telescoping (∞, t_pre) moments above plus near-duplicate cross-cohort moments): measured cond 1e17–1e22 for 100% of units on realistic panels, with a spectrum of one exact-null direction (the degenerate self-pair, below) plus a cluster of statistically-null directions at relative eigenvalue 1e-5–1e-8, far below the ~1e-2 sampling noise of the covariance entries — and no clean spectral gap. The prior pseudoinverse fallback therefore sat on the rcond-cutoff cliff: ANY floating-point-level change (BLAS reordering, platform change, a 1-ulp data perturbation) redrew per-cell weights and moved per-cell ATT(g,t) at up to ~1e-2 relative (measured 1.2e-4 per-cell rel for a 1-ulp outcome perturbation on the pre-v3.7 code; overall ATT stable at ~1e-9 because the redraw averages out across units and cells). The ridge solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` (trace-scaled: scale-equivariant, O(H), no SVD) — a smooth regularization with sensitivity bounded by 1/omega_ridge, not a cutoff. **Why the deviation is safe:** every moment individually identifies ATT(g,t), so any fixed weights summing to 1 keep the estimator consistent; the ridge trades a numerically ill-defined efficiency optimum for a stable one, and the plug-in EIF treatment of estimated weights (Remark 4.2) is unaffected. **Calibration evidence (2026-07):** default 1e-6 chosen by 1-ulp stability (per-cell rel <= 3e-9 vs 1.2e-4 legacy; candidates 1e-4/1e-6/1e-8 all pass the <=1e-6 target) with the HRS Table 6 anchors as a hard gate (all anchors within the published-value tolerance at every candidate; worst deviation 0.0257 SE, unchanged from legacy; shift <= 0.0001 SE); Monte Carlo on covariate-confounded DGPs shows bias/RMSE/SE-calibration/coverage statistically identical to legacy (ridge marginally better point metrics). **One-time value shift:** per-cell and event-study values on the covariate path move within the pre-existing indeterminacy band when upgrading (worst observed post-treatment cell shift ~0.6 of its own SE at n=500, shrinking with n: overall-ATT shift 1.6e-2 → 4e-3 → 1.1e-3 rel at n=500/1k/2k); the no-covariates path is essentially unchanged (~1e-7). `omega_ridge=0` restores the ENTIRE legacy code path bit-for-bit (both the omega constructions and the inv/pinv weights, including per-cell condition-number warnings and the legacy O(n^2 H^2) runtime). **Degenerate self-pair:** for PRE-treatment cells (t < g), the pair (g'=g, t_pre=t) telescopes to the identically-zero moment 0=0 (the exact-null Omega* direction). The legacy pseudoinverse truncated it, spreading weight over noisy moments (spurious pre-treatment placebos of ~5% of the effect size, pure noise amplification); a naive ridge would instead load all weight on the zero-variance moment, making placebos deterministically zero and silently disabling the pre-trend diagnostic. The default ridge path therefore drops this zero-information pair (fit-level filter; post-treatment cells never contain it), restoring honest data-driven placebos; `omega_ridge=0` keeps the legacy pair set. **Warnings/diagnostics:** the no-covariates path still computes per-cell condition numbers (cheap at (H,H)) for `results.omega_condition_numbers` and consolidates cells with cond > 1e12 into ONE fit-level warning (count + max cond) instead of the legacy per-cell pseudoinverse warnings; the covariate path intentionally computes NO per-unit condition numbers — they would cost exactly the per-unit SVDs the v3.7 rewrite removes, near-singularity there is structural (~always true under PT-All, so a warning would be always-on noise rather than signal), and the ridge handles it by design (the legacy covariate path likewise had no per-unit diagnostics). The scalability warning threshold moved from n > 5000 (legacy per-cell warning) to n > 50,000 (one fit-level warning; the kernel stage is still intrinsically O(n^2) but with a ~100x lower constant and tile-bounded memory). **Cross-cell table hoisting (2026-07 follow-up):** the tiled implementation now builds the kernel-covariance tables once per comparison group per unit-tile instead of per (g, t) cell — every Omega* term is `s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | group)`, keyed only by wide-outcome columns, so the per-cell H(H+1)/2-pair tables dedup to distinct product columns per group (~26x fewer kernel GEMM columns on a PT-All fit), and the kernel weight matrices are built and freed one group at a time (the tile memory budget is governed by the largest single group rather than the sum, giving proportionally fatter tiles at large n). Each cell's Omega* is then gathered from the group tables in the same per-entry operation order as the per-cell construction (value-exact, locked by test); this is a pure implementation change — measured results move only at floating-point reassociation level (post-treatment cells ~1e-12 relative, overall ATT ~1e-13; the no-covariates path is byte-identical), and `omega_ridge=0` still routes the entire legacy path. - **Note:** Bootstrap aggregation uses fixed cohort-size weights for overall/event-study reaggregation, matching the CallawaySantAnna bootstrap pattern (staggered_bootstrap.py:281 computes `bootstrap_overall = bootstrap_atts_gt[:, post_indices] @ weights`; L297 uses the same fixed-weight pattern for event study). The analytical path includes a WIF correction; fixed-weight bootstrap captures the same sampling variability through per-cell EIF perturbation without re-estimating aggregation weights, consistent with both the library's CS implementation and the R `did` package. - **Overall ATT convention**: The library's `overall_att` uses cohort-size-weighted averaging of post-treatment (g,t) cells, matching the CallawaySantAnna simple aggregation. This differs from the paper's ES_avg (Eq 2.3), which uniformly averages over event-time horizons. ES_avg can be computed from event study output as `mean(event_study_effects[e]["effect"] for e >= 0)` diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 578189ba..b6939e05 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -277,9 +277,56 @@ with n as the indeterminacy band tightens): overall ATT 1.6e-2 / 4.0e-3 / 1.1e-3 at n=500/1k/2k; worst post-treatment cell <= 0.58 of its own SE; nocov path ~1e-7. 1-ulp per-cell stability after: <= 3e-9 rel (before: 1.2e-4). -Remaining lever (TODO row): cross-cell kernel-table hoisting (Term 5 tables are -(g,t)-independent; Term 2 depends only on t) — ~10x on the intrinsic per-cell GEMM -factor at 100k-unit scale; not needed below ~50k units. +The lever identified at the time — cross-cell kernel-table hoisting (Term 5 +tables are (g,t)-independent; Term 2 depends only on t), estimated ~10x on the +intrinsic per-cell GEMM factor at 100k-unit scale — was filed as a TODO row and +is RESOLVED by the follow-up below. + +### Cross-cell kcov-table hoisting + per-group W lifecycle (2026-07 follow-up) + +Resolves the hoisting TODO row above. Instrumented stage split of the v3.7 tiled +path (`.bench-local/edid_hoist_spike.py`; same 20p/5-cohort/5-cov scenario): + +| stage | n=2k | n=10k | n=30k | scaling | +|---|---|---|---|---| +| t5 + t2 kcov GEMMs | 0.95s | 21.2s | 217.2s | O(n^2), dominant | +| ridge_solve | 3.5s | 16.9s | 48.5s | linear | +| assembly (Python (j,k) loop) | 1.4s | 6.6s | 28.3s | linear x tile-count overhead | +| fit total | 7.7s | 57.2s | 359s | | + +Duplication: t5 computed 108,300 kcov columns per tile where 1,140 are distinct +(95x); t2 18,050 vs 3,610 (5x) — ~26.6x combined GEMM-traffic reduction. A +"base-table" bilinear factorization (level covariances combined 4-term, ~90x) +was measured and REJECTED: ~500x fp-error amplification on persistent panels +(`.bench-local/edid_hoist_cancellation.py`), which through the ridge's 1/lambda +sensitivity cap lands AT the 1e-6 stability contract. The shipped variant hoists +the DIFFERENCED-column tables (numerics identical in kind to the per-cell +construction). Assembly variants were micro-benchmarked: full (H,H) index +gathers were memory-bandwidth-bound and SLOWER than the legacy Python loop; +the shipped triu-strip gather through compact per-cell column slices is ~5x +faster than the naive vectorization and faster than the loop, while preserving +the legacy per-entry operation order value-exactly (term 1 folded into the +compact slice pre-gather; mirror preserved at the copy level). Kernel W +matrices are built and freed one group at a time, so the tile budget is set by +the largest single group (~3-5x fatter tiles at 100k). + +Measured arms (3-rep medians; BEFORE = pristine main fcd77683 via +baseline-worktree PYTHONPATH; `.bench-local/edid_hoist_arms_results.jsonl`): + +| scenario | BEFORE | AFTER | speedup | maxrss B->A | +|---|---|---|---|---| +| conditional n=2k | 7.71s | 6.96s | 1.11x | 0.59 -> 0.59 GB | +| conditional n=10k | 57.2s | 35.7s | 1.60x | 0.76 -> 0.79 GB | +| conditional n=30k | 358.4s | 128.8s | 2.78x | 0.96 -> 1.12 GB | +| conditional n=100k | 85.6 min (v3.7 measured, 1 rep) | 17.8 min (1 rep) | ~4.8x | 4.0 -> 4.02 GB | +| nocov n=2k (guard) | 0.44s | 0.44s | 1.00x (byte-identical) | flat | +| survey n=1k | 3.30s | 3.03s | 1.09x | flat | + +The kernel-covariance stage itself: 21.2s -> 0.53s at 10k (40x). Deltas: +post-treatment cells ~1e-12 rel (max 6e-12), overall ATT ~1e-13, nocov exactly +0. Post-hoist the top stages at every scale are `_ridge_solve_weights` (batched +LAPACK LU over (n x H x H) stacks; 17.2s at 10k — Rust batched-Cholesky is the +natural follow-up, TODO row filed) and the nuisance sieves. --- diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index 2458aa93..91eb1cd1 100644 --- a/tests/test_efficient_did.py +++ b/tests/test_efficient_did.py @@ -794,6 +794,93 @@ def test_tile_forced_twin_survey_weights(self, monkeypatch): r_tiled = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) self._assert_close_surfaces(r_tiled, r_single, rtol=1e-10, atol=1e-12) + def test_table_build_count_proportional_to_tiles(self, monkeypatch): + """v3.8 hoisting proof: the per-group kcov table is built once per + (group-with-keys, tile) - NOT once per cell. Forcing one-unit tiles + must scale the build count by exactly n_units.""" + import diff_diff.efficient_did_covariates as cov_mod + + n_units = 40 + df = self._cov_df(n_units=n_units) + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + + calls = {"n": 0} + orig_build = cov_mod._build_group_kcov_table + + def counting_build(*args, **kwargs): + calls["n"] += 1 + return orig_build(*args, **kwargs) + + monkeypatch.setattr(cov_mod, "_build_group_kcov_table", counting_build) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + single_tile_calls = calls["n"] + # one build per group-with-keys (2 cohorts + never-treated), far + # fewer than the number of H >= 2 cells + assert single_tile_calls == 3, single_tile_calls + + calls["n"] = 0 + monkeypatch.setattr(cov_mod, "_TARGET_OMEGA_TILE_BYTES", 1) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + assert calls["n"] == single_tile_calls * n_units, calls["n"] + + def test_pt_post_builds_no_tables(self, monkeypatch): + """Under pt_assumption="post" every cell has a single pair (H == 1), + so the hoisted path must build ZERO kcov tables and ZERO kernel + matrices - guarding the no-regression contract for that path.""" + import diff_diff.efficient_did_covariates as cov_mod + + df = self._cov_df(n_units=120) + calls = {"tables": 0, "kernels": 0} + orig_build = cov_mod._build_group_kcov_table + orig_kwm = cov_mod._kernel_weights_matrix + + def counting_build(*args, **kwargs): + calls["tables"] += 1 + return orig_build(*args, **kwargs) + + def counting_kwm(*args, **kwargs): + calls["kernels"] += 1 + return orig_kwm(*args, **kwargs) + + monkeypatch.setattr(cov_mod, "_build_group_kcov_table", counting_build) + monkeypatch.setattr(cov_mod, "_kernel_weights_matrix", counting_kwm) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r = EfficientDiD(pt_assumption="post").fit( + df, + "y", + "unit", + "time", + "first_treat", + covariates=["x1", "x2"], + aggregate="all", + ) + assert calls["tables"] == 0 + assert calls["kernels"] == 0 + assert np.isfinite(r.att) + + def test_mid_size_tile_twin_with_chunked_products(self, monkeypatch): + """Multi-tile execution at a realistic tile width, with the product + GEMM forced through multiple column chunks: results must match the + single-tile fit (rel 1e-10).""" + import diff_diff.efficient_did_covariates as cov_mod + + df = self._cov_df(n_units=150) + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_single = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + monkeypatch.setattr(cov_mod, "_TARGET_OMEGA_TILE_BYTES", 1_000_000) + monkeypatch.setattr(cov_mod, "_KCOV_PRODUCT_CHUNK", 3) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_tiled = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + self._assert_close_surfaces(r_tiled, r_single, rtol=1e-10, atol=1e-12) + def test_one_ulp_stability_conditional(self): """THE stability contract: a 1-ulp outcome perturbation moves per-cell and event-study ATTs and SEs by <= 1e-6 relative under the diff --git a/tests/test_methodology_efficient_did.py b/tests/test_methodology_efficient_did.py index 8c6ee9b0..5275f79b 100644 --- a/tests/test_methodology_efficient_did.py +++ b/tests/test_methodology_efficient_did.py @@ -473,6 +473,203 @@ def test_row_tile_equals_full(self): np.testing.assert_allclose(tiled, full, rtol=1e-13, atol=1e-12) +class TestKcovTableHoisting: + """v3.8 cross-cell hoisted kernel-covariance tables: the per-group table + built by ``_build_group_kcov_table`` from the ``_collect_kcov_specs`` + registry must reproduce, column by column, the per-cell ``_kcov_batch`` + construction it dedups (every Omega* term is + ``s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | group)``), and the gather + assembly must reproduce the legacy per-entry (j, k) loop value-exactly. + + The spec sets below are synthetic/defensive at the helper level: they + exercise key shapes (reversed orientation, per-cell ``y1_col``) beyond + what a production fit can currently reach, to lock the key contracts.""" + + T = 8 + + @classmethod + def _setup(cls, seed=0, weighted=False): + from diff_diff.efficient_did_covariates import _kernel_weights_matrix + + rng = np.random.default_rng(seed) + n = 80 + outcome = rng.normal(size=(n, cls.T)).cumsum(axis=1) + x = rng.normal(size=(n, 3)) + first = np.full(n, np.inf) + first[:25] = 4.0 + first[25:45] = 6.0 + masks = {4.0: first == 4.0, 6.0: first == 6.0, np.inf: np.isinf(first)} + uw = np.exp(rng.normal(0, 0.4, n)) if weighted else None + s_full = {gk: rng.uniform(1.0, 5.0, n) for gk in masks} + + def w_for(gkey): + gw = uw[masks[gkey]] if uw is not None else None + return _kernel_weights_matrix(x, x[masks[gkey]], 0.9, group_weights=gw) + + return outcome, masks, s_full, w_for + + @staticmethod + def _period_to_col(): + return {float(p): p - 1 for p in range(1, TestKcovTableHoisting.T + 1)} + + @staticmethod + def _default_specs(): + # (gp, tpre) pairs mixing never-treated, same-cohort (term 3/4 + # active), and another cohort (a second term-5 group); the two cells + # share t2/t5 keys, exercising the cross-cell dedup. + return [ + { + "g": 4.0, + "t": 5.0, + "t_col": 4, + "y1_col": 0, + "pairs": [(np.inf, 2.0), (np.inf, 3.0), (4.0, 2.0), (6.0, 3.0)], + }, + { + "g": 4.0, + "t": 6.0, + "t_col": 5, + "y1_col": 0, + "pairs": [(np.inf, 2.0), (np.inf, 3.0), (4.0, 3.0)], + }, + ] + + def _tables_and_oracle(self, specs, seed=0, weighted=False): + from diff_diff.efficient_did_covariates import ( + _build_group_kcov_table, + _collect_kcov_specs, + _kcov_batch, + ) + + outcome, masks, s_full, w_for = self._setup(seed, weighted) + registry, cell_asm = _collect_kcov_specs(specs, self._period_to_col(), np.inf) + tables, oracles = {}, {} + for gkey, reg in registry.items(): + y_g = outcome[masks[gkey]] + tables[gkey] = _build_group_kcov_table(w_for(gkey), y_g, reg, s_full[gkey]) + dk = reg["diff_keys"] + a = y_g[:, dk[reg["prod_i"], 0]] - y_g[:, dk[reg["prod_i"], 1]] + b = y_g[:, dk[reg["prod_j"], 0]] - y_g[:, dk[reg["prod_j"], 1]] + oracles[gkey] = s_full[gkey][:, None] * _kcov_batch(w_for(gkey), a, b) + return registry, cell_asm, tables, oracles + + @pytest.mark.parametrize("weighted", [False, True]) + def test_group_tables_match_per_cell_construction(self, weighted): + registry, _, tables, oracles = self._tables_and_oracle( + self._default_specs(), seed=1, weighted=weighted + ) + assert set(registry) == {4.0, 6.0, np.inf} + for gkey, reg in registry.items(): + n_p = reg["zero_col"] + np.testing.assert_allclose(tables[gkey][:, :n_p], oracles[gkey], rtol=1e-12, atol=1e-12) + assert np.all(tables[gkey][:, n_p] == 0.0) + + def test_reversed_orientation_pre_periods_straddle_t(self): + """Ordered-key contract (review-critical): a pre-treatment cell whose + pre-periods straddle t puts BOTH (t, earlier) and (t, later) columns + in the table; sort-canonicalizing the difference key would sign-flip + the mixed product KCov(Y_t - Y_2, Y_t - Y_5) at t=3.""" + specs = [ + {"g": 4.0, "t": 3.0, "t_col": 2, "y1_col": 0, "pairs": [(np.inf, 2.0), (np.inf, 5.0)]}, + ] + registry, _, tables, oracles = self._tables_and_oracle(specs, seed=2) + reg = registry[np.inf] + keys = {tuple(k) for k in reg["diff_keys"]} + assert (2, 1) in keys and (2, 4) in keys # (t_col, p_col) ordered + np.testing.assert_allclose( + tables[np.inf][:, : reg["zero_col"]], + oracles[np.inf], + rtol=1e-12, + atol=1e-12, + ) + + def test_multi_y1_cells_share_registry(self): + """Defensive/synthetic: production reaches per-cohort ``y1_col`` only + under pt_assumption="post" (all H == 1, no tables built), but the key + registry carries y1 so two cells with different baselines must dedup + without collision.""" + specs = self._default_specs() + specs[1] = dict(specs[1], y1_col=1) + registry, cell_asm, tables, oracles = self._tables_and_oracle(specs, seed=3) + for gkey, reg in registry.items(): + np.testing.assert_allclose( + tables[gkey][:, : reg["zero_col"]], + oracles[gkey], + rtol=1e-12, + atol=1e-12, + ) + # distinct y1 -> distinct term-1 columns for the two cells + a0 = cell_asm[(4.0, 5.0)] + a1 = cell_asm[(4.0, 6.0)] + assert a0["t1_col"] != a1["t1_col"] + + def test_assembly_matches_legacy_loop_value_exactly(self): + """The triu-strip gather assembly (incl. the term-1 fold, zero-column + masking, and the mirror write) must equal the legacy per-entry + (j, k) loop on identical table inputs - value-exact + (assert_array_equal; gather + add is not BLAS, and the op sequence + is preserved up to commutativity of the first add).""" + from diff_diff.efficient_did_covariates import _assemble_omega_tile + + specs = self._default_specs() + registry, cell_asm, tables, _ = self._tables_and_oracle(specs, seed=4) + + # key -> column lookup rebuilt from the registry (both pair orders) + col_of = {} + for gkey, reg in registry.items(): + dk = reg["diff_keys"] + for c in range(reg["zero_col"]): + d1 = tuple(dk[reg["prod_i"][c]]) + d2 = tuple(dk[reg["prod_j"][c]]) + col_of[(gkey, d1, d2)] = c + col_of[(gkey, d2, d1)] = c + + p2c = self._period_to_col() + for spec in specs: + g, t = spec["g"], spec["t"] + t_col, y1 = spec["t_col"], spec["y1_col"] + pairs = spec["pairs"] + H = len(pairs) + j_c = [p2c[tp] for _, tp in pairs] + j_gp = [gp for gp, _ in pairs] + n_rows = tables[g].shape[0] + + ref = np.empty((n_rows, H, H)) + for j in range(H): + for k in range(j, H): + cj, ck = j_c[j], j_c[k] + val = ( + tables[g][:, col_of[(g, (t_col, y1), (t_col, y1))]] + + tables[np.inf][:, col_of[(np.inf, (t_col, cj), (t_col, ck))]] + ) + if j_gp[j] == g: + val = val - tables[g][:, col_of[(g, (t_col, y1), (cj, y1))]] + if j_gp[k] == g: + val = val - tables[g][:, col_of[(g, (t_col, y1), (ck, y1))]] + if j_gp[j] == j_gp[k]: + gk = np.inf if np.isinf(j_gp[j]) else j_gp[j] + val = val + tables[gk][:, col_of[(gk, (cj, y1), (ck, y1))]] + ref[:, j, k] = val + if j != k: + ref[:, k, j] = val + + got = _assemble_omega_tile(cell_asm[(g, t)], tables, g, np.inf, n_rows) + np.testing.assert_array_equal(got, ref) + + def test_h1_cells_contribute_no_keys(self): + """H == 1 cells skip Omega* entirely: a spec set of only single-pair + cells yields an empty registry (no tables, no kernel matrices).""" + from diff_diff.efficient_did_covariates import _collect_kcov_specs + + specs = [ + {"g": 4.0, "t": 5.0, "t_col": 4, "y1_col": 3, "pairs": [(np.inf, 3.0)]}, + {"g": 6.0, "t": 7.0, "t_col": 6, "y1_col": 5, "pairs": [(np.inf, 5.0)]}, + ] + registry, cell_asm = _collect_kcov_specs(specs, self._period_to_col(), np.inf) + assert registry == {} + assert cell_asm == {} + + # ============================================================================= # Eq 3.9 — no-covariate generated outcome (telescoping) # =============================================================================