Archive · Quarto template · Source bundle
Comparison prepares a common analysis sample and dispatches independently to each enabled estimator. Covariates, weights and unbalanced or repeated-cross-section inputs can reach CS-DID even when a matrix method rejects them. The Credit configuration reuses validated regression results and sample-matched inference files; it does not refit those regressions on each render.
Crabbymetrics receives N × T cell-mean outcomes, a known treatment schedule and a separate observation mask. MC fits observed untreated entries; the missing-data SDID extension renormalizes donor/time weights over available observations; ridge trains on common pre-period overlap among donors observed at each prediction period. Cell mass weights reported ATT/event paths, not the fitting loss. Missing outcomes are not replaced by zeros, and observed treated targets are not dropped when overlap fails. The exact branch build is recorded with each new fit; complete-panel saved fits remain unchanged.
CS-DID reports equal-event-time and weighted observed-treated-cell summaries. Native simple and group aggregations are retained as diagnostics. Large intercept-only fits use equivalent cluster sufficient statistics. Multisynth uses fixed-fit wild-bootstrap intervals without a jackknife.
SHA-256: 3d6f87d28c1a24f0efe21b6d9b06d77866b0d16abcfdaaaef47465cc9cc1078e
1 """Configurable, checkpointed counterfactual comparison around archive regressions."""
2
3 from __future__ import annotations
4
5 import hashlib
6 import json
7 import subprocess
8 import sys
9 import time
10 from importlib.metadata import version
11 from pathlib import Path
12
13 import matplotlib.pyplot as plt
14 import pandas as pd
15 from IPython.display import Markdown, display
16
17 HERE = Path(__file__).resolve().parent
18 ROOT = HERE.parents[1]
19 sys.path.insert(0, str(ROOT / "templates/full_archive"))
20 from replication_engine import Replication, clean, save_json
21
22 LABELS = {
23 "csdid": "CS-DID",
24 "sdid": "SDID",
25 "matrix-completion": "Matrix completion",
26 "horizontal-ridge": "Horizontal ridge",
27 "multisynth-cohort": "Multisynth (cohort-averaged)",
28 "multisynth-individual": "Multisynth (individual)",
29 "mtgp": "MTGP MAP",
30 "trop": "TROP nuclear norm",
31 }
32 DEFAULTS = {
33 "csdid": {
34 "enabled": True,
35 "control_group": "notyettreated",
36 "aggregation": "dynamic",
37 },
38 "sdid": {"enabled": True},
39 "matrix-completion": {
40 "enabled": True,
41 "lambda_fraction": 0.25,
42 "max_iterations": 1000,
43 "tolerance": 1e-12,
44 },
45 "horizontal-ridge": {"enabled": True, "penalty": 1.0},
46 "multisynth-cohort": {
47 "enabled": True,
48 "time_cohort": True,
49 "nu": None,
50 "lambda": 0,
51 "eps_abs": 1e-6,
52 "eps_rel": 1e-6,
53 },
54 "multisynth-individual": {
55 "enabled": True,
56 "time_cohort": False,
57 "nu": None,
58 "lambda": 0,
59 "eps_abs": 1e-6,
60 "eps_rel": 1e-6,
61 },
62 "mtgp": {
63 "enabled": False,
64 "rank": 3,
65 "num_steps": 1500,
66 "learning_rate": 0.03,
67 "seed": 0,
68 },
69 "trop": {
70 "enabled": False,
71 "cv_n_trials": 50,
72 "tuning_cycles": 20,
73 "n_jobs": 1,
74 "random_seed": 0,
75 },
76 }
77
78
79 def digest(path):
80 return hashlib.sha256(Path(path).read_bytes()).hexdigest()
81
82
83 def markdown_table(rows):
84 display(Markdown(pd.DataFrame(rows).to_markdown(index=False, floatfmt=".5g")))
85
86
87 class Comparison:
88 def __init__(self, settings, options, project=ROOT):
89 self.root = Path(project)
90 self.options = options
91 self.settings = dict(settings)
92 source_slug = self.settings.get(
93 "specification_id", self.settings.get("slug", "study")
94 )
95 self.settings["output_dir"] = options.get(
96 "output_dir", f"analyses/expensive/{source_slug}"
97 )
98 self.base = Replication(self.settings, self.root)
99 self.output = self.base.output
100 self.methods = {
101 k: v | options.get("estimators", {}).get(k, {}) for k, v in DEFAULTS.items()
102 }
103 self.extra = {}
104
105 def prepare(self):
106 return self.base.prepare()
107
108 def _receipt(self, data, method, parameters):
109 dependencies = ["numpy", "pandas"]
110 if method == "csdid":
111 dependencies += ["csdid", "drdid"]
112 elif method in {"sdid", "matrix-completion", "horizontal-ridge"}:
113 dependencies += ["crabbymetrics"]
114 receipt = {
115 "panel_sha256": hashlib.sha256(
116 pd.util.hash_pandas_object(data, index=False).values.tobytes()
117 ).hexdigest(),
118 "method": method,
119 "parameters": parameters,
120 "code_sha256": {
121 p.name: digest(p) for p in HERE.iterdir() if p.suffix in {".py", ".R"}
122 },
123 "versions": {p: version(p) for p in dependencies},
124 }
125 if "crabbymetrics" in dependencies:
126 from matrix_models import build_identity
127
128 receipt["crabbymetrics_build"] = build_identity()
129 return receipt
130
131 def _eligible(self, method, data, xs, parameters=None):
132 from eligibility import eligibility
133
134 return eligibility(method, data, xs, parameters, self.base.config)
135
136 def estimate(self, reuse=True, retained=None):
137 baseline_path = self.options.get("regression_results")
138 if baseline_path:
139 baseline = json.loads((self.root / baseline_path).read_text())
140 assert baseline["source"] == self.base.results["source"], (
141 "Regression source mismatch"
142 )
143 for key, current in self.base.results["specifications"].items():
144 saved = baseline["specifications"][key]
145 saved_sample = {
146 "rows_without_never_control_time_fe_removed": 0
147 } | saved["sample"]
148 current_sample = {
149 "rows_without_never_control_time_fe_removed": 0
150 } | current["sample"]
151 assert current_sample == saved_sample, "Regression sample mismatch"
152 allowed = {"output_dir", "slug", "description", "labels"}
153 for key, value in self.base.config.items():
154 if key not in allowed:
155 assert (
156 {
157 "require_never_control_periods": False,
158 "frequency_weights": False,
159 }
160 | baseline["config"]
161 ).get(key) == value, f"Regression configuration mismatch: {key}"
162 self.base.results["specifications"] = baseline["specifications"]
163 elif reuse and (self.output / "estimates.json").exists():
164 self.base.load_results()
165 else:
166 from sparse_models import fit
167
168 large = any(
169 len(d) * len(d.loc[d.g.gt(0), ["g", "event_time"]].drop_duplicates())
170 > 100_000_000
171 for d, _ in self.base.panels.values()
172 )
173 if large:
174 for key, (data, xs) in self.base.panels.items():
175 fit(self.base, key, data, xs)
176 else:
177 self.base.estimate()
178 self.base.f_tests()
179 self.base.export()
180 for key, (data, xs) in self.base.panels.items():
181 self.extra[key] = []
182 folder = self.output / key
183 folder.mkdir(parents=True, exist_ok=True)
184 for slug, opts in self.methods.items():
185 if not opts["enabled"]:
186 continue
187 params = {k: v for k, v in opts.items() if k != "enabled"}
188 previous = (retained or {}).get(key, {}).get(LABELS[slug])
189 if previous is not None:
190 self.extra[key].append(previous)
191 continue
192 result = self._attempt(slug, data, xs, params, folder, reuse)
193 self.extra[key].append(result)
194 print(key, LABELS[slug], result["status"], flush=True)
195 return self.export()
196
197 def _attempt(self, slug, data, xs, params, folder, reuse):
198 """Isolate eligibility, cache, dependency and numerical errors by method."""
199 from eligibility import UnidentifiedTarget, missing_support_error
200
201 try:
202 reason = self._eligible(slug, data, xs, params)
203 if reason:
204 return {
205 "method": LABELS[slug],
206 "status": "unsupported",
207 "reason": reason,
208 }
209 receipt = self._receipt(data, slug, params)
210 cache = folder / f"{slug}.json"
211 stamp = folder / f"{slug}-receipt.json"
212 if (
213 reuse
214 and cache.exists()
215 and stamp.exists()
216 and json.loads(stamp.read_text()) == receipt
217 ):
218 return json.loads(cache.read_text())
219 result = self._fit(slug, data, params, folder, xs=xs)
220 result["status"] = "estimated"
221 save_json(cache, result)
222 save_json(stamp, receipt)
223 return result
224 except Exception as error: # noqa: BLE001 — each method has its own checkpoint
225 result = {
226 "method": LABELS[slug],
227 "status": "unidentified"
228 if isinstance(error, UnidentifiedTarget)
229 or (
230 isinstance(error, ValueError)
231 and missing_support_error(slug, str(error))
232 )
233 else "failed",
234 "reason": f"{type(error).__name__}: {error}",
235 }
236 if slug in {"sdid", "matrix-completion", "horizontal-ridge"}:
237 from matrix_models import build_identity, plain
238
239 result["parameters"] = params
240 result["crabbymetrics_build"] = build_identity()
241 if (
242 isinstance(error, AssertionError)
243 and error.args
244 and isinstance(error.args[0], dict)
245 and error.args[0].get("converged") is False
246 ):
247 result["diagnostics"] = plain(error.args[0])
248 result["reason"] = (
249 f"{LABELS[slug]} did not converge after "
250 f"{error.args[0].get('iterations')} iterations: "
251 f"{error.args[0].get('termination_reason')}"
252 )
253 save_json(folder / f"{slug}-failure.json", result)
254 return result
255
256 def _fit(self, slug, data, params, folder, xs=()):
257 if slug == "csdid":
258 from csdid_model import fit_csdid
259
260 return fit_csdid(
261 data,
262 folder,
263 params
264 | {
265 "frequency_weights": self.base.config.get(
266 "frequency_weights", False
267 )
268 },
269 xs=xs,
270 )
271 if slug in ["sdid", "matrix-completion", "horizontal-ridge"]:
272 from matrix_inputs import prepare_matrix
273 from matrix_models import fit_one
274
275 import crabbymetrics as cm
276
277 factory = {
278 "sdid": cm.SyntheticDID,
279 "matrix-completion": cm.MatrixCompletion,
280 "horizontal-ridge": cm.HorizontalPanelRidge,
281 }[slug]
282 matrix = prepare_matrix(data)
283 matrix.save(folder)
284 return fit_one(
285 LABELS[slug],
286 factory(**params),
287 params,
288 matrix.y,
289 matrix.w,
290 folder,
291 matrix=matrix,
292 )
293 data.to_csv(folder / "panel-local.csv", index=False)
294 save_json(folder / f"{slug}-parameters.json", params)
295 if slug.startswith("multisynth"):
296 command = [
297 "Rscript",
298 str(HERE / "multisynth_model.R"),
299 str(folder),
300 str(folder / f"{slug}-parameters.json"),
301 str(HERE),
302 ]
303 else:
304 if not self.options.get("mtgpax_repo"):
305 raise ValueError("Configure mtgpax_repo for optional MTGP/TROP")
306 repo = Path(self.options["mtgpax_repo"])
307 python = Path(
308 self.options.get("mtgpax_python", str(repo / ".venv/bin/python"))
309 )
310 if not (repo / "mtgpax").exists() or not python.exists():
311 raise FileNotFoundError(
312 "Configure mtgpax_repo and mtgpax_python for optional MTGP/TROP"
313 )
314 command = [
315 str(python),
316 str(HERE / "mtgpax_worker.py"),
317 str(repo),
318 str(folder),
319 slug,
320 ]
321 started = time.perf_counter()
322 process = subprocess.run(
323 command,
324 check=False,
325 cwd=self.root,
326 text=True,
327 capture_output=True,
328 timeout=self.options.get("timeout_seconds", 3600),
329 )
330 (folder / f"{slug}.log").write_text(process.stdout + process.stderr)
331 if process.returncode:
332 raise RuntimeError(f"{slug} failed; see {slug}.log")
333 result = json.loads((folder / f"{slug}.json").read_text())
334 result["seconds"] = time.perf_counter() - started
335 return result
336
337 def rows(self):
338 rows = []
339 for key, spec in self.base.results["specifications"].items():
340 for result in spec.get("estimates", []) + self.extra.get(key, []):
341 name = result["method"]
342 ok = result.get("status", "estimated") == "estimated"
343 path = result.get(
344 "event_study", spec.get("event_studies", {}).get(name)
345 )
346 rows.append(
347 {
348 "Outcome / treatment": key,
349 "Estimator": name,
350 "Estimate (SDs)": result.get("estimate"),
351 "SE": result.get("se"),
352 "Aggregation": "Coefficient"
353 if name == "Static TWFE"
354 else "Equal post-event times"
355 if name == "CS-DID"
356 else self.base.config["post_weighting"]
357 if name in ["Dynamic TWFE", "BJS / did2s", "Sun–Abraham"]
358 else "Treated cells",
359 "Event study": "Yes" if path else "No",
360 "Status": "estimated" if ok else result["status"],
361 "Note": result.get("reason", ""),
362 }
363 )
364 if name == "CS-DID" and ok:
365 simple = next(
366 r
367 for r in result["aggregations"]["summaries"]
368 if r.get("matches_treated_unit_year_target", False)
369 )
370 rows.append(
371 rows[-1]
372 | {
373 "Estimator": "CS-DID (treated cells)",
374 "Estimate (SDs)": simple["estimate"],
375 "SE": simple["se"],
376 "Aggregation": "Treated cells",
377 "Event study": "Same CS-DID path",
378 }
379 )
380 return rows
381
382 def table(self):
383 markdown_table(self.rows())
384
385 def event_plots(self):
386 window = self.options.get("display_window", self.base.config["display_window"])
387 for key, spec in self.base.results["specifications"].items():
388 series = list(spec.get("event_studies", {}).items())
389 series += [
390 (r["method"], r["event_study"])
391 for r in self.extra.get(key, [])
392 if r.get("status") == "estimated" and r.get("event_study")
393 ]
394 for name, path in series:
395 frame = pd.DataFrame(path).sort_values("event_time")
396 offset = 0.0
397 plot = frame.loc[frame.event_time.between(*window)]
398 fig, ax = plt.subplots(figsize=(10, 4), layout="constrained")
399 ax.plot(
400 plot.event_time, plot.estimate - offset, marker="o", markersize=3
401 )
402 if {"lower", "upper"} <= set(plot.columns):
403 ax.fill_between(
404 plot.event_time,
405 plot.lower - offset,
406 plot.upper - offset,
407 alpha=0.15,
408 )
409 ax.axhline(0, color="gray", lw=0.7)
410 ax.axvline(-0.5, color="gray", ls=":")
411 ax.set(
412 title=f"{key}: {name}",
413 xlabel="Periods since adoption",
414 ylabel="Effect (SDs)",
415 xlim=window,
416 )
417 plt.show()
418 plt.close(fig)
419
420 def tuning_table(self):
421 rows = []
422 for key, models in self.extra.items():
423 for model in models:
424 rows.append(
425 {
426 "Specification": key,
427 "Estimator": model["method"],
428 "Settings": json.dumps(model.get("parameters", {})),
429 "Pooling": json.dumps(model.get("partial_pooling", {})),
430 }
431 )
432 markdown_table(rows)
433
434 def export(self):
435 result = {
436 "schema": "expensive-replication-v1",
437 "config": self.base.config,
438 "options": self.options,
439 "source": self.base.results["source"],
440 "regressions": self.base.results["specifications"],
441 "alternative_models": self.extra,
442 "combined_table": self.rows(),
443 "adapter_sha256": {
444 p.name: digest(p) for p in HERE.iterdir() if p.suffix in {".py", ".R"}
445 },
446 }
447 path = self.output / "comparison.json"
448 save_json(path, clean(result))
449 return path
SHA-256: f0eb686ff3f13e89c8f3d85f49890601c4a7f9a611a2938753be599b8ceac9a6
1 """Capabilities of individual adapters, not a shared intersection of methods."""
2
3 import numpy as np
4 from matrix_inputs import CRABBY_MATRIX_METHODS
5
6
7 class UnidentifiedTarget(ValueError):
8 """The requested target cannot be recovered without changing its support."""
9
10
11 def missing_support_error(method, reason):
12 """Recognize explicit branch-wheel support errors, not solver failures."""
13 if method not in CRABBY_MATRIX_METHODS:
14 return False
15 return any(
16 message in reason
17 for message in (
18 "no observed untreated outcomes",
19 "observed untreated unit-period graph is disconnected",
20 "no observed donor/treated pre-period overlap",
21 "no observed treated post outcomes",
22 "no observed positive-weight donors",
23 "no positive donor weight for a pre-period contrast",
24 "no donor/pre-period overlap for observed treated outcomes",
25 )
26 )
27
28
29 def eligibility(method, data, covariates, parameters=None, config=None):
30 """Return an adapter-specific reason, or None to attempt this estimator."""
31 parameters, config = parameters or {}, config or {}
32 standard_fe = (
33 data.groupby("unit").unit_fe.nunique().max() == 1
34 and data.unit_fe.nunique() == data.unit.nunique()
35 and data.groupby("time").time_fe.nunique().max() == 1
36 and data.time_fe.nunique() == data.time.nunique()
37 )
38 if not standard_fe:
39 return f"{method}: matching the configured nonstandard fixed effects is not implemented."
40 repeated = data.duplicated(["unit", "time"]).any()
41 if method == "csdid":
42 if config.get("frequency_weights") and covariates:
43 return "CS-DID frequency-compressed cluster inference with covariates is not implemented."
44 # RC inference is observation-level; the configured cluster retains
45 # dependence across repeated observations of the treatment-level unit.
46 if not repeated and data.groupby("unit").cluster.nunique().max() > 1:
47 return "CS-DID panel influence functions require a stable cluster ID within unit."
48 return None
49 if covariates:
50 return f"{method}: this adapter does not implement covariate adjustment; controls retained."
51 if config.get("frequency_weights") or (
52 method not in CRABBY_MATRIX_METHODS and data.weight.nunique() != 1
53 ):
54 return f"{method}: this adapter does not implement the configured observation weights."
55 if repeated and method not in CRABBY_MATRIX_METHODS:
56 return f"{method}: requires one observation per unit-period; no observations collapsed."
57 cells = len(data[["unit", "time"]].drop_duplicates())
58 if (
59 method not in CRABBY_MATRIX_METHODS
60 and cells != data.unit.nunique() * data.time.nunique()
61 ):
62 n, t = data.unit.nunique(), data.time.nunique()
63 return f"{method}: requires a complete finite outcome matrix after unit-period aggregation; {n} x {t} grid has {n * t - cells} missing cells, which are not zero-filled."
64 periods = np.sort(data.time.unique())
65 if not np.array_equal(periods, np.arange(1, len(periods) + 1)):
66 return (
67 f"{method}: event-path adapter requires a consecutive prepared period grid."
68 )
69 if not data.g.eq(0).any():
70 return f"{method}: the current donor adapter requires never-treated units."
71 if method.startswith("multisynth"):
72 if data.groupby("unit").cluster.nunique().max() > 1:
73 return "Multisynth bootstrap requires a stable cluster ID within unit."
74 pre = data.loc[data.g.gt(0) & data.time.lt(data.g)].groupby("unit").size()
75 if pre.min() < 2:
76 return "Fixed-effect multisynth requires at least two pre-treatment outcomes per treated unit."
77 return None
SHA-256: f9dad11ea2f0aa4206fe75af045d25e2fe0b23b8c1f4f99faf603e570aafd466
1 """Shared estimation and reporting routines for the editable Quarto template."""
2
3 from __future__ import annotations
4
5 import hashlib
6 import json
7 import os
8 import subprocess
9 import tempfile
10 from importlib.metadata import version
11 from pathlib import Path
12 from unittest.mock import patch
13
14 import matplotlib.pyplot as plt
15 import numpy as np
16 import pandas as pd
17 import polars as pl
18 import pyfixest as pf
19 from frequency import sandwich as frequency_sandwich
20 from IPython.display import Markdown, display
21 from pyfixest.did import visualize
22 from scipy import stats
23 from scipy.sparse import coo_matrix
24 from scipy.sparse.csgraph import connected_components
25 from sparse_did2s import covariance as sparse_covariance
26
27 METHODS = ["Static TWFE", "Dynamic TWFE", "BJS / did2s", "Sun–Abraham"]
28
29
30 def text(value):
31 display(Markdown(value))
32
33
34 def table(value):
35 display(Markdown(pd.DataFrame(value).to_markdown(index=False, floatfmt=".5g")))
36
37
38 def clean(value):
39 if isinstance(value, dict):
40 return {str(k): clean(v) for k, v in value.items()}
41 if isinstance(value, (list, tuple, np.ndarray)):
42 return [clean(x) for x in value]
43 if isinstance(value, (np.integer, np.bool_)):
44 return value.item()
45 if isinstance(value, (float, np.floating)):
46 return float(value) if np.isfinite(value) else None
47 return value
48
49
50 def save_json(path, value):
51 Path(path).write_text(json.dumps(clean(value), indent=2, allow_nan=False) + "\n")
52
53
54 def linear(fit, r, df):
55 estimate = float(r @ fit._beta_hat)
56 variance = float(r @ fit._vcov @ r)
57 if not np.isfinite(variance) or variance < -1e-9:
58 raise ValueError("Invalid contrast covariance")
59 se = np.sqrt(max(0, variance))
60 critical = stats.t.ppf(0.975, df)
61 return {
62 "estimate": estimate,
63 "se": se,
64 "pvalue": float(2 * stats.t.sf(abs(estimate / se), df)) if se else None,
65 "lower": estimate - critical * se,
66 "upper": estimate + critical * se,
67 }
68
69
70 def keep_first_stage_singletons(*args, **kwargs):
71 # One untreated observation identifies a unit intercept for imputation.
72 # did2s does not expose feols' singleton option, so scope it to this call.
73 kwargs.setdefault("fixef_rm", "none")
74 return pf.feols(*args, **kwargs)
75
76
77 class Replication:
78 def __init__(self, config, project):
79 self.root = Path(project)
80 self.config = {
81 "outcomes": ["y"],
82 "treatments": ["d"],
83 "unit": "unit",
84 "time": "time",
85 "cluster": "cluster",
86 "covariates": [],
87 "baseline_by_time": [],
88 "weight": None,
89 "period_step": 1,
90 "period_values": None,
91 "display_window": [-10, 10],
92 "seed": 20261008,
93 "panelview_max_units": 100,
94 "test_mode": "classical",
95 "test_horizon": None,
96 "post_horizon": None,
97 "post_weighting": "observations",
98 "sample_query": None,
99 "reversal_policy": "error",
100 "labels": {},
101 "allow_repeated_cells": False,
102 "time_fe": None,
103 "unit_fe": None,
104 "frequency_weights": False,
105 "source_sd": None,
106 "require_never_control_periods": False,
107 "latest_cohort_control": False,
108 } | config
109 self.path = self.root / self.config["data"]
110 self.output = self.root / self.config["output_dir"]
111 self.output.mkdir(parents=True, exist_ok=True)
112 if self.path.suffix == ".csv":
113 self.df = pl.read_csv(self.path)
114 elif self.path.suffix in {".parquet", ".pq"}:
115 self.df = pl.read_parquet(self.path)
116 elif self.path.suffix == ".dta":
117 self.df = pl.from_pandas(
118 pd.read_stata(self.path, convert_categoricals=False)
119 )
120 else:
121 raise ValueError("Input must be CSV, Parquet or Stata")
122 self._run_code_hashes = self.code_hashes()
123 self.results = {
124 "schema": "did-replication-v2",
125 "config": self.config,
126 "source": {
127 "file": self.path.name,
128 "bytes": self.path.stat().st_size,
129 "sha256": hashlib.file_digest(
130 self.path.open("rb"), "sha256"
131 ).hexdigest(),
132 },
133 "versions": {
134 p: version(p)
135 for p in ["pyfixest", "polars", "duckreg", "numpy", "pandas"]
136 },
137 "specifications": {},
138 }
139 self.panels, self.fits = {}, {}
140
141 def label(self, name):
142 return self.config["labels"].get(name, name)
143
144 def describe(self):
145 c = self.config
146 text(
147 f"**{c['study']}** — {c['citation']}\n\n{c['description']}\n\n"
148 f"Archive specification: `{c['source_specification_id']}`. "
149 f"Input: {self.path.stat().st_size / 1e6:.3f} MB."
150 )
151 rows = [("Y", x, self.label(x)) for x in c["outcomes"]]
152 rows += [
153 (
154 "W",
155 x,
156 f"{self.label(x)}; values {sorted(self.df[x].unique().to_list())}; cardinality {self.df[x].n_unique()}",
157 )
158 for x in c["treatments"]
159 ]
160 rows += [
161 ("i", c["unit"], self.label(c["unit"])),
162 ("t", c["time"], self.label(c["time"])),
163 (
164 "X_it",
165 ", ".join(
166 c["covariates"] + [x + " × period" for x in c["baseline_by_time"]]
167 )
168 or "None",
169 "Plus unit and period fixed effects",
170 ),
171 (
172 "Cluster",
173 c["cluster"] or c["unit"],
174 self.label(c["cluster"] or c["unit"]),
175 ),
176 ]
177 table(pd.DataFrame(rows, columns=["Role", "Column", "Definition"]))
178
179 def prepare(self):
180 c = self.config
181 raw = self.df.to_pandas()
182 # Peter's scale is the unweighted SD of the entire prepared source sample.
183 scales = {y: float(raw[y].std(ddof=1)) for y in c["outcomes"]}
184 if c["source_sd"] is not None:
185 scales.update(c["source_sd"])
186 if c["sample_query"]:
187 raw = raw.query(c["sample_query"]).copy()
188 u, t, cl = c["unit"], c["time"], c["cluster"] or c["unit"]
189 assert not raw[[u, t, cl]].isna().any().any()
190 if not c["allow_repeated_cells"]:
191 assert not raw.duplicated([u, t]).any(), "Input must be a unit-period panel"
192 assert raw.groupby([u, t])[c["treatments"]].nunique().max().max() <= 1
193 if c["period_values"] is not None:
194 idx = raw[t].map({v: j + 1 for j, v in enumerate(c["period_values"])})
195 assert idx.notna().all()
196 else:
197 idx = (raw[t] - raw[t].min()) / c["period_step"] + 1
198 assert np.allclose(idx, np.round(idx)), "Supply the complete period grid"
199 self.period_labels = dict(zip(idx.astype(int), raw[t]))
200 samples = []
201 for w in c["treatments"]:
202 base = pd.DataFrame(
203 {
204 "unit": pd.factorize(raw[u], sort=True)[0] + 1,
205 "time": idx.astype(int),
206 "d": raw[w].astype("int64"),
207 "cluster": pd.factorize(raw[cl], sort=True)[0] + 1,
208 },
209 index=raw.index,
210 )
211 assert base.d.notna().all() and set(base.d.unique()) <= {0, 1}
212 base["unit_fe"] = (
213 pd.factorize(raw[c["unit_fe"]], sort=True)[0] + 1
214 if c["unit_fe"]
215 else base.unit
216 )
217 base["time_fe"] = (
218 pd.factorize(raw[c["time_fe"]], sort=True)[0] + 1
219 if c["time_fe"]
220 else base.time
221 )
222 base = base.sort_values(["unit", "time"])
223 base["g"] = (
224 base.time.where(base.d.eq(1))
225 .groupby(base.unit)
226 .transform("min")
227 .fillna(0)
228 .astype(int)
229 )
230 latest_control = None
231 if not base.g.eq(0).any() and c["latest_cohort_control"]:
232 latest_control = int(base.g.max())
233 base = base.loc[base.time.lt(latest_control)].copy()
234 base.loc[base.g.eq(latest_control), "g"] = 0
235 base["event_time"] = (base.time - base.g).where(base.g.gt(0))
236 base["ever"] = base.g.gt(0).astype(int)
237 base["relative_time"] = base.event_time.fillna(-1).astype(int)
238 reversals = base.g.gt(0) & base.event_time.ge(0) & base.d.eq(0)
239 reversal_count = int(reversals.sum())
240 if reversal_count:
241 assert c["reversal_policy"] == "drop_post_untreated", (
242 "Nonabsorbing treatment"
243 )
244 base = base.loc[~reversals].copy()
245 always = base.loc[
246 base.g.gt(0) & base.g.le(base.groupby("unit").time.transform("min")),
247 "unit",
248 ].unique()
249 base = base.loc[~base.unit.isin(always)].copy()
250 xcols = []
251 for j, x in enumerate(c["covariates"]):
252 name = f"x{j}"
253 base[name] = pd.to_numeric(raw.loc[base.index, x])
254 xcols.append(name)
255 for x in c["baseline_by_time"]:
256 values = pd.to_numeric(raw.loc[base.index, x])
257 assert values.groupby(base.unit).nunique().max() <= 1, (
258 "Not a unit baseline"
259 )
260 values = (values - values.mean()) / values.std(ddof=1)
261 for period in sorted(base.time.unique())[1:]:
262 name = f"x{len(xcols)}"
263 base[name] = values * base.time.eq(period)
264 xcols.append(name)
265 base["weight"] = raw.loc[base.index, c["weight"]] if c["weight"] else 1.0
266 for y in c["outcomes"]:
267 assert np.isfinite(scales[y]) and scales[y] > 0
268 data = base.assign(y=raw.loc[base.index, y] / scales[y])
269 data = data.replace([np.inf, -np.inf], np.nan).dropna(
270 subset=["y", "weight", *xcols]
271 )
272 assert data.weight.gt(0).all()
273 data = data.loc[
274 data.g.eq(0) | data.g.gt(data.groupby("unit").time.transform("min"))
275 ].copy()
276 baseline_units = set(data.loc[data.event_time.eq(-1), "unit"])
277 baseline_excluded = data.loc[
278 data.g.gt(0) & ~data.unit.isin(baseline_units), "unit"
279 ].nunique()
280 data = data.loc[data.g.eq(0) | data.unit.isin(baseline_units)].copy()
281 # Remove recursively uninformative unit/period fixed-effect singletons.
282 # Recheck pre-treatment support after these removals.
283 while True:
284 before = len(data)
285 data = data.loc[
286 data.groupby("unit").unit.transform("size").gt(1)
287 & data.groupby("time_fe").time_fe.transform("size").gt(1)
288 ].copy()
289 data = data.loc[
290 data.g.eq(0)
291 | data.g.gt(data.groupby("unit").time.transform("min"))
292 ].copy()
293 if len(data) == before:
294 break
295 never_support_removed = 0
296 if c["require_never_control_periods"]:
297 periods = data.loc[data.g.eq(0), "time_fe"].unique()
298 before = len(data)
299 data = data.loc[data.time_fe.isin(periods)].copy()
300 never_support_removed = before - len(data)
301 # Predictions require untreated unit and time FE in one connected component.
302 # Iteratively retain a common counterfactual-supported sample.
303 unsupported_count = 0
304 for _ in range(20):
305 old_n = len(data)
306 data = data.loc[
307 data.groupby("unit_fe").unit_fe.transform("size").gt(1)
308 & data.groupby("time_fe").time_fe.transform("size").gt(1)
309 ].copy()
310 ui = pd.Categorical(data.unit_fe).codes.astype(np.int64)
311 ti = pd.Categorical(data.time_fe).codes.astype(np.int64)
312 ni, nt = ui.max() + 1, ti.max() + 1
313 zero = data.d.eq(0).to_numpy()
314 graph = coo_matrix(
315 (np.ones(zero.sum()), (ui[zero], ni + ti[zero])),
316 shape=(ni + nt, ni + nt),
317 )
318 _, comp = connected_components(graph, directed=False)
319 support = comp[ui] == comp[ni + ti]
320 unsupported_count += int((~support).sum())
321 data = data.loc[support].copy()
322 observed_g = (
323 data.time.where(data.d.eq(1))
324 .groupby(data.unit)
325 .transform("min")
326 .fillna(0)
327 .astype(int)
328 )
329 invalid = data.loc[observed_g.ne(data.g), "unit"].unique()
330 valid_baseline = data.loc[data.event_time.eq(-1), "unit"].unique()
331 data = data.loc[
332 ~data.unit.isin(invalid)
333 & (data.g.eq(0) | data.unit.isin(valid_baseline))
334 ].copy()
335 if len(data) == old_n:
336 break
337 recovered = (
338 data.time.where(data.d.eq(1))
339 .groupby(data.unit)
340 .transform("min")
341 .fillna(0)
342 .astype(int)
343 )
344 assert recovered.equals(data.g), (
345 "Missing outcomes obscure the original adoption date"
346 )
347 assert data.g.eq(0).any() and data.loc[data.g.gt(0), "g"].nunique() >= 1
348 assert data.groupby("unit").size().min() >= 2
349 data = data.reset_index(drop=True)
350 if c["frequency_weights"]:
351 data["__frequency"] = data.weight
352 key = f"{y}__{w}"
353 self.panels[key] = (data, xcols)
354 units = data[["unit", "g"]].drop_duplicates()
355 sample = {
356 "outcome": y,
357 "treatment": w,
358 "source_rows": self.df.height,
359 "source_sd": scales[y],
360 "rows": len(data),
361 "units": len(units),
362 "clusters": data.cluster.nunique(),
363 "treated_units": int(units.g.gt(0).sum()),
364 "never_treated_units": int(units.g.eq(0).sum()),
365 "treated_cohorts": units.loc[units.g.gt(0), "g"].nunique(),
366 "latest_cohort_control_index": latest_control,
367 "unique_unit_period_cells": len(
368 data[["unit", "time"]].drop_duplicates()
369 ),
370 "unsupported_prediction_rows_removed": unsupported_count,
371 "rows_without_never_control_time_fe_removed": never_support_removed,
372 "initially_treated_units_removed": len(always),
373 "post_adoption_untreated_rows_removed": reversal_count,
374 "units_without_minus_one_removed": baseline_excluded,
375 "additional_rows_removed": len(base) - len(data),
376 }
377 self.results["specifications"][key] = {"sample": sample}
378 samples.append(sample)
379 return pd.DataFrame(samples)
380
381 def panelview(self):
382 for key, (data, _) in self.panels.items():
383 n = min(self.config["panelview_max_units"], data.unit.nunique())
384 chosen = np.random.default_rng(self.config["seed"]).choice(
385 sorted(data.unit.unique()), n, replace=False
386 )
387 fig, ax = plt.subplots(figsize=(9, 4))
388 plot_quilt = visualize._plot_panelview
389
390 def numeric_quilt(treatment_quilt, plot_quilt=plot_quilt, **kwargs):
391 # Boolean pivots with missing cells otherwise become object arrays.
392 return plot_quilt(treatment_quilt.astype(float), **kwargs)
393
394 with patch("pyfixest.did.visualize._plot_panelview", numeric_quilt):
395 pf.panelview(
396 data.loc[data.unit.isin(chosen)].drop_duplicates(["unit", "time"]),
397 unit="unit",
398 time="time",
399 treat="d",
400 sort_by_timing=True,
401 noticks=True,
402 ax=ax,
403 title="Treatment adoption",
404 )
405 periods = sorted(data.time.unique())
406 positions = np.unique(
407 np.linspace(0, len(periods) - 1, min(7, len(periods)), dtype=int)
408 )
409 ax.set_xticks(
410 positions, [str(self.period_labels[periods[k]]) for k in positions]
411 )
412 self._show(fig)
413 text(
414 f"{key}: {n} of {data.unit.nunique()} units shown (seed {self.config['seed']}). "
415 "All eligible units enter estimation. Purple: untreated; yellow: treated; white: unobserved."
416 )
417
418 def estimate(self):
419 for key, (data, xs) in self.panels.items():
420 print(f"Fitting {key}: {len(data):,} rows", flush=True)
421 extra = (" + " + " + ".join(xs)) if xs else ""
422 args = {
423 "data": data,
424 "vcov": {"CRV1": "cluster"},
425 "weights": "weight" if self.config["weight"] else None,
426 }
427 static = pf.feols("y ~ d" + extra + " | unit_fe + time_fe", **args)
428 dynamic = pf.feols(
429 "y ~ i(relative_time, ever, ref=-1)" + extra + " | unit_fe + time_fe",
430 **args,
431 )
432 # Explicit second-stage indicators: never-treated rows are zero in every column.
433 events = sorted(data.event_time.dropna().astype(int).unique())
434 columns = {
435 e: f"es_{'m' + str(-e) if e < 0 else 'p' + str(e)}" for e in events
436 }
437 indicators = pd.DataFrame(
438 {
439 name: data.event_time.eq(e).astype(float)
440 for e, name in columns.items()
441 }
442 )
443 with (
444 patch("pyfixest.did.did2s.feols", keep_first_stage_singletons),
445 patch("pyfixest.did.estimation._did2s_vcov", sparse_covariance),
446 ):
447 did = pf.did2s(
448 pd.concat([data, indicators], axis=1),
449 yname="y",
450 first_stage="~ "
451 + (" + ".join(xs) if xs else "0")
452 + " | unit_fe + time_fe",
453 second_stage="~ " + " + ".join(columns.values()),
454 treatment="d",
455 cluster="cluster",
456 weights=args["weights"],
457 )
458 sa = pf.event_study(
459 data=data,
460 yname="y",
461 idname="unit",
462 tname="time",
463 gname="g",
464 cluster="cluster",
465 estimator="saturated",
466 att=False,
467 )
468 if (
469 xs
470 or self.config["weight"]
471 or self.config["time_fe"]
472 or self.config["unit_fe"]
473 ):
474 lhs, _fe = sa._fml.split("|")
475 sa = pf.feols(
476 lhs + extra + " | unit_fe + time_fe",
477 data=sa._data,
478 vcov={"CRV1": "cluster"},
479 weights=args["weights"],
480 )
481 assert all(m._N == len(data) for m in [static, dynamic, did, sa])
482 if self.config["frequency_weights"]:
483 for model in (static, dynamic, sa):
484 model._vcov = frequency_sandwich(model, data)
485 fits = dict(zip(METHODS, [static, dynamic, did, sa]))
486 self.fits[key] = fits
487 self._summarize(key, data, fits, columns)
488
489 def _summarize(self, key, data, fits, columns):
490 df = (
491 data.groupby("cluster").weight.first().sum()
492 if self.config["frequency_weights"]
493 else data.cluster.nunique()
494 ) - 1
495 events = sorted(data.event_time.dropna().astype(int).unique())
496 mass = data.groupby(["g", "event_time"]).weight.sum()
497 sa = fits["Sun–Abraham"]
498 parts = (
499 pd.Series(sa.coef().index)
500 .str.extract(
501 r"^rel_time::(?P<e>-?\d+(?:\.\d+)?):first_treated_period::(?P<g>\d+)$"
502 )
503 .astype(float)
504 )
505 sa_rows = {}
506 for e in sorted(parts.e.dropna().unique()):
507 hit = parts.e.eq(e)
508 r = np.zeros(len(parts))
509 r[hit] = [mass.loc[(int(g), e)] for g in parts.loc[hit, "g"]]
510 sa_rows[int(e)] = r / r.sum()
511 package_aggregation_gap = None
512 if not (
513 self.config["covariates"]
514 or self.config["baseline_by_time"]
515 or self.config["weight"]
516 or self.config["time_fe"]
517 or self.config["unit_fe"]
518 ):
519 official = sa.aggregate(agg="period", weighting="shares").astype(float)
520 expected = np.array(
521 [linear(sa, r, df)["estimate"] for r in sa_rows.values()]
522 )
523 package_aggregation_gap = float(
524 np.max(np.abs(expected - official["Estimate"].to_numpy()))
525 )
526 assert package_aggregation_gap < 1e-7, (
527 "Saturated aggregation needs a support review"
528 )
529 # Renormalize over estimated cohort cells; omitted columns get no target mass.
530 dynamic = fits["Dynamic TWFE"]
531 dyn_events = (
532 pd.Series(dynamic.coef().index)
533 .str.extract(r"^relative_time::(-?\d+(?:\.\d+)?)", expand=False)
534 .astype(float)
535 )
536 matrices = {
537 "Dynamic TWFE": {
538 e: dyn_events.eq(e).astype(float).to_numpy()
539 for e in events
540 if dyn_events.eq(e).any()
541 },
542 "BJS / did2s": {
543 e: (fits["BJS / did2s"].coef().index == name).astype(float)
544 for e, name in columns.items()
545 },
546 "Sun–Abraham": sa_rows,
547 }
548 post = data.loc[data.d.eq(1)].groupby("event_time").weight.sum()
549 if self.config["post_horizon"] is not None:
550 post = post.loc[post.index <= self.config["post_horizon"]]
551 common = sorted(
552 set(post.index).intersection(*(set(m) for m in matrices.values()))
553 )
554 assert common
555 weights = post.reindex(common).to_numpy(dtype=float, copy=True)
556 if self.config["post_weighting"] == "equal_periods":
557 weights = np.ones_like(weights)
558 else:
559 assert self.config["post_weighting"] == "observations"
560 weights /= weights.sum()
561 pooled = [
562 dict(
563 method="Static TWFE",
564 **linear(
565 fits["Static TWFE"],
566 (fits["Static TWFE"].coef().index == "d").astype(float),
567 df,
568 ),
569 )
570 ]
571 paths = {}
572 for method, matrix in matrices.items():
573 model = fits[method]
574 r = sum(w * matrix[e] for e, w in zip(common, weights))
575 pooled.append(dict(method=method, **linear(model, r, df)))
576 paths[method] = [
577 dict(event_time=e, **linear(model, r, df)) for e, r in matrix.items()
578 ]
579 # Independent imputation check of every post-treatment did2s coefficient.
580 xs = self.panels[key][1]
581 first_formula = (
582 "y ~ " + (" + ".join(xs) if xs else "0") + " | unit_fe + time_fe"
583 )
584 first = pf.feols(
585 first_formula,
586 data=data.loc[data.d.eq(0)],
587 weights="weight" if self.config["weight"] else None,
588 fixef_tol=1e-10,
589 fixef_rm="none",
590 )
591 first.fixef()
592 residual = data.y.to_numpy() - first.predict(newdata=data)
593 mass_by_event = data.groupby("event_time").weight.sum()
594 weighted = (
595 pd.Series(residual * data.weight.to_numpy()).groupby(data.event_time).sum()
596 )
597 bjs = weighted / mass_by_event
598 check = max(
599 abs(row["estimate"] - bjs.loc[row["event_time"]])
600 for row in paths["BJS / did2s"]
601 if row["event_time"] >= 0
602 )
603 assert check < 1e-5, f"did2s/imputation discrepancy: {check}"
604 output = self.results["specifications"][key]
605 output.update(
606 estimates=pooled,
607 event_studies=paths,
608 target={
609 "event_times": common,
610 "weights": weights.tolist(),
611 "weighting": self.config["post_weighting"],
612 },
613 checks={
614 "did2s_vs_direct_imputation_max_post_difference": check,
615 "package_aggregate_vs_estimable_cell_weights": package_aggregation_gap,
616 },
617 inference={
618 "standard_errors": "Cluster-robust; did2s uses two-stage covariance",
619 "confidence_intervals": "Pointwise t, clusters minus one degrees of freedom",
620 },
621 )
622 pd.DataFrame(pooled).to_csv(self.output / f"{key}-estimates.csv", index=False)
623
624 def estimate_table(self):
625 for key, result in self.results["specifications"].items():
626 text(f"**{key}**")
627 table(
628 [
629 x
630 for x in result["estimates"]
631 if x["method"] not in result.get("method_status", {})
632 ]
633 )
634 target = result["target"]
635 text(
636 f"Post-treatment event times {min(target['event_times']):g}–{max(target['event_times']):g}; "
637 f"{target['weighting'].replace('_', ' ')} weights. All estimates are in source-outcome SD units."
638 )
639
640 def event_plots(self):
641 for result in self.results["specifications"].values():
642 static = result["estimates"][0]
643 for method, path in result["event_studies"].items():
644 if method in result.get("method_status", {}):
645 text(result["method_status"][method])
646 continue
647 frame = pd.DataFrame(path).sort_values("event_time")
648 lo, hi = self.config["display_window"]
649 frame = frame.loc[frame.event_time.between(lo, hi)]
650 fig, ax = plt.subplots(figsize=(8, 3.6))
651 ax.errorbar(
652 frame.event_time,
653 frame.estimate,
654 yerr=[frame.estimate - frame.lower, frame.upper - frame.estimate],
655 fmt="o",
656 ms=3,
657 capsize=2,
658 color="#147d87",
659 )
660 if method == "Dynamic TWFE":
661 ax.scatter(
662 [-1], [0], facecolor="white", edgecolor="#147d87", zorder=4
663 )
664 ax.axhline(static["estimate"], color="#a15236", label="Static TWFE")
665 ax.axhspan(
666 static["lower"], static["upper"], alpha=0.08, color="#a15236"
667 )
668 ax.legend()
669 elif method == "Sun–Abraham":
670 ax.scatter(
671 [-1], [0], facecolor="white", edgecolor="#147d87", zorder=4
672 )
673 ax.axhline(0, color="gray", lw=0.7)
674 ax.axvline(-0.5, color="gray", ls="--", lw=0.7)
675 ax.set(
676 title=method,
677 xlabel="Periods relative to adoption",
678 ylabel="Outcome SD units",
679 )
680 self._show(fig)
681
682 def f_tests(self):
683 for key, (data, xs) in self.panels.items():
684 if data.loc[data.g.gt(0), "g"].nunique() < 2:
685 self.results["specifications"][key]["f_test"] = {
686 "status": "not_identified",
687 "pvalue": None,
688 "reason": "Only one treated cohort remains on the common supported sample.",
689 }
690 continue
691 outfile = self.output / f"{key}-f-test.json"
692 with tempfile.TemporaryDirectory() as td:
693 input_path = Path(td) / "panel.csv"
694 data.to_csv(input_path, index=False)
695 conf = {
696 "reference_repo": str(self.root / "TestingInEventStudies"),
697 "data": str(input_path),
698 "unit": "unit",
699 "time": "time",
700 "outcome": "y",
701 "treatment": "d",
702 "cluster": "cluster",
703 "fixed_effect_time": "time_fe",
704 "fixed_effect_unit": "unit_fe",
705 "frequency_weights": self.config["frequency_weights"],
706 "covariates": xs,
707 "weight": "weight" if self.config["weight"] else None,
708 "test_mode": self.config["test_mode"],
709 "test_horizon": self.config["test_horizon"],
710 "output": str(outfile),
711 }
712 result = subprocess.run(
713 [
714 str(self.root / "environments/reference/.venv/bin/python"),
715 str(Path(__file__).with_name("reference_f_test.py")),
716 json.dumps(conf),
717 ],
718 env=dict(
719 os.environ,
720 MPLBACKEND="Agg",
721 OPENBLAS_NUM_THREADS="1",
722 OMP_NUM_THREADS="1",
723 NUMBA_NUM_THREADS="2",
724 PYTHONDONTWRITEBYTECODE="1",
725 ),
726 check=False,
727 capture_output=True,
728 text=True,
729 )
730 if result.returncode:
731 if "No post-treatment cohort deviations to test" in result.stderr:
732 self.results["specifications"][key]["f_test"] = {
733 "status": "not_identified",
734 "pvalue": None,
735 "reason": "No identified positive-event-time cohort deviations remain to test.",
736 }
737 continue
738 raise RuntimeError(result.stderr)
739 test = json.loads(outfile.read_text())
740 sa = self.fits[key]["Sun–Abraham"]
741 sa_residual = np.asarray(sa._u_hat).ravel()
742 sa_rss = float(sa_residual @ sa_residual)
743 gap = abs(sa_rss - test["rss_unrestricted"]) / max(1, sa_rss)
744 if gap >= 1e-7:
745 self.results["specifications"][key]["method_status"] = {
746 "Sun–Abraham": "unvalidated: saturated and reference deviations fits disagree"
747 }
748 self.results["specifications"][key]["f_test"] = {
749 "status": "not_identified",
750 "pvalue": None,
751 "reason": "The saturated and imported deviations designs do not span the same fitted model on this sparse panel. Test withheld.",
752 "relative_rss_difference": gap,
753 }
754 continue
755 self.results["specifications"][key]["checks"][
756 "saturated_vs_deviations_relative_rss_difference"
757 ] = gap
758 self.results["specifications"][key]["f_test"] = test
759
760 def f_report(self):
761 for key, result in self.results["specifications"].items():
762 test = result["f_test"]
763 if test.get("status") == "not_identified":
764 text(f"**{key}: homogeneity test not identified.** {test['reason']}")
765 continue
766 text(
767 f"**{key}: F({test['df_num']}, {test['df_den']}) = {test['F']:.4f}; p = {test['pvalue']:.4g}.** "
768 f"{test['inference']}. {test['design_columns']} design columns."
769 )
770 frame = pd.DataFrame(test["deviations"])
771 if frame.empty:
772 continue
773 lo, hi = self.config["display_window"]
774 frame = frame.loc[frame.event_time.between(lo, hi)]
775 fig, ax = plt.subplots(figsize=(8, 3.7))
776 for cohort, subset in frame.groupby("cohort"):
777 subset = subset.sort_values("event_time")
778 ax.plot(subset.event_time, subset.estimate, lw=0.9, alpha=0.65)
779 ax.axhline(0, color="black", lw=0.8)
780 ax.axvline(-0.5, color="gray", ls="--")
781 ax.set(
782 title="Cohort deviations from the reference event-time path",
783 xlabel="Periods relative to adoption",
784 ylabel="Deviation (outcome SD units)",
785 )
786 self._show(fig)
787 text(
788 "Each line is a cohort deviation. The joint test restricts positive-event-time deviations to zero; event time 0 remains unrestricted."
789 )
790
791 def code_hashes(self):
792 return {
793 name: hashlib.sha256(
794 Path(__file__).with_name(name).read_bytes()
795 ).hexdigest()
796 for name in [
797 "replication_engine.py",
798 "reference_f_test.py",
799 "sparse_did2s.py",
800 "frequency.py",
801 "sparse_models.py",
802 ]
803 }
804
805 def load_results(self):
806 cached = json.loads((self.output / "estimates.json").read_text())
807 assert cached["config"] == self.config, "Cached configuration changed"
808 assert cached["source"] == self.results["source"], "Cached source changed"
809 assert cached["versions"] == self.results["versions"], (
810 "Cached environment changed"
811 )
812 assert cached["code_sha256"] == self.code_hashes(), (
813 "Cached implementation changed"
814 )
815 self.results = cached
816
817 def export(self):
818 self.results["code_sha256"] = self._run_code_hashes
819 save_json(self.output / "estimates.json", self.results)
820 return self.output / "estimates.json"
821
822 @staticmethod
823 def _show(fig):
824 fig.tight_layout()
825 plt.show()
826 plt.close(fig)
SHA-256: 81aea1cb43f664d135670feb22a2b2e66f4c50f58a61ea49567bda9ed9f7930f
1 """Sparse least-squares representation of the same event-study regressions.
2
3 The saturated and restricted models are constructed from observed cohort/event
4 cells. Positive-event cells are pooled across cohorts in the restricted model;
5 leads and event zero remain unrestricted. This is a change of basis of the
6 TestingInEventStudies deviations regression, checked against its dense fit.
7 """
8
9 import numpy as np
10 import pandas as pd
11 import scipy.sparse as sp
12 from scipy import stats
13 from scipy.sparse.csgraph import connected_components
14 from scipy.sparse.linalg import lsmr
15
16
17 def solve(a, y):
18 scale = np.sqrt(np.asarray(a.power(2).sum(axis=0)).ravel())
19 assert np.all(scale > 0)
20 b = a.multiply(1 / scale).tocsr()
21 z = lsmr(b, y, atol=1e-10, btol=1e-10, maxiter=6000)
22 if z[1] not in (0, 1, 2, 4, 5):
23 raise ValueError(f"LSMR convergence {z[1]}, iterations {z[2]}")
24 beta = z[0] / scale
25 return (
26 beta,
27 np.asarray(y - a @ beta),
28 {"iterations": int(z[2]), "stop": int(z[1]), "normal_residual": float(z[4])},
29 )
30
31
32 def contrast(a, contrast_weights, u, cluster, df_k):
33 scale = np.sqrt(np.asarray(a.power(2).sum(axis=0)).ravel())
34 b = a.multiply(1 / scale).tocsr()
35 target = contrast_weights / scale
36 z = lsmr(b.T, target, atol=1e-10, btol=1e-10, maxiter=6000)
37 if z[1] not in (0, 1, 2, 4, 5):
38 raise ValueError(f"Contrast solve failed {z[1]}")
39 discrepancy = np.linalg.norm(b.T @ z[0] - target) / max(
40 1e-12, np.linalg.norm(target)
41 )
42 if discrepancy > 1e-6:
43 raise ValueError(f"Nonestimable contrast: {discrepancy}")
44 score = np.bincount(cluster, weights=z[0] * u)
45 g = cluster.max() + 1
46 n = len(u)
47 variance = score @ score * g / (g - 1) * (n - 1) / (n - df_k)
48 return float(np.sqrt(max(0, variance)))
49
50
51 def summary(estimate, se, g):
52 if se is None:
53 return {
54 "estimate": float(estimate),
55 "se": None,
56 "lower": None,
57 "upper": None,
58 "pvalue": None,
59 }
60 crit = stats.t.ppf(0.975, g - 1)
61 return {
62 "estimate": float(estimate),
63 "se": float(se),
64 "lower": float(estimate - crit * se),
65 "upper": float(estimate + crit * se),
66 "pvalue": float(2 * stats.t.sf(abs(estimate / se), g - 1)) if se else None,
67 }
68
69
70 def design(data, cells):
71 # All unit FE levels, all but the first time FE. Never-treated connectivity
72 # and an observed -1 for every treated unit guarantee the saturated rank.
73 n = len(data)
74 ui = pd.Categorical(data.unit_fe).codes.astype(np.int64)
75 ti = pd.Categorical(data.time_fe).codes.astype(np.int64)
76 ni = ui.max() + 1
77 nt = ti.max() + 1
78 rows = [np.arange(n), np.flatnonzero(ti > 0)]
79 cols = [ui, ni + ti[ti > 0] - 1]
80 cols.append(ni + nt - 1 + cells[cells >= 0])
81 rows.append(np.flatnonzero(cells >= 0))
82 a = sp.csr_matrix(
83 (np.ones(sum(map(len, rows))), (np.concatenate(rows), np.concatenate(cols))),
84 shape=(n, int(ni + nt + max(cells.max(), -1))),
85 )
86 return a, int(ni + nt - 1)
87
88
89 def fit(run, key, data, xs):
90 assert not xs and not run.config["weight"], (
91 "Sparse large-panel backend currently supports unweighted models without covariates"
92 )
93 n = len(data)
94 y = data.y.to_numpy()
95 cluster = pd.Categorical(data.cluster).codes.astype(np.int64)
96 g = cluster.max() + 1
97 never = data.g.eq(0)
98 ui = pd.Categorical(data.loc[never, "unit_fe"]).codes.astype(np.int64)
99 ti = pd.Categorical(
100 data.loc[never, "time_fe"], categories=sorted(data.time_fe.unique())
101 ).codes.astype(np.int64)
102 nu = ui.max() + 1
103 nt = data.time_fe.nunique()
104 graph = sp.coo_matrix((np.ones(len(ui)), (ui, nu + ti)), shape=(nu + nt, nu + nt))
105 assert connected_components(graph, directed=False, return_labels=False) == 1, (
106 "Never-treated graph does not identify the full saturated time path"
107 )
108 events = data.event_time.to_numpy()
109 cohorts = data.g.to_numpy()
110 post = data.loc[data.d.eq(1)].groupby("event_time").size()
111 if run.config["post_horizon"] is not None:
112 post = post.loc[post.index <= run.config["post_horizon"]]
113 weights = post.astype(float)
114 if run.config["post_weighting"] == "equal_periods":
115 weights[:] = 1
116 weights /= weights.sum()
117 target = weights.to_dict()
118 window = run.config["display_window"]
119 paths = {}
120 estimates = []
121 checks = {}
122 fits = {}
123 labels_dynamic = sorted(set(events[np.isfinite(events)]) - {-1})
124 lookup = {e: j for j, e in enumerate(labels_dynamic)}
125 dynamic_codes = pd.Series(events).map(lookup).fillna(-1).to_numpy(int)
126 valid = (cohorts > 0) & (events != -1)
127 lo = int(np.nanmin(events))
128 span = int(np.nanmax(events)) - lo + 1
129 packed = (
130 cohorts[valid].astype(np.int64) * span + events[valid].astype(np.int64) - lo
131 )
132 packed_labels, inverse = np.unique(packed, return_inverse=True)
133 sat_codes = np.full(n, -1, dtype=np.int32)
134 sat_codes[valid] = inverse
135 labels_sat = [(int(v // span), int(v % span + lo)) for v in packed_labels]
136 del packed, inverse
137 restricted_labels = [("pooled", e) if e > 0 else (int(c), e) for c, e in labels_sat]
138 mapping, unique_restricted = pd.factorize(
139 pd.Index(restricted_labels, tupleize_cols=False), sort=False
140 )
141 restricted_codes = np.full(n, -1)
142 restricted_codes[valid] = mapping[sat_codes[valid]]
143 unit_nested = data.groupby("unit_fe").cluster.nunique().max() == 1
144 for method, codes, labels in [
145 ("Static TWFE", np.where(data.d.eq(1), 0, -1), [0]),
146 ("Dynamic TWFE", dynamic_codes, labels_dynamic),
147 ("Sun–Abraham", sat_codes, list(labels_sat)),
148 ]:
149 print("SPARSE", key, method, n, len(labels), flush=True)
150 a, offset = design(data, codes)
151 beta, u, check = solve(a, y)
152 checks[method] = check
153 rss = float(u @ u)
154 contrasts = {}
155 if method == "Static TWFE":
156 contrasts[0] = np.array([1.0])
157 elif method == "Dynamic TWFE":
158 for e in labels:
159 contrasts[e] = np.eye(1, len(labels), labels.index(e)).ravel()
160 else:
161 counts = np.bincount(codes[codes >= 0], minlength=len(labels))
162 for e in sorted({e for c, e in labels}):
163 idx = np.array([v == e for c, v in labels])
164 r = counts * idx
165 contrasts[e] = r / r.sum()
166 pooled = (
167 contrasts[0]
168 if method == "Static TWFE"
169 else sum(w * contrasts[e] for e, w in target.items())
170 )
171 contrast_weights = np.r_[np.zeros(offset), pooled]
172 df_k = a.shape[1] - (data.unit_fe.nunique() if unit_nested else 0)
173 se = contrast(a, contrast_weights, u, cluster, df_k)
174 estimates.append(dict(method=method, **summary(contrast_weights @ beta, se, g)))
175 if method != "Static TWFE":
176 path = []
177 for e, r in contrasts.items():
178 # Full coefficient paths are exported; solve displayed intervals.
179 ee = float(r @ beta[offset:])
180 err = None
181 if window[0] <= e <= window[1]:
182 err = contrast(a, np.r_[np.zeros(offset), r], u, cluster, df_k)
183 path.append(dict(event_time=int(e), **summary(ee, err, g)))
184 paths[method] = path
185 fits[method] = {
186 "rss": rss,
187 "rank": a.shape[1],
188 "offset": offset,
189 "beta": beta[offset:],
190 }
191 del a, beta, u
192 print("SPARSE", key, "restricted deviations", flush=True)
193 ar, offset = design(data, restricted_codes)
194 br, ur, check = solve(ar, y)
195 q = len(labels_sat) - len(unique_restricted)
196 df = n - fits["Sun–Abraham"]["rank"]
197 rss_u = fits["Sun–Abraham"]["rss"]
198 rss_r = float(ur @ ur)
199 F = (rss_r - rss_u) / q / (rss_u / df)
200 test = {
201 "F": F,
202 "df_num": q,
203 "df_den": df,
204 "pvalue": float(stats.f.sf(F, q, df)),
205 "inference": "Homoskedastic nested-model F",
206 "rss_unrestricted": rss_u,
207 "rss_restricted": rss_r,
208 "design_columns": len(labels_sat),
209 "deviations": [],
210 "backend": "sparse equivalent deviations restriction",
211 "restricted_solver": check,
212 }
213 del ar, br, ur
214 print("SPARSE", key, "BJS imputation", flush=True)
215 first, offset = design(data, np.full(n, -1))
216 zero = data.d.eq(0).to_numpy()
217 a0 = first[zero]
218 b0, u0, check = solve(a0, y[zero])
219 residual = y - first @ b0
220 evs = sorted(set(events[np.isfinite(events)]))
221 means = {e: float(residual[events == e].mean()) for e in evs}
222 second = residual.copy()
223 for e in evs:
224 second[events == e] -= means[e]
225
226 def bjs_se(ew):
227 aa = np.zeros(n)
228 for e, w in ew.items():
229 aa[events == e] = w / np.count_nonzero(events == e)
230 target = np.asarray(first.T @ aa).ravel()
231 scale = np.sqrt(np.asarray(a0.power(2).sum(axis=0)).ravel())
232 assert (scale > 0).all()
233 zz = lsmr(
234 a0.multiply(1 / scale).T,
235 target / scale,
236 atol=1e-10,
237 btol=1e-10,
238 maxiter=6000,
239 )
240 score = np.bincount(cluster, weights=aa * second, minlength=g) - np.bincount(
241 cluster[zero], weights=zz[0] * u0, minlength=g
242 )
243 return float(np.sqrt(score @ score))
244
245 pooled = sum(w * means[e] for e, w in target.items())
246 bjs = summary(pooled, bjs_se(target), g)
247 estimates.insert(2, dict(method="BJS / did2s", **bjs))
248 paths["BJS / did2s"] = [
249 dict(
250 event_time=int(e),
251 **summary(
252 means[e], bjs_se({e: 1}) if window[0] <= e <= window[1] else None, g
253 ),
254 )
255 for e in evs
256 ]
257 result = run.results["specifications"][key]
258 result.update(
259 estimates=estimates,
260 event_studies=paths,
261 f_test=test,
262 target={
263 "event_times": list(target),
264 "weights": list(target.values()),
265 "weighting": run.config["post_weighting"],
266 },
267 checks={"sparse_solver_checks": checks, "restricted_design_rank_difference": q},
268 inference={
269 "backend": "Sparse equivalent normal equations; BJS uses the did2s two-stage score formula",
270 "standard_errors": "Cluster robust; displayed event times have pointwise intervals",
271 },
272 )
273 return result
SHA-256: b9fd5536d4886e99a707d2b162db5b11da3ff463930eca5b7590d27500096a3f
1 """Callaway–Sant'Anna DID on the existing Credit Markets analysis sample."""
2
3 import contextlib
4 import io
5 import json
6 import time
7 import warnings
8 from importlib.metadata import version
9
10 import numpy as np
11 import pandas as pd
12 from csdid.att_gt import ATTgt
13 from csdid_aggregations import summarize_aggregations
14 from matrix_models import plain
15 from scipy.stats import norm
16
17
18 def independent_checks(data, cells, estimate, dynamic_if, package_se, control_group):
19 """Direct not-yet-treated DIDs and dynamic IF with random cohort shares."""
20 units = data.drop_duplicates("unit").set_index("unit").sort_index()
21 panel = data.pivot(index="unit", columns="time", values="y").sort_index()
22 groups = units.g.to_numpy()
23 n_units = len(units)
24 cell_ifs = {}
25 error = []
26 for row in cells.itertuples():
27 g, t = int(row.group), int(row.year)
28 selected = groups == g
29 controls = groups == 0
30 if control_group == "notyettreated":
31 controls = controls | ((groups > max(t, g - 1)) & ~selected)
32 delta = (panel[t] - panel[g - 1]).to_numpy()
33 treated_mean = delta[selected].mean()
34 control_mean = delta[controls].mean()
35 error.append(row.att - (treated_mean - control_mean))
36 cell_ifs[g, t] = (
37 selected * (delta - treated_mean) / selected.mean()
38 - controls * (delta - control_mean) / controls.mean()
39 )
40 event_ifs, event_atts = [], []
41 for _, block in cells.loc[cells.event_time.ge(0)].groupby("event_time"):
42 mass = block.n.sum()
43 event_att = np.average(block.att, weights=block.n)
44 event_if = np.zeros(n_units)
45 for row in block.itertuples():
46 event_if += row.n / mass * cell_ifs[int(row.group), int(row.year)]
47 event_if += n_units / mass * (groups == row.group) * (row.att - event_att)
48 event_ifs.append(event_if)
49 event_atts.append(event_att)
50 if_direct = np.mean(event_ifs, axis=0)
51 assert abs(estimate - np.mean(event_atts)) < 1e-12
52 package_units = np.asarray(dynamic_if["unit_ids"])
53 direct_ordered = (
54 pd.Series(if_direct, index=panel.index).loc[package_units].to_numpy()
55 )
56 if_error = float(np.max(np.abs(direct_ordered - dynamic_if["values"])))
57 cluster_sums = pd.Series(if_direct).groupby(units.cluster.to_numpy()).sum()
58 direct_se = float(np.linalg.norm(cluster_sums) / n_units)
59 checks = {
60 "max_group_time_vs_direct_did_difference": float(np.max(np.abs(error))),
61 "max_dynamic_if_vs_direct_if_difference": if_error,
62 "independent_cluster_se": direct_se,
63 "cluster_se_difference": abs(direct_se - package_se),
64 }
65 assert checks["max_group_time_vs_direct_did_difference"] < 1e-10
66 assert if_error < 1e-9
67 assert abs(direct_se - package_se) < 1e-10
68 return checks
69
70
71 def fit_csdid(data, folder, options=None, xs=()):
72 # Dense package IFs scale as sampling units × cohorts × periods, with
73 # several simultaneous copies. Use the checked sufficient-statistic path
74 # for large intercept-only designs, without excluding an estimator.
75 sampling_units = (
76 len(data) if data.duplicated(["unit", "time"]).any() else data.unit.nunique()
77 )
78 dense_bytes = (
79 8 * sampling_units * data.loc[data.g.gt(0), "g"].nunique() * data.time.nunique()
80 )
81 if not xs and (
82 (options or {}).get("frequency_weights", False)
83 or dense_bytes > (options or {}).get("dense_if_max_bytes", 8 * 1024**2)
84 ):
85 from csdid_streaming import fit_intercept
86
87 output = fit_intercept(data, options)
88 output["checks"]["avoided_dense_if_bytes"] = int(dense_bytes)
89 (folder / "csdid.json").write_text(
90 json.dumps(output, indent=2, allow_nan=False) + "\n"
91 )
92 return output
93 regular = (
94 not xs
95 and data.weight.nunique() == 1
96 and not data.duplicated(["unit", "time"]).any()
97 and len(data) == data.unit.nunique() * data.time.nunique()
98 and np.array_equal(
99 np.sort(data.time.unique()), np.arange(1, data.time.nunique() + 1)
100 )
101 and data.g.eq(0).any()
102 )
103 if not regular:
104 from csdid_general import fit_general
105
106 return fit_general(data, folder, options or {}, xs)
107 options = {"control_group": "notyettreated", "aggregation": "dynamic"} | (
108 options or {}
109 )
110 assert options["aggregation"] == "dynamic", (
111 "CS-DID adapter reports dynamic aggregation"
112 )
113 start = time.perf_counter()
114 panel = data[["unit", "time", "g", "y", "cluster", "d"]].copy()
115 n_units, n_rows = panel.unit.nunique(), len(panel)
116 n_clusters, treated_cells = panel.cluster.nunique(), int(panel.d.sum())
117 console = io.StringIO()
118 with (
119 warnings.catch_warnings(record=True) as caught,
120 contextlib.redirect_stdout(console),
121 ):
122 warnings.simplefilter("always")
123 fit = ATTgt(
124 yname="y",
125 tname="time",
126 idname="unit",
127 gname="g",
128 data=panel,
129 xformla="y ~ 1",
130 control_group=options["control_group"],
131 panel=True,
132 allow_unbalanced_panel=False,
133 anticipation=0,
134 clustervar="cluster",
135 cband=False,
136 compute_inffunc=True,
137 ).fit(est_method="dr", base_period="universal", bstrap=False)
138 simple = fit.aggte(typec="simple", bstrap=False, cband=False).atte
139 dynamic = fit.aggte(typec="dynamic", bstrap=False, cband=False).atte
140 (folder / "csdid-package-summary.log").write_text(console.getvalue())
141 assert fit.MP["n"] == n_units and len(fit.dp["data"]) == n_rows
142 cells = pd.DataFrame(
143 {key: fit.results[key] for key in ["group", "year", "att", "se"]}
144 )
145 assert (
146 len(cells) == panel.loc[panel.g.gt(0), "g"].nunique() * panel.time.nunique()
147 and np.isfinite(cells.att).all()
148 )
149 assert cells.groupby("group").size().eq(panel.time.nunique()).all()
150 size = panel.drop_duplicates("unit").groupby("g").size()
151 cells["n"] = cells.group.map(size).astype(int)
152 cells["event_time"] = cells.year - cells.group
153 post = cells.event_time.ge(0)
154 assert cells.loc[post, "n"].sum() == treated_cells
155 estimate = float(dynamic["overall_att"])
156 scalar_from_cells = float(
157 np.average(cells.loc[post, "att"], weights=cells.loc[post, "n"])
158 )
159 assert abs(float(simple["overall_att"]) - scalar_from_cells) < 1e-12
160
161 critical = float(norm.ppf(0.975))
162 path = []
163 for event, att, se in zip(
164 dynamic["egt"], dynamic["att_egt"], np.ravel(dynamic["se_egt"]), strict=True
165 ):
166 block = cells.loc[cells.event_time.eq(event)]
167 assert abs(att - np.average(block.att, weights=block.n)) < 1e-12
168 path.append(
169 {
170 "event_time": int(event),
171 "estimate": float(att),
172 "se": float(se),
173 "lower": float(att - critical * se) if event != -1 else 0.0,
174 "upper": float(att + critical * se) if event != -1 else 0.0,
175 "relative_to_minus_one": float(att),
176 "n": int(block.n.sum()),
177 }
178 )
179 assert next(r["estimate"] for r in path if r["event_time"] == -1) == 0
180 dta = fit.dp["data"].loc[fit.dp["data"].time.eq(fit.dp["tlist"][0])]
181 checks = independent_checks(
182 panel,
183 cells,
184 estimate,
185 {
186 "unit_ids": dta.unit.to_numpy(),
187 "values": np.ravel(dynamic["inf_function"]["dynamic_inf_func"]),
188 },
189 float(dynamic["overall_se"]),
190 options["control_group"],
191 )
192 scalar_from_events = np.mean([r["estimate"] for r in path if r["event_time"] >= 0])
193 assert abs(estimate - scalar_from_events) < 1e-12
194 checks["scalar_vs_equal_event_average"] = abs(estimate - scalar_from_events)
195 checks["simple_vs_treated_cell_average"] = abs(
196 float(simple["overall_att"]) - scalar_from_cells
197 )
198 checks["group_time_cells"] = len(cells)
199 checks["post_treatment_cells"] = int(post.sum())
200 checks["treated_unit_years"] = int(cells.loc[post, "n"].sum())
201 aggregations = summarize_aggregations(fit, panel, cells)
202 output = plain(
203 {
204 "method": "CS-DID",
205 "aggregations": aggregations,
206 "class": "csdid.att_gt.ATTgt",
207 "package": {"csdid": version("csdid"), "drdid": version("drdid")},
208 "parameters": {
209 "est_method": "dr",
210 "control_group": options["control_group"],
211 "xformla": "y ~ 1",
212 "base_period": "universal",
213 "anticipation": 0,
214 "panel": True,
215 "allow_unbalanced_panel": False,
216 "clustervar": "cluster",
217 "clusters": n_clusters,
218 "bstrap": False,
219 "cband": False,
220 "aggregate": "dynamic",
221 "event_aggregate": "dynamic",
222 "balance_e": None,
223 "post_event_min": 0,
224 "post_event_max": max(r["event_time"] for r in path),
225 },
226 "partial_pooling": {
227 "parameter": None,
228 "value": None,
229 "status": "not_applicable",
230 "fit_scope": "Group-time DIDs; no synthetic-control pooling parameter",
231 },
232 "estimate": estimate,
233 "se": float(dynamic["overall_se"]),
234 "lower": estimate - critical * dynamic["overall_se"],
235 "upper": estimate + critical * dynamic["overall_se"],
236 "event_study": path,
237 "group_time_estimates": cells.assign(
238 adoption_period=cells.group,
239 calendar_period=cells.year,
240 ).to_dict("records"),
241 "package_dynamic_equal_event_att": float(dynamic["overall_att"]),
242 "treated_cell_att": float(simple["overall_att"]),
243 "treated_cell_se": float(simple["overall_se"]),
244 "checks": checks,
245 "warnings": sorted({str(item.message) for item in caught}),
246 "seconds": time.perf_counter() - start,
247 "target": "Dynamic aggregation: equal weights across all observed nonnegative event times; cohort-size weights within each event time",
248 "donor_rule": options["control_group"],
249 "inference": "Configured-cluster analytical influence-function SE; 95% pointwise normal intervals; includes estimated group-share uncertainty",
250 }
251 )
252 (folder / "csdid.json").write_text(
253 json.dumps(output, indent=2, allow_nan=False) + "\n"
254 )
255 print(f"CS-DID: {estimate:.10f}; configured-cluster SE: {output['se']:.10f}")
256 print("Checks:", checks)
257 return output
SHA-256: 07619c3b3dfee796ae9e4842afdbc6e367a74cece254ede73f7c403dfc924f6d
1 """CS-DID with controls, weights, repeated cross-sections and unbalanced panels.
2
3 The original rows are retained. Native dynamic/simple/group summaries are kept
4 separate from the weighted observed-treated-cell target used in the comparison.
5 """
6
7 import contextlib
8 import io
9 import json
10 import time
11 import warnings
12 from importlib.metadata import version
13
14 import numpy as np
15 import pandas as pd
16 from csdid.att_gt import ATTgt
17 from csdid.attgt_fnc import drdid_trim
18 from csdid.attgt_fnc.compute_att_gt_shared import (
19 drdid_design_singular,
20 overlap_check_fail,
21 rcond_check_fail,
22 )
23 from eligibility import UnidentifiedTarget
24 from matrix_models import plain
25 from scipy.linalg import qr
26 from scipy.stats import norm
27
28
29 def independent_columns(x):
30 """Select a scaled basis for the same cell-specific covariate column space."""
31 x = np.asarray(x, dtype=float)
32 scale = np.linalg.norm(x, axis=0)
33 scaled = x / np.where(scale > 0, scale, 1)
34 _, r, pivots = qr(scaled, mode="economic", pivoting=True)
35 diagonal = np.abs(np.diag(r))
36 tolerance = max(x.shape) * np.finfo(float).eps * diagonal.max(initial=0)
37 rank = int((diagonal > tolerance).sum())
38 # Preserve the intercept first: DRDID detects it by column position.
39 chosen = sorted(pivots[:rank])
40 basis = scaled[:, chosen]
41 if rank == 0:
42 raise ValueError("Empty covariate column space")
43 # Put a literal intercept first even when pivoting selected another constant.
44 constant = np.flatnonzero(np.ptp(basis, axis=0) < 1e-14)
45 if len(constant):
46 j = int(constant[0])
47 basis = np.column_stack([np.ones(len(x)), np.delete(basis, j, axis=1)])
48 return basis
49
50
51 def rank_reduced_dr(*args, covariates, i_weights, **kwargs):
52 """Native DR after removing only redundant columns; retain its support guards."""
53 x = independent_columns(covariates)
54 treated = np.asarray(args[2] if args else kwargs["D"])
55 post = kwargs.get("post")
56 if rcond_check_fail(x[treated == 0]):
57 raise ValueError(
58 "Control covariate design is singular in this group-time comparison"
59 )
60 if overlap_check_fail(x, treated):
61 raise ValueError("Propensity overlap fails in this group-time comparison")
62 if drdid_design_singular(x, i_weights, treated, "dr", post):
63 raise ValueError("A required group-period covariate design is singular")
64 estimator = drdid_trim.drdid_panel if args else drdid_trim.drdid_rc
65 return estimator(*args, covariates=x, i_weights=i_weights, **kwargs)
66
67
68 def bounds(estimate, se):
69 return {
70 "estimate": float(estimate),
71 "se": float(se),
72 "lower": float(estimate - norm.ppf(0.975) * se),
73 "upper": float(estimate + norm.ppf(0.975) * se),
74 }
75
76
77 def cluster_se(values, clusters):
78 sums = pd.Series(np.ravel(values)).groupby(np.asarray(clusters)).sum()
79 return float(np.linalg.norm(sums) / len(values))
80
81
82 def treated_cell_summary(fit, data, cells):
83 """Delta-method IF for the ratio of weighted observed treated-cell effects."""
84 dp = fit.dp
85 idname = dp["idname"]
86 ordered = dp["data"].drop_duplicates(idname)
87 ids = ordered[idname].to_numpy()
88 n = len(ids)
89 if dp["true_rep_cross_section"]:
90 source_ids = np.arange(len(data))
91 else:
92 source_ids = data.unit.to_numpy()
93 positions = pd.Index(ids).get_indexer(source_ids)
94 assert (positions >= 0).all() and n == fit.MP["n"]
95 weights = data.loc[data.d.eq(1)].groupby(["g", "time"]).weight.sum()
96 cells = cells.copy()
97 cells["mass"] = [weights.get((g, t), 0.0) for g, t in zip(cells.group, cells.year)]
98 missing = cells.loc[cells.mass.gt(0) & ~np.isfinite(cells.att)]
99 requested = set(weights.index)
100 covered = set(zip(cells.group, cells.year))
101 if len(missing) or requested - covered:
102 raise UnidentifiedTarget(
103 f"CS-DID cannot estimate {len(missing) + len(requested - covered)} "
104 "observed post-treatment group-time cells with the configured controls; "
105 "no target cells were dropped. See csdid-group-time.json and package log."
106 )
107 use = cells.mass.gt(0).to_numpy()
108 post = cells.loc[use]
109 mass = post.mass.to_numpy()
110 denominator = mass.sum()
111 estimate = float(np.average(post.att, weights=mass))
112 influence = np.asarray(fit.MP["inffunc"]["inffunc"])[:, use] @ (mass / denominator)
113 cell_att = dict(zip(zip(post.group, post.year), post.att))
114 adjustment = np.zeros(n)
115 for (g, t), indexes in data.loc[data.d.eq(1)].groupby(["g", "time"]).groups.items():
116 indexes = np.asarray(indexes)
117 np.add.at(
118 adjustment,
119 positions[indexes],
120 data.loc[indexes, "weight"].to_numpy() * (cell_att[g, t] - estimate),
121 )
122 influence += n / denominator * adjustment
123 se = cluster_se(influence, ordered.cluster)
124 return {
125 "typec": "treated_cells",
126 "aggregation": "Observed treated observations, configured weights",
127 "matches_treated_unit_year_target": True,
128 "weight_sum": float(denominator),
129 "post_cells": len(post),
130 **bounds(estimate, se),
131 }, cells
132
133
134 def fit_general(data, folder, options, xs):
135 started = time.perf_counter()
136 options = {"control_group": "notyettreated", "aggregation": "dynamic"} | options
137 data = data.reset_index(drop=True).copy()
138 repeated = data.duplicated(["unit", "time"]).any()
139 varying_weights = data.groupby("unit").weight.nunique().max() > 1
140 panel = not repeated
141 formula = "y ~ 1" + (" + " + " + ".join(xs) if xs else "")
142 console = io.StringIO()
143 try:
144 with (
145 warnings.catch_warnings(record=True) as caught,
146 contextlib.redirect_stdout(console),
147 ):
148 warnings.simplefilter("always")
149 fit = ATTgt(
150 yname="y",
151 tname="time",
152 idname=None if repeated else "unit",
153 gname="g",
154 data=data,
155 xformla=formula,
156 control_group=options["control_group"],
157 panel=panel,
158 allow_unbalanced_panel=True,
159 weights_name="weight",
160 fix_weights="varying" if varying_weights else None,
161 clustervar="cluster",
162 cband=False,
163 compute_inffunc=True,
164 faster_mode=options.get("faster_mode", True),
165 ).fit(
166 est_method=rank_reduced_dr if xs else "dr",
167 base_period="universal",
168 bstrap=False,
169 )
170 retained = fit.dp["data"]
171 if len(retained) != len(data) or set(retained.index) != set(data.index):
172 raise UnidentifiedTarget(
173 "CS-DID preprocessing changed the prepared sample; result withheld."
174 )
175 for column in ["y", "g", "time", "weight"]:
176 np.testing.assert_array_equal(
177 retained.sort_index()[column], data[column]
178 )
179 cells = pd.DataFrame(
180 {k: fit.results[k] for k in ["group", "year", "att", "se"]}
181 )
182 cells["event_time"] = cells.year - cells.group
183 (folder / "csdid-group-time.json").write_text(
184 json.dumps(plain(cells.to_dict("records")), indent=2)
185 )
186 matched, cells = treated_cell_summary(fit, data, cells)
187 summaries, native = [], {}
188 for kind in ["dynamic", "simple", "group"]:
189 # Only pre-treatment diagnostics may be unavailable; every
190 # observed post-treatment target cell was checked above.
191 agg = fit.aggte(typec=kind, bstrap=False, cband=False, na_rm=True).atte
192 native[kind] = agg
193 summaries.append(
194 {
195 "typec": kind,
196 "aggregation": f"Native CS-DID {kind}",
197 "matches_treated_unit_year_target": False,
198 **bounds(agg["overall_att"], agg["overall_se"]),
199 }
200 )
201 dynamic = native["dynamic"]
202 ordered = fit.dp["data"].drop_duplicates(fit.dp["idname"])
203 direct_se = cluster_se(
204 dynamic["inf_function"]["dynamic_inf_func"], ordered.cluster
205 )
206 np.testing.assert_allclose(direct_se, dynamic["overall_se"], atol=1e-10)
207 finally:
208 messages = sorted({str(w.message) for w in locals().get("caught", [])})
209 (folder / "csdid-package-summary.log").write_text(
210 console.getvalue() + "\n" + "\n".join(messages)
211 )
212 path = []
213 for e, att, se in zip(
214 dynamic["egt"], dynamic["att_egt"], np.ravel(dynamic["se_egt"]), strict=True
215 ):
216 block = cells.loc[cells.event_time.eq(e)]
217 path.append(
218 {
219 "event_time": int(e),
220 **bounds(att, se),
221 "relative_to_minus_one": float(att),
222 "n": float(block.mass.sum()),
223 }
224 )
225 post = [r["estimate"] for r in path if r["event_time"] >= 0]
226 np.testing.assert_allclose(np.mean(post), dynamic["overall_att"], atol=1e-10)
227 # Native dynamic weights need not be the observed cell weights. Do not
228 # reconstruct the matched scalar from the native event path on such inputs.
229 output = plain(
230 {
231 "method": "CS-DID",
232 "class": "csdid.att_gt.ATTgt",
233 "package": {p: version(p) for p in ["csdid", "drdid"]},
234 "parameters": {
235 "est_method": "dr",
236 "covariate_basis": "cell-specific rank-reduced column space"
237 if xs
238 else "intercept",
239 "control_group": options["control_group"],
240 "xformla": formula,
241 "panel": panel,
242 "effective_panel": fit.dp["panel"],
243 "allow_unbalanced_panel": True,
244 "weights_name": "weight",
245 "fix_weights": "varying" if varying_weights else None,
246 "base_period": "universal",
247 "clustervar": "cluster",
248 "aggregate": "dynamic",
249 "faster_mode": options.get("faster_mode", True),
250 },
251 "data_structure": "repeated cross-section"
252 if repeated
253 else "panel"
254 if fit.dp["panel"]
255 else "unbalanced or varying-weight panel; native repeated-cross-section comparisons",
256 "aggregations": {
257 "summaries": summaries + [matched],
258 "same_group_time_effects": True,
259 "additional_covariates": list(xs),
260 },
261 **bounds(dynamic["overall_att"], dynamic["overall_se"]),
262 "event_study": path,
263 "group_time_estimates": cells.to_dict("records"),
264 "treated_cell_att": matched["estimate"],
265 "treated_cell_se": matched["se"],
266 "checks": {
267 "input_rows": len(data),
268 "retained_rows": len(retained),
269 "all_observed_post_cells_estimated": True,
270 "dynamic_cluster_se_error": abs(direct_se - dynamic["overall_se"]),
271 },
272 "target": "Equal average of nonnegative native CS-DID event-time effects",
273 "donor_rule": options["control_group"],
274 "inference": "Configured-cluster analytical influence-function intervals; observed-cell ratio includes estimated weight-share uncertainty",
275 "warnings": sorted({str(w.message) for w in caught}),
276 "seconds": time.perf_counter() - started,
277 }
278 )
279 (folder / "csdid.json").write_text(
280 json.dumps(output, indent=2, allow_nan=False) + "\n"
281 )
282 return output
SHA-256: 378198e98c44bd46862601bdcb514b19d9ea4038cddc38dad7fdba9b0aa0216b
1 """Intercept-only CS-DID using cluster sufficient statistics, not dense unit IFs.
2
3 Two-period DR with only an intercept is the weighted difference in differences.
4 Accumulate its cluster influence contributions and aggregation-share terms
5 directly. This preserves observations and the configured clustering while avoiding
6 the package's observations-by-group-times influence matrix on large panels.
7 """
8
9 import time
10
11 import numpy as np
12 import pandas as pd
13 from csdid_general import bounds
14 from eligibility import UnidentifiedTarget
15 from matrix_models import plain
16
17
18 def fit_intercept(data, options=None):
19 started = time.perf_counter()
20 control = (options or {}).get("control_group", "notyettreated")
21 repeated = data.duplicated(["unit", "time"]).any()
22 balanced = not repeated and len(data) == data.unit.nunique() * data.time.nunique()
23 varying = data.groupby("unit").weight.nunique().max() > 1
24 clusters, _ = pd.factorize(data.cluster, sort=True)
25 nc = int(clusters.max()) + 1
26 frequency = (options or {}).get("frequency_weights", False)
27 multiplicity = np.ones(nc)
28 if frequency:
29 counts = data.groupby("cluster", sort=True).weight.agg(["first", "nunique"])
30 if counts["nunique"].max() != 1 or not np.allclose(
31 counts["first"], counts["first"].round()
32 ):
33 raise ValueError(
34 "Frequency weights must be constant integer cluster multiplicities"
35 )
36 multiplicity = counts["first"].to_numpy()
37
38 def standard_error(contribution):
39 return float(np.linalg.norm(contribution / np.sqrt(multiplicity)))
40
41 work = data[["g", "time", "y", "weight"]].copy()
42 work["cluster"] = clusters
43 work["wy"] = work.y * work.weight
44 stats = (
45 work.groupby(["g", "time", "cluster"], sort=True)
46 .agg(w=("weight", "sum"), wy=("wy", "sum"), count=("y", "size"))
47 .reset_index()
48 )
49 periods = np.sort(data.time.unique())
50 cohorts = np.sort(data.loc[data.g.gt(0), "g"].unique())
51 if not data.g.eq(0).any():
52 raise UnidentifiedTarget(
53 "No controls remain for the full all-post target; sample not truncated."
54 )
55 by_time = {t: frame for t, frame in stats.groupby("time")}
56 if repeated:
57 group_mass = work.groupby(["g", "cluster"]).weight.sum()
58 else:
59 units = data.groupby("unit", sort=True).agg(
60 g=("g", "first"), cluster=("cluster", "first"), weight=("weight", "mean")
61 )
62 # Native balanced-panel aggte uses first-period weights, which equal
63 # unit-mean weights unless varying weights route it through RC.
64 codes = pd.Index(np.sort(data.cluster.unique())).get_indexer(units.cluster)
65 units["cluster"] = codes
66 group_mass = units.groupby(["g", "cluster"]).weight.sum()
67 mass_vectors = {}
68 for g in cohorts:
69 values = group_mass.loc[g]
70 q = np.zeros(nc)
71 q[values.index.to_numpy(dtype=int)] = values.to_numpy()
72 mass_vectors[g] = q
73 cohort_mass = {g: q.sum() for g, q in mass_vectors.items()}
74
75 def component(g, t, base, treated):
76 frame = by_time[t]
77 select = frame.g.eq(g) if treated else frame.g.eq(0)
78 if not treated and control == "notyettreated":
79 select |= frame.g.gt(max(t, base)) & frame.g.ne(g)
80 return frame.loc[select]
81
82 cells = []
83 for g in cohorts:
84 before = periods[periods < g]
85 if not len(before):
86 raise UnidentifiedTarget("A treated cohort has no pre-treatment period")
87 base = before[-1]
88 for t in periods:
89 blocks = [
90 component(g, t, base, True),
91 component(g, base, t, True),
92 component(g, t, base, False),
93 component(g, base, t, False),
94 ]
95 mass = blocks[0].w.sum() if t >= g else 0.0
96 if any(b.empty for b in blocks):
97 if mass > 0:
98 raise UnidentifiedTarget(
99 f"CS-DID lacks group-period support for cohort {g}, period {t}; target not trimmed."
100 )
101 continue
102 # Native intercept-only overlap guards and propensity trimming.
103 nt = blocks[0]["count"].sum() + blocks[1]["count"].sum()
104 nn = sum(b["count"].sum() for b in blocks)
105 wt = blocks[0].w.sum() + blocks[1].w.sum()
106 ww = sum(b.w.sum() for b in blocks)
107 if nt / nn >= 0.999 or wt / ww >= 0.995:
108 if mass > 0:
109 raise UnidentifiedTarget(
110 f"CS-DID propensity overlap fails at cohort {g}, period {t}."
111 )
112 continue
113 means = [b.wy.sum() / b.w.sum() for b in blocks]
114 att = float(np.dot([1, -1, -1, 1], means)) if t != base else 0.0
115 cells.append(
116 dict(group=g, year=t, event_time=t - g, att=att, mass=mass, base=base)
117 )
118 cells = pd.DataFrame(cells)
119 post = cells.loc[cells.event_time.ge(0)]
120 requested_mass = data.loc[data.d.eq(1), "weight"].sum()
121 np.testing.assert_allclose(post.mass.sum(), requested_mass)
122 post_lengths = post.groupby("group").size()
123 simple_den = sum(cohort_mass[g] for g in post.group)
124 group_den = sum(cohort_mass[g] for g in post_lengths.index)
125 targets = {
126 "simple": float(
127 sum(cohort_mass[r.group] * r.att for r in post.itertuples()) / simple_den
128 ),
129 "group": float(
130 sum(cohort_mass[g] * b.att.mean() for g, b in post.groupby("group"))
131 / group_den
132 ),
133 "treated_cells": float(np.average(post.att, weights=post.mass)),
134 }
135 contributions = {kind: np.zeros(nc) for kind in targets}
136 dynamic_contribution = np.zeros(nc)
137 path = []
138 positive_events = post.event_time.nunique()
139 for e, event in cells.groupby("event_time", sort=True):
140 denominator = sum(cohort_mass[g] for g in event.group)
141 average = (
142 sum(cohort_mass[r.group] * r.att for r in event.itertuples()) / denominator
143 )
144 event_contribution = np.zeros(nc)
145 for r in event.itertuples():
146 value = np.zeros(nc)
147 blocks = [
148 component(r.group, r.year, r.base, True),
149 component(r.group, r.base, r.year, True),
150 component(r.group, r.year, r.base, False),
151 component(r.group, r.base, r.year, False),
152 ]
153 for sign, block in zip([1, -1, -1, 1], blocks, strict=True):
154 w = block.w.sum()
155 mean = block.wy.sum() / w
156 np.add.at(
157 value,
158 block.cluster.to_numpy(dtype=int),
159 sign * (block.wy - mean * block.w).to_numpy() / w,
160 )
161 q = mass_vectors[r.group]
162 event_contribution += cohort_mass[
163 r.group
164 ] / denominator * value + q / denominator * (r.att - average)
165 if e >= 0:
166 for kind, den, factor in [
167 ("simple", simple_den, 1),
168 ("group", group_den, 1 / post_lengths[r.group]),
169 ]:
170 contributions[kind] += (
171 factor
172 / den
173 * (cohort_mass[r.group] * value + q * (r.att - targets[kind]))
174 )
175 cellmass = blocks[0]
176 contributions["treated_cells"] += r.mass / requested_mass * value
177 np.add.at(
178 contributions["treated_cells"],
179 cellmass.cluster.to_numpy(dtype=int),
180 cellmass.w.to_numpy()
181 / requested_mass
182 * (r.att - targets["treated_cells"]),
183 )
184 if e >= 0:
185 dynamic_contribution += event_contribution / positive_events
186 path.append(
187 dict(
188 event_time=int(e),
189 **bounds(average, standard_error(event_contribution)),
190 n=float(event.mass.sum()),
191 relative_to_minus_one=float(average),
192 )
193 )
194 dynamic = float(np.mean([p["estimate"] for p in path if p["event_time"] >= 0]))
195 summaries = [
196 {
197 "typec": "dynamic",
198 "aggregation": "Equal post-event times",
199 "matches_treated_unit_year_target": False,
200 **bounds(dynamic, standard_error(dynamic_contribution)),
201 }
202 ]
203 for kind, estimate in targets.items():
204 summaries.append(
205 {
206 "typec": kind,
207 "aggregation": "Observed treated observations, configured weights"
208 if kind == "treated_cells"
209 else f"Native-equivalent CS-DID {kind}",
210 "matches_treated_unit_year_target": kind == "treated_cells",
211 **bounds(estimate, standard_error(contributions[kind])),
212 }
213 )
214 matched = summaries[-1]
215 return plain(
216 {
217 "method": "CS-DID",
218 **bounds(dynamic, standard_error(dynamic_contribution)),
219 "event_study": path,
220 "group_time_estimates": cells.to_dict("records"),
221 "parameters": {
222 "est_method": "dr",
223 "xformla": "y ~ 1",
224 "control_group": control,
225 "base_period": "universal",
226 "clustervar": "cluster",
227 "weights_name": "weight",
228 "backend": "cluster-sufficient-statistics",
229 "panel": not repeated,
230 "effective_panel": balanced and not varying,
231 "frequency_weights": frequency,
232 "effective_clusters": int(multiplicity.sum()),
233 },
234 "data_structure": "repeated cross-section"
235 if repeated
236 else "panel"
237 if balanced and not varying
238 else "unbalanced or varying-weight panel; repeated-cross-section comparisons",
239 "aggregations": {
240 "summaries": summaries,
241 "same_group_time_effects": True,
242 "additional_covariates": [],
243 },
244 "treated_cell_att": matched["estimate"],
245 "treated_cell_se": matched["se"],
246 "checks": {
247 "input_rows": len(data),
248 "retained_rows": len(data),
249 "all_observed_post_cells_estimated": True,
250 "post_weight_sum": float(post.mass.sum()),
251 },
252 "target": "Equal average of nonnegative native CS-DID event-time effects",
253 "donor_rule": control,
254 "inference": "Configured-cluster analytical influence functions, including aggregation-share uncertainty; sufficient-statistic calculation of intercept-only CS-DID",
255 "seconds": time.perf_counter() - started,
256 }
257 )
SHA-256: cbed0fb13450d9398f9a59d53dbfdf8e275a74daa01158c38c0f3396a3e8c6ee
1 """Explicit CS-DID scalar targets and independently checked influence functions."""
2
3 import contextlib
4 import hashlib
5 import inspect
6 import io
7
8 import numpy as np
9 import pandas as pd
10 from scipy.stats import norm
11
12
13 def summarize_aggregations(fit, data, cells):
14 units = data.drop_duplicates("unit").set_index("unit").sort_index()
15 panel = data.pivot(index="unit", columns="time", values="y").sort_index()
16 groups = units.g.to_numpy()
17 n = len(units)
18 control_group = fit.dp["control_group"]
19 post = cells.loc[cells.event_time.ge(0)].copy().reset_index(drop=True)
20 sizes = post.groupby("group").n.first()
21 lengths = post.groupby("group").size()
22 event_mass = post.groupby("event_time").n.sum()
23 event_count = len(event_mass)
24 cell_ifs = []
25 for row in post.itertuples():
26 selected = groups == row.group
27 controls = groups == 0
28 if control_group == "notyettreated":
29 controls |= (groups > max(row.year, row.group - 1)) & ~selected
30 delta = (panel[row.year] - panel[row.group - 1]).to_numpy()
31 cell_ifs.append(
32 selected * (delta - delta[selected].mean()) / selected.mean()
33 - controls * (delta - delta[controls].mean()) / controls.mean()
34 )
35 cell_ifs = np.asarray(cell_ifs).T
36 package_ids = (
37 fit.dp["data"]
38 .loc[fit.dp["data"]["time"].eq(fit.dp["tlist"][0]), "unit"]
39 .to_numpy()
40 )
41 labels = {
42 "dynamic": "Equal post-event times",
43 "simple": "Equal treated unit-years",
44 "group": "Cohort-size-weighted post-period means",
45 }
46 if_keys = {
47 "dynamic": "dynamic_inf_func",
48 "simple": "simple_att",
49 "group": "selective_inf_func",
50 }
51 results = []
52 for kind in ["dynamic", "simple", "group"]:
53 with contextlib.redirect_stdout(io.StringIO()):
54 agg = fit.aggte(typec=kind, bstrap=False, cband=False).atte
55 if kind == "dynamic":
56 weights = post.n / post.event_time.map(event_mass) / event_count
57 elif kind == "simple":
58 weights = post.n / post.n.sum()
59 else:
60 weights = post.n / post.group.map(lengths) / sizes.sum()
61 manual_att = float(weights @ post.att)
62 manual_if = cell_ifs @ weights.to_numpy()
63 if kind == "dynamic":
64 for event, block in post.groupby("event_time"):
65 mass = event_mass[event]
66 average = np.average(block.att, weights=block.n)
67 for row in block.itertuples():
68 manual_if += (
69 n
70 / mass
71 / event_count
72 * (groups == row.group)
73 * (row.att - average)
74 )
75 else:
76 divisor = post.n.sum() if kind == "simple" else sizes.sum()
77 for row in post.itertuples():
78 coefficient = 1 if kind == "simple" else 1 / lengths[row.group]
79 manual_if += (
80 n
81 / divisor
82 * coefficient
83 * (groups == row.group)
84 * (row.att - manual_att)
85 )
86 package_if = np.ravel(agg["inf_function"][if_keys[kind]])
87 reordered = pd.Series(manual_if, index=units.index).loc[package_ids].to_numpy()
88 country_sums = pd.Series(manual_if).groupby(units.cluster.to_numpy()).sum()
89 manual_se = float(np.linalg.norm(country_sums) / n)
90 att, se = float(agg["overall_att"]), float(agg["overall_se"])
91 checks = {
92 "point_error": abs(att - manual_att),
93 "influence_function_max_error": float(
94 np.max(np.abs(package_if - reordered))
95 ),
96 "cluster_se_error": abs(se - manual_se),
97 }
98 assert checks["point_error"] < 1e-11, (kind, checks)
99 assert checks["influence_function_max_error"] < 1e-8, (kind, checks)
100 assert checks["cluster_se_error"] < 1e-10, (kind, checks)
101 assert np.isclose(weights.sum(), 1)
102 results.append(
103 {
104 "typec": kind,
105 "aggregation": labels[kind],
106 "estimate": att,
107 "se": se,
108 "lower": att - norm.ppf(0.975) * se,
109 "upper": att + norm.ppf(0.975) * se,
110 "matches_treated_unit_year_target": kind == "simple",
111 "checks": checks,
112 "cohort_total_weights": [
113 {
114 "cohort": int(g),
115 "units": int(sizes[g]),
116 "post_periods": int(lengths[g]),
117 "weight": float(weights[post.group.eq(g)].sum()),
118 }
119 for g in sizes.index
120 ],
121 }
122 )
123 import csdid.aggte_fnc.compute_aggte as module
124
125 source = inspect.getfile(module)
126 with open(source, "rb") as stream:
127 code_hash = hashlib.sha256(stream.read()).hexdigest()
128 return {
129 "summaries": results,
130 "control_group": control_group,
131 "same_group_time_effects": True,
132 "additional_covariates": [],
133 "compute_aggte_sha256": code_hash,
134 "inference": "Clustered analytical influence functions including estimated cohort-share uncertainty; 95% pointwise normal intervals",
135 }
136
137
138 def aggregation_table(model):
139 return [
140 {
141 "Aggregation": r["aggregation"],
142 "Estimate (SDs)": r["estimate"],
143 "SE": r["se"],
144 "95% CI": f"[{r['lower']:.4f}, {r['upper']:.4f}]",
145 }
146 for r in model["aggregations"]["summaries"]
147 ]
SHA-256: 3f675276f9022147996dd5b62762721be2dd436d5db171491a286b05628176ae
1 """Construct outcome, treatment and observation-mass matrices from long inputs.
2
3 The fitting unit is the configured panel unit, not an individual row. Repeated
4 observations are averaged within unit-period using their observation weights.
5 Crabbymetrics fits its standard cell-mean objective; cell masses are retained for
6 ATT/event aggregation, not misrepresented as supported precision weights.
7 """
8
9 from dataclasses import dataclass
10
11 import numpy as np
12
13 CRABBY_MATRIX_METHODS = {"sdid", "matrix-completion", "horizontal-ridge"}
14
15
16 @dataclass
17 class MatrixInput:
18 y: np.ndarray
19 w: np.ndarray
20 mass: np.ndarray
21 counts: np.ndarray
22 units: np.ndarray
23 periods: np.ndarray
24 rows: int
25
26 @property
27 def observed(self):
28 return self.counts > 0
29
30 def audit(self):
31 seen = self.observed
32 return {
33 "source_rows": self.rows,
34 "matrix_shape": list(self.y.shape),
35 "observed_cells": int(seen.sum()),
36 "missing_cells": int((~seen).sum()),
37 "repeated_cells": int((self.counts > 1).sum()),
38 "minimum_cell_count": int(self.counts[seen].min()),
39 "maximum_cell_count": int(self.counts.max()),
40 "observation_mass": float(self.mass.sum()),
41 "treated_observation_mass": float(self.mass[self.w == 1].sum()),
42 "within_cell_aggregation": "Observation-weighted outcome mean",
43 "fitting_weights": "Standard crabbymetrics objective on cell means; no cell-count precision weighting",
44 "reporting_weights": "Sum of original observation weights per treated unit-period",
45 }
46
47 def save(self, folder):
48 np.savez_compressed(
49 folder / "matrix-input-local.npz",
50 Y=self.y,
51 W=self.w,
52 observation_mass=self.mass,
53 row_counts=self.counts,
54 observed=self.observed,
55 units=self.units,
56 periods=self.periods,
57 )
58
59
60 def prepare_matrix(data):
61 """Collapse once, preserve both axes, and never fill missing outcomes with zero."""
62 keys = ["unit", "time"]
63 if data[keys + ["g", "d", "y", "weight"]].isna().any().any():
64 raise ValueError(
65 "Matrix source has missing keys, treatment, outcome or weights"
66 )
67 if (
68 not np.isfinite(data[["y", "weight"]].to_numpy()).all()
69 or not data.weight.gt(0).all()
70 ):
71 raise ValueError(
72 "Matrix source outcomes/positive observation weights must be finite"
73 )
74 adoption = data.groupby("unit", sort=True).g.agg(["first", "nunique"])
75 if adoption["nunique"].max() != 1:
76 raise ValueError("Treatment adoption varies within the configured panel unit")
77 implied = (data.g.gt(0) & data.time.ge(data.g)).astype(int)
78 if not np.array_equal(data.d.to_numpy(), implied.to_numpy()):
79 raise ValueError("Treatment is inconsistent with unit-level absorbing adoption")
80 units = adoption.index.to_numpy()
81 periods = np.sort(data.time.unique())
82 groups = data.groupby(keys, sort=True)
83 counts = groups.size().unstack().reindex(index=units, columns=periods).fillna(0)
84 mass = groups.weight.sum().unstack().reindex(index=units, columns=periods).fillna(0)
85 if data.weight.nunique() == 1:
86 # Preserve the direct-pivot values bit-for-bit for singleton panels.
87 means = groups.y.mean()
88 else:
89 numerator = (data.y * data.weight).groupby([data.unit, data.time]).sum()
90 means = numerator / groups.weight.sum()
91 y = means.unstack().reindex(index=units, columns=periods).to_numpy(dtype=float)
92 g = adoption["first"].to_numpy()
93 w = ((g[:, None] > 0) & (periods[None, :] >= g[:, None])).astype(float)
94 result = MatrixInput(
95 y,
96 w,
97 mass.to_numpy(dtype=float),
98 counts.to_numpy(dtype=int),
99 units,
100 periods,
101 len(data),
102 )
103 assert result.counts.sum() == len(data)
104 assert np.isclose(result.mass.sum(), data.weight.sum())
105 assert np.isclose(
106 result.mass[result.w == 1].sum(), data.loc[data.d.eq(1), "weight"].sum()
107 )
108 return result
SHA-256: dc7188bf9bf709f0803169dfc3b1198e8b9d2b79d5904ebc5f010a48edde2ee2
1 """Fit matrix-panel counterfactuals on the prepared unit-period matrix."""
2
3 import hashlib
4 import json
5 import time
6 from importlib.metadata import distribution, version
7 from pathlib import Path
8
9 import numpy as np
10
11
12 def build_identity():
13 """Distinguish the branch wheel from PyPI's same-version release."""
14 import crabbymetrics as cm
15
16 binary = next(Path(cm.__file__).parent.glob("*.so"))
17 binary_hash = hashlib.sha256(binary.read_bytes()).hexdigest()
18 root = Path(__file__).resolve().parents[2]
19 record = root / "configs/crabbymetrics-build.json"
20 build = json.loads(record.read_text()) if record.exists() else {}
21 installed = json.loads(
22 distribution("crabbymetrics").read_text("direct_url.json") or "{}"
23 )
24 commit = installed.get("vcs_info", {}).get("commit_id")
25 if not commit and build.get("extension_sha256") == binary_hash:
26 commit = build["source_commit"]
27 return {
28 "source_commit": commit,
29 "extension_sha256": binary_hash,
30 }
31
32
33 def plain(value):
34 if isinstance(value, dict):
35 return {k: plain(v) for k, v in value.items()}
36 if isinstance(value, (np.ndarray, list, tuple)):
37 return [plain(v) for v in value]
38 if isinstance(value, (np.integer, np.bool_)):
39 return value.item()
40 if isinstance(value, (np.floating, float)):
41 return float(value) if np.isfinite(value) else None
42 return value
43
44
45 def pooling_settings(name):
46 joint = name == "Matrix completion"
47 return {
48 "parameter": None,
49 "value": None,
50 "status": "not_applicable",
51 "fit_scope": "Joint untreated-cell fit"
52 if joint
53 else "Separate adoption-cohort fits",
54 "within_cohort_target": None if joint else "Mean treated outcome history",
55 "cross_cohort_partial_pooling": False,
56 "shared_structure": "Low-rank matrix and unit/time effects" if joint else None,
57 "note": "Regularization penalties are not partial-pooling parameters.",
58 }
59
60
61 def summarize_effects(corrected, effects, W, matrix=None):
62 controls = W.sum(axis=1) == 0
63 adoption = np.where(controls, -1, W.argmax(axis=1))
64 mass = np.ones_like(W) if matrix is None else matrix.mass
65 counts = np.ones_like(W, dtype=int) if matrix is None else matrix.counts
66 treated = (W == 1) & (mass > 0)
67 if not np.isfinite(corrected[treated]).all():
68 raise ValueError("Counterfactual missing for observed treated post outcomes")
69 scalar = float(np.average(corrected[treated], weights=mass[treated]))
70 relative = np.arange(W.shape[1])[None, :] - adoption[:, None]
71 path = []
72 for event in sorted(np.unique(relative[~controls])):
73 mask = (~controls[:, None]) & (relative == event) & (mass > 0)
74 # Some diagnostic pre-period predictions have no donor overlap. They
75 # are omitted, but every observed post outcome is required above.
76 mask &= np.isfinite(corrected) & np.isfinite(effects)
77 if not mask.any():
78 continue
79 path.append(
80 {
81 "event_time": int(event),
82 "estimate": float(np.average(corrected[mask], weights=mass[mask])),
83 "raw_gap": float(np.average(effects[mask], weights=mass[mask])),
84 "n": int(counts[mask].sum()),
85 "n_cells": int(mask.sum()),
86 "weight_sum": float(mass[mask].sum()),
87 }
88 )
89 baseline = next((row["estimate"] for row in path if row["event_time"] == -1), None)
90 for row in path:
91 row["relative_to_minus_one"] = (
92 row["estimate"] - baseline if baseline is not None else None
93 )
94 post = [row for row in path if row["event_time"] >= 0]
95 error = abs(
96 np.average(
97 [r["estimate"] for r in post], weights=[r["weight_sum"] for r in post]
98 )
99 - scalar
100 )
101 assert np.isfinite(scalar) and error < 1e-9
102 return scalar, path, float(error)
103
104
105 def correct_effects(name, result, W):
106 effects = np.asarray(result["treatment_effect"])
107 controls = W.sum(axis=1) == 0
108 adoption = np.where(controls, -1, W.argmax(axis=1))
109 cohorts = np.unique(adoption[~controls])
110 corrected = effects.copy()
111 correction = []
112 # SDID's native scalar subtracts a time-weighted pre-treatment gap.
113 # Its package event summary is the uncorrected synthetic-control gap.
114 if name == "SDID":
115 for j, cohort in enumerate(cohorts):
116 units = adoption == cohort
117 values = effects[units, :cohort]
118 available = np.isfinite(values)
119 gap = np.divide(
120 np.where(available, values, 0).sum(axis=0),
121 available.sum(axis=0),
122 out=np.full(cohort, np.nan),
123 where=available.sum(axis=0) > 0,
124 )
125 weights = np.asarray(result["time_weights"])[j, :cohort]
126 used = weights > 0
127 if not np.isfinite(gap[used]).all():
128 raise ValueError("SDID has a weighted pre-period without observed gaps")
129 pre_gap = float(gap[used] @ weights[used])
130 corrected[units] -= pre_gap
131 correction.append({"cohort": int(cohort), "pre_gap": pre_gap})
132 return effects, corrected, correction
133
134
135 def fit_one(name, model, parameters, Y, W, folder, matrix=None):
136 start = time.perf_counter()
137 model.fit(Y, W, observed=np.isfinite(Y) if matrix is None else matrix.observed)
138 result = model.summary()
139 effects, corrected, correction = correct_effects(name, result, W)
140 scalar = float(result["att"])
141 common_att = float(corrected[(W == 1) & np.isfinite(Y)].mean())
142 assert np.isfinite(scalar) and np.isfinite(common_att)
143 assert np.isclose(scalar, common_att, atol=1e-9)
144 native_att = scalar
145 scalar, path, aggregation_error = summarize_effects(corrected, effects, W, matrix)
146 diagnostics = {
147 k: result[k]
148 for k in [
149 "converged",
150 "iterations",
151 "termination_reason",
152 "objective",
153 "lambda_l",
154 "lambda_fraction",
155 "pre_rmse",
156 "singular_values",
157 "zeta_omega",
158 "zeta_lambda",
159 "penalty",
160 "history_objective",
161 "history_rmse",
162 "n_observed",
163 "n_missing",
164 "n_treated_observed",
165 "n_untreated_observed",
166 "missing_method",
167 "missing_diagnostics",
168 ]
169 if k in result
170 }
171 if name == "Matrix completion":
172 assert result["converged"], diagnostics
173 output = plain(
174 {
175 "method": name,
176 "class": type(model).__name__,
177 "parameters": parameters,
178 "partial_pooling": pooling_settings(name),
179 "estimate": scalar,
180 "event_study": path,
181 "package_event_study": result["event_study"]["weighted"],
182 "sdid_pre_gap_correction": correction,
183 "diagnostics": diagnostics,
184 "checks": {
185 "scalar_vs_treated_cell_average": aggregation_error,
186 "native_scalar_vs_unweighted_matrix_average": abs(
187 common_att - native_att
188 ),
189 },
190 "native_matrix_att": native_att,
191 "matrix_preparation": None if matrix is None else matrix.audit(),
192 "seconds": time.perf_counter() - start,
193 "crabbymetrics": version("crabbymetrics"),
194 "crabbymetrics_build": build_identity(),
195 "normalization": "Original source outcome / source SD; no mean subtraction",
196 "target": "All observed treated observations, using configured observation weights; cell mass equals sum of row weights",
197 }
198 )
199 slug = name.lower().replace(" ", "-")
200 (folder / f"{slug}.json").write_text(
201 json.dumps(output, indent=2, allow_nan=False) + "\n"
202 )
203 np.savez_compressed(folder / f"{slug}-effects-local.npz", effects=corrected)
204 print(
205 name,
206 scalar,
207 "seconds",
208 round(output["seconds"], 2),
209 "diagnostics",
210 {
211 k: v
212 for k, v in diagnostics.items()
213 if k in ["converged", "iterations", "termination_reason", "pre_rmse"]
214 },
215 flush=True,
216 )
217 return output
SHA-256: 3ae0740c5b75b14f563dacc9961688e8286655866a6b435045a3ef6f453c1604
1 # Cohort-level, partially pooled SCM from the unmodified augsynth package.
2 .libPaths(c(normalizePath('.Rlib'), .libPaths()))
3 suppressPackageStartupMessages(library(augsynth))
4 suppressPackageStartupMessages(library(jsonlite))
5
6 args <- commandArgs(trailingOnly = TRUE)
7 stopifnot(length(args) == 3L)
8 out <- args[1]
9 options <- fromJSON(args[2])
10 cohort_mode <- isTRUE(options$time_cohort)
11 label <- if (cohort_mode) 'Multisynth (cohort-averaged)' else 'Multisynth (individual)'
12 slug <- if (cohort_mode) 'multisynth-cohort' else 'multisynth-individual'
13 source(file.path(args[3], 'multisynth_sparse.R'))
14 panel <- read.csv(file.path(out, 'panel-local.csv'))
15 names(panel)[names(panel) == 'unit'] <- 'stock_id'
16 start <- proc.time()[['elapsed']]
17 # Record solver status without changing package source or optimization.
18 solver_checks <- list()
19 trace('solve_osqp', where = asNamespace('osqp'), print = FALSE,
20 exit = quote({
21 ans <- returnValue()
22 fields <- c('status', 'status_val', 'iter', 'obj_val',
23 'prim_res', 'dual_res', 'run_time')
24 solver_checks[[length(solver_checks) + 1L]] <<-
25 ans$info[intersect(fields, names(ans$info))]
26 }))
27
28 # Full observed horizon retains the existing all-treated-unit-year target.
29 # At this horizon the eligible donor set is exactly the never-treated units.
30 periods <- length(unique(panel$time))
31 cohorts <- sort(unique(panel$g[panel$g > 0]))
32 expected_units <- length(unique(panel$stock_id[panel$g > 0]))
33 expected_cells <- sum(panel$d)
34 expected_donors <- length(unique(panel$stock_id[panel$g == 0]))
35 expected_fits <- if (cohort_mode) length(cohorts) else expected_units
36 fit_model <- function() multisynth(
37 y ~ d, stock_id, time, panel,
38 n_leads = periods - min(cohorts) + 1L, n_lags = max(cohorts) - 1L,
39 nu = options$nu, lambda = options$lambda,
40 fixedeff = TRUE, scm = TRUE, time_cohort = cohort_mode,
41 eps_abs = options$eps_abs, eps_rel = options$eps_rel
42 )
43 fit <- if (cohort_mode) fit_model() else with_sparse_multisynth(fit_model())
44 untrace('solve_osqp', where = asNamespace('osqp'))
45 stopifnot(inherits(fit, 'multisynth'), identical(fit$time_cohort, cohort_mode),
46 length(fit$grps) == expected_fits, ncol(fit$weights) == expected_fits)
47 stopifnot(all(vapply(fit$donors, function(x) sum(x) == expected_donors, logical(1))))
48 stopifnot(all(vapply(fit$donors, function(x) {
49 all(!is.finite(fit$data$trt[x]))
50 }, logical(1))))
51 stopifnot(all(vapply(solver_checks, function(x) {
52 identical(x$status, 'solved')
53 }, logical(1))))
54
55
56 prediction <- predict(fit, att = TRUE)
57 events <- seq.int(-fit$n_lags, fit$n_leads - 1L)
58 stopifnot(nrow(prediction) == length(events) + 1L)
59 # Prediction columns are cohorts; preserve treated-unit counts in aggregation.
60 effects <- prediction[-nrow(prediction), -1L, drop = FALSE]
61 members <- if (cohort_mode) {
62 lapply(fit$grps, function(g) which(fit$data$trt == g))
63 } else {
64 as.list(which(is.finite(fit$data$trt)))
65 }
66 sizes <- vapply(members, length, integer(1))
67 stopifnot(sum(sizes) == expected_units)
68 counts <- as.numeric(is.finite(effects) %*% sizes)
69 weighted_effects <- sweep(effects, 2L, sizes, '*')
70 path <- data.frame(
71 event_time = events,
72 estimate = rowSums(weighted_effects, na.rm = TRUE) / counts,
73 n = counts
74 )
75 stopifnot(all(path$n > 0), all(is.finite(path$estimate)))
76 path$relative_to_minus_one <- path$estimate - path$estimate[path$event_time == -1]
77 post <- path$event_time >= 0
78 att <- weighted.mean(path$estimate[post], path$n[post])
79 stopifnot(sum(path$n[post]) == expected_cells)
80 stopifnot(max(abs(path$estimate - prediction[-nrow(prediction), 1L])) < 1e-10)
81
82 # Reconstruct cohort effects in calendar time from weights and fixed effects.
83 Y <- cbind(fit$data$X, fit$data$y)
84 post_sum <- 0
85 post_count <- 0L
86 path_difference <- 0
87 cohort_paths <- list()
88 cohort_diagnostics <- list()
89 for (j in seq_along(fit$grps)) {
90 g <- fit$grps[j]
91 treated <- members[[j]]
92 counterfactual <- colMeans(fit$y0hat[[j]][treated, , drop = FALSE]) +
93 as.numeric(crossprod(fit$weights[, j], fit$residuals[[j]]))
94 gap <- colMeans(Y[treated, , drop = FALSE]) - counterfactual
95 observed <- (g + 1L):ncol(Y)
96 post_sum <- post_sum + sizes[j] * sum(gap[observed])
97 post_count <- post_count + sizes[j] * length(observed)
98 available <- is.finite(effects[, j])
99 calendar <- events[available] + g + 1L
100 path_difference <- max(path_difference,
101 abs(gap[calendar] - effects[available, j]))
102 cohort_paths[[as.character(j)]] <- data.frame(
103 cohort = g, event_time = events[available],
104 estimate = effects[available, j], n = sizes[j]
105 )
106 }
107 for (g in sort(unique(fit$grps))) {
108 cohort_diagnostics[[as.character(g + 1L)]] <- list(
109 adoption_period = g + 1L,
110 treated_units = sum(fit$data$trt == g),
111 donor_units = expected_donors, nu = fit$nu, lambda = fit$lambda,
112 time_cohort = fit$time_cohort,
113 pooling_scope = paste('One global nu shared by', expected_fits, 'fits')
114 )
115 }
116
117 stopifnot(post_count == expected_cells, abs(post_sum / post_count - att) < 1e-10,
118 path_difference < 1e-10)
119
120 result <- list(
121 method = label, class = 'augsynth::multisynth',
122 parameters = list(
123 nu = fit$nu, nu_selection = if (is.null(options$nu)) 'Package data-derived default' else 'Configured value',
124 lambda = fit$lambda, fixedeff = TRUE, scm = fit$scm,
125 time_cohort = fit$time_cohort, n_leads = fit$n_leads,
126 n_lags = fit$n_lags, eps_abs = options$eps_abs, eps_rel = options$eps_rel,
127 V = 'Identity (package default)',
128 matrix_assembly = if (cohort_mode) 'Unmodified package default' else 'Equivalent sparse matrix construction'
129 ),
130 partial_pooling = list(
131 parameter = 'nu', value = fit$nu, status = 'used',
132 fit_scope = paste('Joint optimization of', expected_fits, if (cohort_mode) 'cohort fits' else 'unit fits'),
133 cross_cohort_partial_pooling = TRUE,
134 note = 'nu=0 weights separate cohort imbalance only; nu=1 weights pooled imbalance only. One shared nu, not a cohort-specific tuning parameter.'
135 ),
136 estimate = att, event_study = path,
137 cohort_event_study = aggregate(estimate ~ cohort + event_time, do.call(rbind, cohort_paths), mean),
138 diagnostics = cohort_diagnostics,
139 fit_diagnostics = list(
140 cohort_fits = length(fit$grps), treated_units = sum(sizes),
141 global_l2 = fit$global_l2, ind_l2 = fit$ind_l2,
142 scaled_global_l2 = fit$scaled_global_l2,
143 scaled_ind_l2 = fit$scaled_ind_l2,
144 sum_weights_range = range(colSums(fit$weights)),
145 minimum_weight = min(fit$weights),
146 maximum_weight = max(fit$weights), solver = solver_checks
147 ),
148 checks = list(
149 treated_cells = post_count,
150 direct_cell_average_difference = abs(post_sum / post_count - att),
151 direct_cohort_path_difference = path_difference,
152 package_equal_unit_att = prediction[nrow(prediction), 1L],
153 aggregation = 'Equal treated-cell weights; not the package final ATT row'
154 ),
155 seconds = proc.time()[['elapsed']] - start,
156 augsynth = as.character(packageVersion('augsynth')),
157 augsynth_commit = packageDescription('augsynth')$RemoteSha,
158 R_version = R.version.string,
159 target = 'All observed treated unit-years, equal cell weights',
160 donor_rule = paste(expected_donors, 'never-treated donors per fit; masks verified')
161 )
162 write_json(result, file.path(out, paste0(slug, '.json')),
163 pretty = TRUE, auto_unbox = TRUE, digits = NA, na = 'null')
164 cat('multisynth:', att, 'nu:', fit$nu, 'lambda:', fit$lambda,
165 'seconds:', result$seconds, '\n')
166
167 # Retain the completed fit locally for inexpensive bootstrap inference.
168 saveRDS(fit, file.path(out, paste0(slug, '-fit-local.rds')))
SHA-256: f9ab3d8d7f7b52a79ccd694c7173f626276a0830f23bb79e527c88c0f0292540
1 # Keep the package's identical QP objective sparse in this R call only.
2 # Dense V'V makes kronecker() allocate a 26,450-square dense block here.
3 with_sparse_multisynth <- function(expr) {
4 ns <- asNamespace('augsynth')
5 original <- get('make_Pmat', ns)
6 sparse <- original
7 old <- quote(Vsq <- t(V) %*% V)
8 new <- quote(Vsq <- Matrix::Matrix(t(V) %*% V, sparse = TRUE))
9 parts <- as.list(body(original))
10 matched <- vapply(parts, identical, logical(1), y = old)
11 stopifnot(sum(matched) == 1L)
12 parts[matched] <- list(new)
13 body(sparse) <- as.call(parts)
14 assignInNamespace('make_Pmat', sparse, ns = 'augsynth')
15 on.exit(assignInNamespace('make_Pmat', original, ns = 'augsynth'))
16 force(expr)
17 }
SHA-256: a55b171313d43a703063c87e668021b75eb3153fa30795b9ae3d9e2721d9b9d0
1 """Full country-cluster pairs bootstrap for the three matrix estimators."""
2
3 import argparse
4 import hashlib
5 import json
6 import os
7 import time
8 from pathlib import Path
9
10 import numpy as np
11 import pandas as pd
12 from matrix_inputs import prepare_matrix
13 from matrix_models import correct_effects
14
15 import crabbymetrics as cm
16
17 parser = argparse.ArgumentParser(description=__doc__)
18 parser.add_argument("folder", type=Path, help="Folder containing panel-local.csv")
19 OUT = parser.parse_args().folder.resolve()
20 B = int(os.environ.get("CREDIT_BOOTSTRAP_REPS", "999"))
21 SEED = 20261009
22 folder = OUT / "inference"
23 folder.mkdir(exist_ok=True)
24 data = pd.read_csv(OUT / "panel-local.csv")
25 assert data.groupby("unit").cluster.nunique().max() == 1, (
26 "Pairs bootstrap requires a stable cluster within matrix unit"
27 )
28 matrix = prepare_matrix(data)
29 assert matrix.observed.all(), "Bootstrap requires complete cell-mean matrices"
30 units = data.drop_duplicates("unit").set_index("unit").reindex(matrix.units)
31 Y, W = matrix.y, matrix.w
32 clusters = units.cluster.to_numpy()
33 groups = np.unique(clusters)
34 members = [np.flatnonzero(clusters == group) for group in groups]
35 rng = np.random.default_rng(SEED)
36 names = ["SDID", "Matrix completion", "Horizontal ridge"]
37 draw_file = folder / "country-bootstrap-draws.csv"
38 receipt = {
39 "replications": B,
40 "seed": SEED,
41 "resampling_unit": "country",
42 "clusters": len(groups),
43 "country_draws_per_replication": len(groups),
44 "sample_sha256": hashlib.sha256((OUT / "panel-local.csv").read_bytes()).hexdigest(),
45 "code_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
46 "all_units_and_years_retained_within_sampled_country": True,
47 "refit_weights_and_data_dependent_tuning": True,
48 "cohort_shares": "Recomputed from each country-resampled sample",
49 "scale": "Fixed original source-SD denominator",
50 "matrix_preparation": matrix.audit(),
51 }
52 existing = []
53 settings_file = folder / "country-bootstrap-settings.json"
54 if draw_file.exists():
55 assert json.loads(settings_file.read_text()) == receipt
56 existing = pd.read_csv(draw_file).to_dict("records")
57 else:
58 settings_file.write_text(json.dumps(receipt, indent=2) + "\n")
59 started = time.perf_counter()
60 for b in range(B):
61 indices = np.concatenate(
62 [members[j] for j in rng.integers(len(groups), size=len(groups))]
63 )
64 if b < len(existing):
65 continue
66 y, w = Y[indices], W[indices]
67 assert np.any(w.sum(axis=1) == 0) and w.sum() > 0
68 row = {"replication": b, "units": len(indices), "treated_cells": int(w.sum())}
69 for name in names:
70 if name == "SDID":
71 model = cm.SyntheticDID()
72 elif name == "Matrix completion":
73 model = cm.MatrixCompletion(
74 lambda_fraction=0.25, max_iterations=5000, tolerance=1e-12
75 )
76 else:
77 model = cm.HorizontalPanelRidge(penalty=1)
78 model.fit(y, w)
79 fit = model.summary()
80 assert np.isfinite(fit["att"]), (b, name)
81 if name == "Matrix completion":
82 assert fit["converged"], (b, name, fit["termination_reason"])
83 _, corrected, _ = correct_effects(name, fit, w)
84 row[name] = float(
85 np.average(corrected[w == 1], weights=matrix.mass[indices][w == 1])
86 )
87 pd.DataFrame([row]).to_csv(
88 draw_file, mode="a", header=not draw_file.exists(), index=False
89 )
90 if (b + 1) % 50 == 0:
91 print(
92 b + 1, "of", B, "seconds", round(time.perf_counter() - started), flush=True
93 )
94 draws = pd.read_csv(draw_file)
95 assert len(draws) == B and draws.replication.tolist() == list(range(B))
96 results = []
97 for name in names:
98 lower, upper = np.quantile(draws[name], [0.025, 0.975])
99 results.append(
100 {
101 "method": name,
102 "se": float(draws[name].std(ddof=1)),
103 "lower": float(lower),
104 "upper": float(upper),
105 "inference": "95% percentile country-cluster pairs bootstrap; all model weights and data-dependent tuning refitted",
106 "replications": B,
107 "failed_or_discarded_draws": 0,
108 }
109 )
110 (folder / "matrix-confidence-intervals.json").write_text(
111 json.dumps({"settings": receipt, "results": results}, indent=2) + "\n"
112 )
113 print(results)
SHA-256: 271442d1ce32a753b9051c2b6fdab14ffc2f8babd1ebcd04ff43e29b24a2dd41
1 # Fixed-fit wild bootstrap matching augsynth::weighted_bootstrap_multi.
2 # No jackknife or QP refits. Shared multipliers preserve configured clusters.
3 bootstrap_multisynth <- function(fit, panel, result, n_boot = 9999L,
4 seed = 20261009L, alpha = .05) {
5 stopifnot(inherits(fit, 'multisynth'), fit$force == 3)
6 if (!'unit' %in% names(panel)) names(panel)[names(panel) == 'stock_id'] <- 'unit'
7 unit_cluster <- unique(panel[c('unit', 'cluster')])
8 stopifnot(!anyDuplicated(unit_cluster$unit))
9 ids <- fit$data$units
10 clusters <- unit_cluster$cluster[match(ids, unit_cluster$unit)]
11 stopifnot(!anyNA(clusters), length(clusters) == nrow(fit$data$X))
12 prediction <- predict(fit, att = TRUE)
13 effects <- prediction[-nrow(prediction), -1L, drop = FALSE]
14 events <- seq.int(-fit$n_lags, fit$n_leads - 1L)
15 members <- if (fit$time_cohort) {
16 lapply(fit$grps, function(g) which(fit$data$trt == g))
17 } else as.list(which(is.finite(fit$data$trt)))
18 sizes <- vapply(members, length, integer(1))
19 counts <- as.numeric(is.finite(effects) %*% sizes)
20 post <- events >= 0
21 target_weights <- ifelse(post, counts, 0) / sum(counts[post])
22 point <- unname(prediction[-nrow(prediction), 1L])
23 estimate <- sum(point * target_weights)
24 stopifnot(abs(estimate - result$estimate) < 1e-9,
25 sum(counts[post]) == sum(panel$d))
26
27 # predict.multisynth is linear in bs_weight. Construct its exact unit
28 # contributions once, then multiply 9,999 draws without repeating J fits.
29 Y <- cbind(fit$data$X, fit$data$y)
30 contribution <- matrix(0, nrow(Y), length(events))
31 for (j in seq_along(members)) {
32 available <- which(is.finite(effects[, j]))
33 calendar <- events[available] + fit$grps[j] + 1L
34 y0 <- if (is.list(fit$y0hat)) fit$y0hat[[j]] else fit$y0hat
35 residual <- if (is.list(fit$residuals)) fit$residuals[[j]] else fit$residuals
36 unit_effect <- -sweep(residual[, calendar, drop = FALSE], 1,
37 fit$weights[, j], '*')
38 treated <- members[[j]]
39 unit_effect[treated, ] <- unit_effect[treated, , drop = FALSE] +
40 (Y[treated, calendar, drop = FALSE] -
41 y0[treated, calendar, drop = FALSE]) / sizes[j]
42 contribution[, available] <- contribution[, available, drop = FALSE] +
43 sweep(unit_effect, 2, sizes[j] / counts[available], '*')
44 }
45 point_error <- max(abs(colSums(contribution) - point))
46 stopifnot(point_error < 1e-9)
47 # Exactly retain the package's centering, including sum(Z)/n_treated.
48 centered <- sweep(contribution, 2, point / sum(sizes), '-')
49 grouped <- rowsum(centered, clusters, reorder = TRUE)
50 cluster_index <- match(as.character(clusters), rownames(grouped))
51 set.seed(seed)
52 multipliers <- replicate(n_boot, augsynth:::rwild_b(nrow(grouped)))
53 if (is.null(dim(multipliers))) multipliers <- matrix(multipliers, nrow = 1L)
54 draws <- crossprod(grouped, multipliers)
55 scalar_draws <- as.numeric(crossprod(target_weights, draws))
56
57 # Validate actual sampled paths and the package's pointwise output under
58 # identical clustered multipliers (not an unrelated approximate bootstrap).
59 check_draws <- min(7L, n_boot)
60 draw_error <- 0
61 for (b in seq_len(check_draws)) {
62 z <- multipliers[cluster_index, b]
63 native <- predict(fit, att = TRUE, bs_weight = z) -
64 sum(z) / sum(sizes) * prediction
65 draw_error <- max(draw_error,
66 abs(native[-nrow(native), 1L] - draws[, b]))
67 }
68 draw_counter <- 0L
69 rweight <- function(n) {
70 stopifnot(n == length(cluster_index))
71 draw_counter <<- draw_counter + 1L
72 multipliers[cluster_index, draw_counter]
73 }
74 native <- augsynth:::weighted_bootstrap_multi(
75 fit, rweight = rweight, n_boot = check_draws, alpha = alpha)
76 se_of <- function(x) sqrt(mean((x - mean(x))^2))
77 native_error <- max(abs(native$se[-nrow(prediction), 1L] -
78 apply(draws[, seq_len(check_draws), drop = FALSE], 1, se_of)))
79 stopifnot(draw_error < 1e-9, native_error < 1e-9)
80 bounds <- function(point, draws) {
81 c(se = se_of(draws), lower = point - unname(quantile(draws, 1-alpha/2)),
82 upper = point - unname(quantile(draws, alpha/2)))
83 }
84 scalar <- bounds(estimate, scalar_draws)
85 event_bounds <- t(vapply(seq_along(events), function(k) {
86 bounds(point[k], draws[k, ])
87 }, numeric(3)))
88 event_path <- data.frame(event_time = events, estimate = point, n = counts,
89 event_bounds, check.names = FALSE)
90 list(estimate = estimate, se = scalar[['se']], lower = scalar[['lower']],
91 upper = scalar[['upper']], event_study = event_path,
92 inference = paste0('95% configured-cluster wild bootstrap, ', n_boot,
93 ' Mammen draws; augsynth fixed-weight prediction bootstrap; no jackknife/refits'),
94 settings = list(n_boot = n_boot, seed = seed, alpha = alpha,
95 clusters = nrow(grouped), units = nrow(Y), treated_units = sum(sizes),
96 treated_cells = sum(counts[post]), time_cohort = fit$time_cohort,
97 nu = fit$nu, lambda = fit$lambda,
98 aggregation = 'Equal treated-unit-year weights; recomputed for every bootstrap draw',
99 centering = 'Exact package predict(bs_weight=Z) - sum(Z)/n_treated * predict()',
100 reestimate_weights = FALSE, reestimate_tuning = FALSE, jackknife = FALSE,
101 note = 'Conditional on fitted weights, pooling/tuning and observed design; not a full-refit pairs bootstrap. Event intervals are pointwise, not simultaneous.'),
102 checks = list(point_max_abs_error = point_error,
103 native_draw_max_abs_error = draw_error, native_se_max_abs_error = native_error),
104 scalar_bootstrap_draws = scalar_draws,
105 augsynth = as.character(packageVersion('augsynth')))
106 }
SHA-256: 94e9c5116a13335ead2db81834cac48593160291aea51c30c12fe3e11e5da1ce
1 """Reporting-only scalar intervals, forest figures and source provenance.
2
3 No estimator is fitted here. Cached inference must match the prepared sample,
4 point-fit settings, package version and event path before it is attached.
5 """
6
7 import hashlib
8 import json
9 import math
10 import re
11 from pathlib import Path
12
13 import matplotlib.pyplot as plt
14
15 SIGNATURE_FIELDS = (
16 "method",
17 "parameters",
18 "estimate",
19 "event_study",
20 "crabbymetrics",
21 "target",
22 )
23
24
25 def sha256(path):
26 return hashlib.sha256(Path(path).read_bytes()).hexdigest()
27
28
29 def fit_signature(model):
30 payload = {k: model.get(k) for k in SIGNATURE_FIELDS}
31 if payload.get("event_study"):
32 payload["event_study"] = [
33 {k: v for k, v in row.items() if k not in ("se", "lower", "upper")}
34 for row in payload["event_study"]
35 ]
36 return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
37
38
39 def attach_intervals(run):
40 """Attach explicitly configured inference; reject stale or unrelated results."""
41 for key, relative in run.options.get("inference_files", {}).items():
42 if key not in run.base.panels:
43 raise ValueError(f"Inference specification is absent: {key}")
44 path = run.root / relative
45 saved = json.loads(path.read_text())
46 data, _ = run.base.panels[key]
47 panel_hash = hashlib.sha256(data.to_csv(index=False).encode()).hexdigest()
48 if saved["settings"]["sample_sha256"] != panel_hash:
49 raise ValueError(
50 f"Inference sample mismatch: {key}; remove inference_files or recompute"
51 )
52 by_method = {r["method"]: r for r in saved["results"]}
53 for model in run.extra.get(key, []):
54 entry = by_method.get(model["method"])
55 if entry is None or model.get("status") != "estimated":
56 continue
57 if entry["fit_signature"] != fit_signature(model):
58 raise ValueError(f"Inference fit mismatch: {key}, {model['method']}")
59 model.update({k: entry[k] for k in ("se", "lower", "upper", "inference")})
60 if entry.get("event_bounds"):
61 bounds = {r["event_time"]: r for r in entry["event_bounds"]}
62 for row in model["event_study"]:
63 row.update(
64 {
65 k: bounds[row["event_time"]][k]
66 for k in ("se", "lower", "upper")
67 }
68 )
69 if entry.get("bootstrap"):
70 model["bootstrap"] = entry["bootstrap"]
71 model["inference_provenance"] = {
72 "path": relative,
73 "sha256": sha256(path),
74 "settings": saved["settings"],
75 "source": saved["source"],
76 }
77
78
79 def scalar_ranges(run):
80 """Join table rows to native bounds; never manufacture CIs from missing SEs."""
81 lookup = {}
82 for key, spec in run.base.results["specifications"].items():
83 for model in spec.get("estimates", []) + run.extra.get(key, []):
84 lookup[key, model["method"]] = model
85 if model["method"] == "CS-DID" and model.get("aggregations"):
86 simple = next(
87 r
88 for r in model["aggregations"]["summaries"]
89 if r.get("matches_treated_unit_year_target", False)
90 )
91 lookup[key, "CS-DID (treated cells)"] = simple | {
92 "inference": model["inference"]
93 }
94 rows = []
95 for row in run.rows():
96 model = lookup[row["Outcome / treatment"], row["Estimator"]]
97 if row["Aggregation"] == "observations":
98 row["Aggregation"] = "Treated cells"
99 if row["Estimator"] == "CS-DID":
100 row["Estimator"] = "CS-DID (equal event-time)"
101 valid = row["Status"] == "estimated" and finite(row["Estimate (SDs)"])
102 lower, upper = model.get("lower"), model.get("upper")
103 has_ci = valid and finite(lower) and finite(upper) and lower <= upper
104 row.update(
105 lower=lower if has_ci else None,
106 upper=upper if has_ci else None,
107 ci_level=0.95 if has_ci else None,
108 inference=model.get("inference", "Configured-cluster regression interval")
109 if has_ci
110 else "Not computed",
111 )
112 rows.append(row)
113 return rows
114
115
116 def finite(value):
117 return value is not None and math.isfinite(value)
118
119
120 def forest_plots(run):
121 """One full-width, exportable forest figure per outcome-treatment pair."""
122 rows = scalar_ranges(run)
123 paths = []
124 for key in run.base.results["specifications"]:
125 subset = [r for r in rows if r["Outcome / treatment"] == key]
126 fig, ax = plt.subplots(figsize=(11, max(4, 0.61 * len(subset) + 1.4)))
127 labels = []
128 for y, row in enumerate(subset):
129 labels.append(f"{row['Estimator']}\n{row['Aggregation']}")
130 point, lo, hi = row["Estimate (SDs)"], row["lower"], row["upper"]
131 valid = row["Status"] == "estimated" and finite(point)
132 color = "#14819a" if row["Estimator"].startswith("CS-DID") else "#34495e"
133 if valid:
134 if lo is not None:
135 ax.hlines(y, lo, hi, color=color, lw=1.6)
136 ax.plot([lo, hi], [y, y], "|", color=color, ms=7)
137 ax.plot(
138 point,
139 y,
140 "o",
141 color=color,
142 mfc=color if lo is not None else "white",
143 ms=6,
144 )
145 text = f"{point:+.3f}" + (" [CI not computed]" if lo is None else "")
146 else:
147 text = f"Not estimated ({row['Status']})"
148 ax.text(
149 1.02,
150 y,
151 text,
152 transform=ax.get_yaxis_transform(),
153 va="center",
154 fontsize=9,
155 )
156 ax.set_yticks(range(len(subset)), labels, fontsize=9)
157 ax.set_ylim(len(subset) - 0.4, -0.7)
158 ax.axvline(0, color="0.65", lw=0.8)
159 ax.set_axisbelow(True)
160 ax.grid(axis="x", alpha=0.14)
161 ax.set(
162 xlabel="Estimate (outcome SDs); available 95% confidence intervals",
163 title=f"{key}: scalar estimates and aggregation targets",
164 )
165 ax.spines[["top", "right", "left"]].set_visible(False)
166 ax.tick_params(axis="y", length=0, pad=12)
167 ax.margins(x=0.09)
168 fig.subplots_adjust(left=0.28, right=0.77, top=0.92, bottom=0.10)
169 stem = run.output / ("forest-" + re.sub(r"[^A-Za-z0-9_-]", "_", key))
170 for suffix in (".png", ".pdf"):
171 path = stem.with_suffix(suffix)
172 fig.savefig(path, dpi=180, bbox_inches="tight", facecolor="white")
173 paths.append(str(path.relative_to(run.root)))
174 plt.show()
175 plt.close(fig)
176 return paths
177
178
179 def export_report(run):
180 path = run.export()
181 payload = json.loads(path.read_text())
182 payload["scalar_ranges"] = scalar_ranges(run)
183 payload["reporting_source"] = {
184 "path": "templates/reporting/comparison_report.py",
185 "sha256": sha256(__file__),
186 }
187 path.write_text(json.dumps(payload, indent=2, allow_nan=False) + "\n")
188 return path