"""

Real-world datasets for Difference-in-Differences analysis.

This module provides functions to load classic econometrics datasets

commonly used for teaching and demonstrating DiD methods.

Canonical data are downloaded from checksum-pinned public sources and cached

locally. A download that fails verification falls back to a checksum-valid

cache entry when one exists, so verified data on disk is never displaced by

generated data. Only when no verified copy is available does the loader warn

and return an explicitly provenance-marked synthetic fallback.

"""

import hashlib

import os

import sys

import warnings

from http.client import HTTPException

from io import BytesIO, StringIO

from pathlib import Path

from tempfile import NamedTemporaryFile

from typing import Callable, Dict, Optional, cast

from urllib.error import HTTPError, URLError

from urllib.request import urlopen

import numpy as np

import pandas as pd

# Cache directory for downloaded datasets

_CACHE_DIR = Path.home() / ".cache" / "diff_diff" / "datasets"

_MAX_DATASET_BYTES = 50 * 1024 * 1024

# These commit-pinned mirrors were verified byte-equivalent to their authoritative

# sources at pinning time: ``public.dat`` is byte-identical to the copy in

# ``njmin.zip`` from David Card's data archive, and ``mpdta.csv`` matches

# ``did::mpdta`` (R package 2.5.1) to CSV round-trip precision (max absolute

# difference 1.07e-14 on ``lemp``). Every cached and downloaded byte sequence is

# verified against the pinned SHA-256 below, so a later change at any mirror can

# never substitute different data: it falls back to the verified cache entry if one

# exists (with a warning naming the integrity failure), and to the loud synthetic

# fallback only when there is no verified copy to fall back to.

_CARD_KRUEGER_SOURCE_URL = (

"https://raw.githubusercontent.com/rafiash/CardKrueger-stata-sample/"

"07bc929f1d6552db117bd27a7cf0d881d16e9494/public.dat"

)

_CARD_KRUEGER_SOURCE_SHA256 = "04bde0cad5540980f32ce099c6dad369e2f05494698071d8a65b3e1cbe9ca53a"

_CASTLE_DOCTRINE_SOURCE_URL = (

"https://raw.githubusercontent.com/scunning1975/mixtape/"

"ca4279a87a6f0759f6b6f02841a53bdd68e27d3c/castle.dta"

)

_CASTLE_DOCTRINE_SOURCE_SHA256 = "804633c161827b6c0824f86f239046386d1a8266a866f83bf5ddb2aa762a5f29"

_MPDTA_SOURCE_URL = (

"https://raw.githubusercontent.com/d2cml-ai/csdid/"

"7ad707385354cb3924b8da94ef7e62a76bf55a4d/data/mpdta.csv"

)

_MPDTA_SOURCE_SHA256 = "2283bea1221a152420f98dfa20f633c5d054ea51d881115c8cd702a97bcd3167"

# ``sid`` follows alphabetical state order with 9 reserved for Washington, DC,

# which is absent from the source panel.

_CASTLE_STATE_BY_SID = dict(

enumerate(

"""

AL AK AZ AR CA CO CT DE _ FL GA HI ID IL IN IA KS KY LA ME MD MA MI MN

MS MO MT NE NV NH NJ NM NY NC ND OH OK OR PA RI SC SD TN TX UT VT VA WA

WV WI WY

""".split(),

start=1,

)

)

class _DatasetSourceError(RuntimeError):

"""Expected failure while fetching, parsing, or validating canonical data."""

def _caller_stacklevel() -> int:

"""``stacklevel`` that attributes a warning to the first frame outside this module.

The text and binary loaders reach the download helper at different depths (the text

path routes through ``_load_verified_dataset`` and a source adapter; the binary path

calls it almost directly), so no fixed ``stacklevel`` points at user code for both.

"""

level = 1

try:

frame = sys._getframe(1)

except ValueError: # pragma: no cover - defensive

return 2

while frame is not None:

if frame.f_globals.get("__name__") != __name__:

return level

frame = frame.f_back

level += 1

return 2 # pragma: no cover - defensive

def _get_cache_path(name: str) -> Path:

"""Get the cache path for a dataset."""

return _CACHE_DIR / f"{name}.csv"

def _download_with_cache(

url: str,

name: str,

sha256: str,

force_download: bool = False,

) -> str:

"""Download UTF-8 text, verify its checksum, and cache it."""

cache_path = _get_cache_path(name)

content = _download_verified_bytes(url, name, sha256, cache_path, force_download)

try:

return content.decode("utf-8")

except UnicodeDecodeError as e:

raise _DatasetSourceError(

f"Dataset '{name}' passed its byte checksum but is not valid UTF-8 text."

) from e

def _read_verified_cache(cache_path: Path, sha256: str) -> Optional[bytes]:

"""Return a bounded, checksum-valid cache entry or None."""

try:

if not cache_path.exists() or cache_path.stat().st_size > _MAX_DATASET_BYTES:

return None

content = cache_path.read_bytes()

except OSError:

return None

if hashlib.sha256(content).hexdigest() == sha256:

return content

return None

def _write_cache_atomically(cache_path: Path, content: bytes, name: str) -> None:

"""Replace a cache entry only after a complete same-directory write."""

temp_path: Optional[Path] = None

try:

cache_path.parent.mkdir(parents=True, exist_ok=True)

with NamedTemporaryFile(

mode="wb",

dir=cache_path.parent,

prefix=f".{cache_path.name}.",

delete=False,

) as temp_file:

# Bind the path before writing: ``delete=False`` means a failed write

# still leaves the file on disk, and the handler below can only clean

# up what it knows about.

temp_path = Path(temp_file.name)

temp_file.write(content)

os.replace(temp_path, cache_path)

except OSError as e:

if temp_path is not None:

try:

temp_path.unlink()

except OSError:

pass

raise _DatasetSourceError(f"Failed to cache dataset '{name}': {e}") from e

def _download_verified_bytes(

url: str,

name: str,

sha256: str,

cache_path: Path,

force_download: bool = False,

) -> bytes:

"""Return checksum-verified bytes from cache or a fresh download.

The cache is read up front and retained even under ``force_download``, so that

EVERY way a fresh download can fail verification - transport, size limit, or

checksum - can still fall back to bytes that already passed the pinned hash.

Falling through to the synthetic frame while verified canonical bytes sit on

disk would be a downgrade, and on the checksum path it would let a tampered

or moved upstream quietly replace real data with generated data.

"""

cached = _read_verified_cache(cache_path, sha256)

if cached is not None and not force_download:

return cached

def _recover(message: str, cause: Optional[BaseException] = None) -> bytes:

"""Prefer verified cached bytes over failing into the synthetic fallback."""

if cached is not None:

return cached

raise _DatasetSourceError(message) from cause

try:

with urlopen(url, timeout=30) as response:

content = response.read(_MAX_DATASET_BYTES + 1)

# ``HTTPException`` is the parent of ``IncompleteRead``, ``BadStatusLine`` and the

# rest of the protocol-level errors ``urlopen`` can surface; none of them derive

# from ``OSError``, so catching the base class keeps the whole family inside the

# documented warn-and-fall-back boundary rather than only the named siblings.

except (HTTPError, HTTPException, OSError, TimeoutError, URLError) as e:

return _recover(

f"Failed to download dataset '{name}' from {url}: {e}\n"

"Check your internet connection or try again later.",

e,

)

if len(content) > _MAX_DATASET_BYTES:

return _recover(

f"Dataset '{name}' downloaded from {url} exceeds the "

f"{_MAX_DATASET_BYTES}-byte safety limit."

)

if hashlib.sha256(content).hexdigest() != sha256:

if cached is not None:

# Canonical bytes are already on disk, so the user keeps real data - but a

# pin mismatch is an integrity event, not a transport hiccup, and must not

# pass unnoticed. Deliberately not a SYNTHETIC warning: nothing synthetic

# is involved, and callers key their provenance checks on that word.

warnings.warn(

f"Upstream copy of dataset '{name}' at {url} no longer matches the "

"pinned SHA-256. Returning the verified cached copy instead. Verify "

"the source revision before updating the pinned checksum; until then "

"treat the upstream file as untrusted.",

UserWarning,

stacklevel=_caller_stacklevel(),

)

return cached

raise _DatasetSourceError(

f"Checksum mismatch for dataset '{name}' downloaded from {url}.\n"

"The upstream file differs from the pinned SHA-256. Verify the "

"source revision before updating the pinned checksum; otherwise "

"treat the download as untrusted."

)

try:

_write_cache_atomically(cache_path, content, name)

except _DatasetSourceError:

# Cache persistence is best-effort after the downloaded bytes have

# already passed the pinned SHA-256 check.

pass

return content

def _get_cache_path_binary(name: str) -> Path:

"""Get the cache path for a binary (Stata .dta) dataset."""

return _CACHE_DIR / f"{name}.dta"

def _download_with_cache_binary(

url: str,

name: str,

sha256: str,

force_download: bool = False,

) -> bytes:

"""Download a binary file (e.g. Stata .dta), verify its checksum, and cache it.

Every byte-load (cache or fresh download) is verified against a pinned

SHA-256. A stale or corrupt cache triggers one re-download. A checksum

mismatch on freshly downloaded bytes falls back to a verified cache entry

when one exists (warning that the upstream no longer matches the pin), and

raises only when there is no verified copy to fall back to.

"""

return _download_verified_bytes(

url,

name,

sha256,

_get_cache_path_binary(name),

force_download,

)

def _load_verified_dataset(

*,

cache_name: str,

source: str,

force_download: bool,

load_source: Optional[Callable[[bool], pd.DataFrame]],

prepare: Callable[[pd.DataFrame], pd.DataFrame],

validate_source: Callable[[pd.DataFrame], None],

validate_fallback: Callable[[pd.DataFrame], None],

fallback: Callable[[], pd.DataFrame],

) -> pd.DataFrame:

"""Load and validate canonical data or return a loud synthetic fallback."""

try:

if load_source is None:

raise _DatasetSourceError("no verified canonical source is configured")

df = prepare(load_source(force_download))

validate_source(df)

except _DatasetSourceError as exc:

warnings.warn(

f"{cache_name} canonical data are unavailable ({exc}); returning a "

"SYNTHETIC fallback. Check `df.attrs['source']` before treating "

"the result as replication data.",

UserWarning,

stacklevel=3,

)

df = fallback()

validate_fallback(df)

df.attrs["source"] = "synthetic_fallback"

return df

df.attrs["source"] = source

return df

def _load_card_krueger_source(force_download: bool) -> pd.DataFrame:

"""Load the checksum-pinned Card-Krueger public flat file."""

content = _download_with_cache(

_CARD_KRUEGER_SOURCE_URL,

"card_krueger",

_CARD_KRUEGER_SOURCE_SHA256,

force_download,

)

columns = """

sheet chain co_owned state southj centralj northj pa1 pa2 shore

ncalls empft emppt nmgrs wage_st inctime firstinc bonus pctaff meals

open hrsopen psoda pfry pentree nregs nregs11 type2 status2 date2

ncalls2 empft2 emppt2 nmgrs2 wage_st2 inctime2 firstin2 special2

meals2 open2r hrsopen2 psoda2 pfry2 pentree2 nregs2 nregs112

""".split()

try:

return pd.read_csv(

StringIO(content),

sep=r"\s+",

names=columns,

na_values=".",

)

except (TypeError, ValueError) as e:

raise _DatasetSourceError(f"Failed to parse Card-Krueger source data: {e}") from e

def _load_castle_doctrine_source(force_download: bool) -> pd.DataFrame:

"""Load the checksum-pinned Cheng-Hoekstra Stata data."""

content = _download_with_cache_binary(

_CASTLE_DOCTRINE_SOURCE_URL,

"castle_doctrine",

_CASTLE_DOCTRINE_SOURCE_SHA256,

force_download,

)

try:

return pd.read_stata(BytesIO(content))

except (OSError, TypeError, ValueError) as e:

raise _DatasetSourceError(f"Failed to parse Castle Doctrine source data: {e}") from e

def _load_mpdta_source(force_download: bool) -> pd.DataFrame:

"""Load the checksum-pinned Callaway-Sant'Anna example data."""

content = _download_with_cache(

_MPDTA_SOURCE_URL,

"mpdta",

_MPDTA_SOURCE_SHA256,

force_download,

)

try:

return pd.read_csv(StringIO(content))

except (TypeError, ValueError) as e:

raise _DatasetSourceError(f"Failed to parse mpdta source data: {e}") from e

def _require_columns(df: pd.DataFrame, dataset: str, columns: set) -> None:

"""Reject empty or structurally incomplete downloaded datasets."""

if df.empty:

raise _DatasetSourceError(f"{dataset} source is empty")

missing = columns - set(df.columns)

if missing:

raise _DatasetSourceError(f"{dataset} source is missing columns: {sorted(missing)}")

def _identity_dataset(df: pd.DataFrame) -> pd.DataFrame:

"""Return an already-normalized dataset unchanged."""

return df

def _require_complete(df: pd.DataFrame, dataset: str, columns: set) -> None:

"""Reject missing values in columns whose public contract is complete."""

missing = df[list(columns)].isna().sum()

missing = missing[missing > 0]

if not missing.empty:

raise _DatasetSourceError(

f"{dataset} source has missing required values: {missing.to_dict()}"

)

def _require_finite(df: pd.DataFrame, dataset: str, columns: set) -> None:

"""Reject non-numeric or non-finite values in numeric contract columns."""

try:

values = df[list(columns)].to_numpy(dtype=float)

except (TypeError, ValueError) as e:

raise _DatasetSourceError(f"{dataset} source has non-numeric values") from e

if not np.isfinite(values).all():

raise _DatasetSourceError(f"{dataset} source has non-finite values in {sorted(columns)}")

def _validate_panel_keys(df: pd.DataFrame, dataset: str, unit: str) -> None:

"""Validate the common unit-time and cohort invariants for panel datasets."""

if df.duplicated([unit, "year"]).any():

raise _DatasetSourceError(f"{dataset} source has duplicate {unit}-year rows")

cohort_counts = df.groupby(unit)["first_treat"].nunique(dropna=False)

if not (cohort_counts == 1).all():

raise _DatasetSourceError(f"{dataset} first_treat is not constant within {unit}")

if not (df["cohort"] == df["first_treat"]).all():

raise _DatasetSourceError(f"{dataset} cohort does not match first_treat")

expected_treated = ((df["first_treat"] > 0) & (df["year"] >= df["first_treat"])).astype(int)

if not (df["treated"] == expected_treated).all():

raise _DatasetSourceError(f"{dataset} treated indicator is inconsistent with first_treat")

def _prepare_card_krueger(df: pd.DataFrame) -> pd.DataFrame:

"""Normalize a Card-Krueger source frame to the public loader schema."""

raw_columns = {

"sheet",

"state",

"chain",

"empft",

"emppt",

"nmgrs",

"wage_st",

"empft2",

"emppt2",

"nmgrs2",

"wage_st2",

}

if raw_columns <= set(df.columns):

store_id = df["sheet"].copy()

_require_complete(df, "card_krueger", {"sheet", "state", "chain"})

_require_finite(df, "card_krueger", {"sheet", "state", "chain"})

if not set(df["state"].unique()) <= {0, 1}:

raise _DatasetSourceError("card_krueger source has unknown state codes")

if not set(df["chain"].unique()) <= {1, 2, 3, 4}:

raise _DatasetSourceError("card_krueger source has unknown chain codes")

for column in (

"empft",

"emppt",

"nmgrs",

"wage_st",

"empft2",

"emppt2",

"nmgrs2",

"wage_st2",

):

converted = pd.to_numeric(df[column], errors="coerce")

if converted.notna().sum() != df[column].notna().sum():

raise _DatasetSourceError(f"card_krueger source has non-numeric values in {column}")

df[column] = converted

duplicate_407 = store_id == 407

if (

duplicate_407.sum() != 2

or set(df.loc[duplicate_407, "state"]) != {0, 1}

or (store_id == 408).any()

):

raise _DatasetSourceError(

"card_krueger source does not match the documented duplicate-407 convention"

)

store_id.loc[duplicate_407 & (df["state"] == 1)] = 408

emp_pre = df["empft"] + df["nmgrs"] + 0.5 * df["emppt"]

emp_post = df["empft2"] + df["nmgrs2"] + 0.5 * df["emppt2"]

return pd.DataFrame(

{

"store_id": store_id.astype(int),

"state": np.where(df["state"] == 1, "NJ", "PA"),

"chain": df["chain"].map({1: "bk", 2: "kfc", 3: "roys", 4: "wendys"}),

"emp_pre": emp_pre,

"emp_post": emp_post,

"wage_pre": df["wage_st"],

"wage_post": df["wage_st2"],

"treated": (df["state"] == 1).astype(int),

"emp_change": emp_post - emp_pre,

}

)

df = df.rename(columns={"sheet": "store_id"}).copy()

if "state" not in df.columns and "nj" in df.columns:

df["state"] = np.where(df["nj"] == 1, "NJ", "PA")

if "treated" not in df.columns and "state" in df.columns:

df["treated"] = (df["state"] == "NJ").astype(int)

if "emp_change" not in df.columns and {"emp_post", "emp_pre"} <= set(df.columns):

df["emp_change"] = df["emp_post"] - df["emp_pre"]

return df

def _validate_card_krueger(df: pd.DataFrame) -> None:

"""Validate the documented Card-Krueger wide-data contract."""

_require_columns(

df,

"card_krueger",

{

"store_id",

"state",

"chain",

"emp_pre",

"emp_post",

"wage_pre",

"wage_post",

"treated",

"emp_change",

},

)

_require_complete(df, "card_krueger", {"store_id", "state", "chain", "treated"})

_require_finite(df, "card_krueger", {"store_id", "treated"})

if df["store_id"].duplicated().any():

raise _DatasetSourceError("card_krueger source has duplicate store_id values")

if set(df["state"].dropna().unique()) != {"NJ", "PA"}:

raise _DatasetSourceError("card_krueger source must contain both NJ and PA only")

if set(df["chain"].dropna().unique()) != {"bk", "kfc", "roys", "wendys"}:

raise _DatasetSourceError("card_krueger source has unexpected restaurant chains")

for column in ("emp_pre", "emp_post", "wage_pre", "wage_post"):

values = pd.to_numeric(df[column], errors="coerce")

if (

values.notna().sum() != df[column].notna().sum()

or not np.isfinite(values.dropna()).all()

or (values.dropna() < 0).any()

):

raise _DatasetSourceError(

f"card_krueger source has invalid non-negative values in {column}"

)

emp_change = pd.to_numeric(df["emp_change"], errors="coerce")

if (

emp_change.notna().sum() != df["emp_change"].notna().sum()

or not np.isfinite(emp_change.dropna()).all()

):

raise _DatasetSourceError("card_krueger source has invalid emp_change values")

expected_treated = (df["state"] == "NJ").astype(int)

if not (df["treated"] == expected_treated).all():

raise _DatasetSourceError("card_krueger treated indicator is inconsistent with state")

expected_change = df["emp_post"] - df["emp_pre"]

if not np.allclose(df["emp_change"], expected_change, equal_nan=True):

raise _DatasetSourceError(

"card_krueger emp_change is inconsistent with emp_pre and emp_post"

)

def _validate_card_krueger_source(df: pd.DataFrame) -> None:

"""Validate source-specific Card-Krueger counts and categories."""

_validate_card_krueger(df)

if len(df) != 410:

raise _DatasetSourceError("card_krueger source must contain 410 stores")

if df.groupby("state").size().to_dict() != {"NJ": 331, "PA": 79}:

raise _DatasetSourceError("card_krueger source has unexpected state counts")

expected_missing = {

"emp_pre": 12,

"emp_post": 14,

"wage_pre": 20,

"wage_post": 21,

"emp_change": 26,

}

if df[list(expected_missing)].isna().sum().to_dict() != expected_missing:

raise _DatasetSourceError("card_krueger source has unexpected missing-value counts")

def _prepare_castle_doctrine(df: pd.DataFrame) -> pd.DataFrame:

"""Normalize a Castle Doctrine source frame to the public loader schema."""

df = df.copy()

if "sid" in df.columns:

state_codes = df["sid"].map(_CASTLE_STATE_BY_SID)

if state_codes.notna().all() and not (state_codes == "_").any():

df["state"] = state_codes

if "first_treat" not in df.columns and "effyear" in df.columns:

try:

df["first_treat"] = df["effyear"].fillna(0).astype(int)

except (TypeError, ValueError) as e:

raise _DatasetSourceError("castle_doctrine source has invalid effyear values") from e

if "cohort" not in df.columns and "first_treat" in df.columns:

df["cohort"] = df["first_treat"]

if {"first_treat", "year"} <= set(df.columns):

df["treated"] = ((df["first_treat"] > 0) & (df["year"] >= df["first_treat"])).astype(int)

if "treatment_exposure" not in df.columns and "cdl" in df.columns:

df["treatment_exposure"] = df["cdl"]

if "homicide_rate" not in df.columns and "homicide" in df.columns:

df["homicide_rate"] = df["homicide"]

if {

"state",

"year",

"first_treat",

"homicide_rate",

"population",

"income",

"treated",

"treatment_exposure",

"cohort",

} <= set(df.columns):

return df[

[

"state",

"year",

"first_treat",

"homicide_rate",

"population",

"income",

"treated",

"treatment_exposure",

"cohort",

]

].copy()

return df

def _validate_castle_doctrine(df: pd.DataFrame) -> None:

"""Validate the documented Castle Doctrine panel contract."""

_require_columns(

df,

"castle_doctrine",

{

"state",

"year",

"first_treat",

"homicide_rate",

"population",

"income",

"treated",

"treatment_exposure",

"cohort",

},

)

_require_complete(

df,

"castle_doctrine",

{

"state",

"year",

"first_treat",

"homicide_rate",

"population",

"income",

"treated",

"cohort",

},

)

_require_finite(

df,

"castle_doctrine",

{

"year",

"first_treat",

"homicide_rate",

"population",

"income",

"treated",

"treatment_exposure",

"cohort",

},

)

if (

(df["homicide_rate"] < 0).any()

or (df["population"] <= 0).any()

or (df["income"] <= 0).any()

):

raise _DatasetSourceError("castle_doctrine source has invalid outcome or covariate values")

if not df["state"].astype(str).str.fullmatch(r"[A-Z]{2}").all():

raise _DatasetSourceError("castle_doctrine source has invalid state abbreviations")

if not df["treatment_exposure"].between(0, 1).all():

raise _DatasetSourceError("castle_doctrine treatment_exposure must be between 0 and 1")

_validate_panel_keys(df, "castle_doctrine", "state")

def _validate_castle_doctrine_source(df: pd.DataFrame) -> None:

"""Validate source-specific Castle Doctrine panel dimensions."""

_validate_castle_doctrine(df)

if len(df) != 550 or df["state"].nunique() != 50:

raise _DatasetSourceError("castle_doctrine source must contain 50 states and 550 rows")

if set(df["year"].unique()) != set(range(2000, 2011)):

raise _DatasetSourceError("castle_doctrine source has unexpected years")

if set(df["first_treat"].unique()) != {0, 2005, 2006, 2007, 2008, 2009}:

raise _DatasetSourceError("castle_doctrine source has unexpected treatment cohorts")

def _validate_divorce_laws(df: pd.DataFrame) -> None:

"""Validate the documented divorce-laws panel contract."""

_require_columns(

df,

"divorce_laws",

{

"state",

"year",

"first_treat",

"divorce_rate",

"female_lfp",

"suicide_rate",

"treated",

"cohort",

},

)

_require_complete(

df,

"divorce_laws",

{

"state",

"year",

"first_treat",

"divorce_rate",

"female_lfp",

"suicide_rate",

"treated",

"cohort",

},

)

_require_finite(

df,

"divorce_laws",

{

"year",

"first_treat",

"divorce_rate",

"female_lfp",

"suicide_rate",

"treated",

"cohort",

},

)

if (df["divorce_rate"] < 0).any() or (df["suicide_rate"] < 0).any():

raise _DatasetSourceError("divorce_laws source has negative outcome values")

if not df["female_lfp"].between(0, 1).all():

raise _DatasetSourceError("divorce_laws female_lfp must be between 0 and 1")

_validate_panel_keys(df, "divorce_laws", "state")

def _prepare_mpdta(df: pd.DataFrame) -> pd.DataFrame:

"""Normalize an mpdta source frame to the public loader schema."""

if "first.treat" in df.columns:

df = df.rename(columns={"first.treat": "first_treat"})

if "cohort" not in df.columns and "first_treat" in df.columns:

df["cohort"] = df["first_treat"]

columns = ["countyreal", "year", "lpop", "lemp", "first_treat", "treat", "cohort"]

if set(columns) <= set(df.columns):

return df[columns].copy()

return df

def _validate_mpdta(df: pd.DataFrame) -> None:

"""Validate the canonical R did::mpdta panel structure."""

_require_columns(

df,

"mpdta",

{"countyreal", "year", "lpop", "lemp", "first_treat", "treat", "cohort"},

)

_require_complete(

df,

"mpdta",

{"countyreal", "year", "lpop", "lemp", "first_treat", "treat", "cohort"},

)

_require_finite(

df,

"mpdta",

{"countyreal", "year", "lpop", "lemp", "first_treat", "treat", "cohort"},

)

if df.duplicated(["countyreal", "year"]).any():

raise _DatasetSourceError("mpdta source has duplicate county-year rows")

if len(df) != 2500 or df["countyreal"].nunique() != 500:

raise _DatasetSourceError("mpdta source must contain 500 counties and 2500 rows")

if set(df["year"].unique()) != {2003, 2004, 2005, 2006, 2007}:

raise _DatasetSourceError("mpdta source has unexpected years")

if set(df["first_treat"].unique()) != {0, 2004, 2006, 2007}:

raise _DatasetSourceError("mpdta source has unexpected treatment cohorts")

cohort_counts = df.groupby("countyreal")["first_treat"].nunique(dropna=False)

if not (cohort_counts == 1).all():

raise _DatasetSourceError("mpdta first_treat is not constant within county")

if not (df["cohort"] == df["first_treat"]).all():

raise _DatasetSourceError("mpdta cohort does not match first_treat")

if not (df["treat"] == (df["first_treat"] > 0).astype(int)).all():

raise _DatasetSourceError("mpdta treat indicator is inconsistent with first_treat")

def clear_cache() -> None:

"""Clear the local dataset cache.

Also removes any ``.<name>.<ext>.<suffix>`` scratch files left behind by an

atomic cache write that was interrupted between creating the temporary file

and replacing the cache entry (a hard kill, for instance). Those are hidden

and do not match the plain ``*.csv`` / ``*.dta`` patterns, so without this

they would accumulate and survive the documented remedy.

"""

if _CACHE_DIR.exists():

for pattern in ("*.csv", "*.dta", ".*.csv.*", ".*.dta.*"):

for f in _CACHE_DIR.glob(pattern):

if f.is_file():

f.unlink()

print(f"Cleared cache at {_CACHE_DIR}")

def load_card_krueger(force_download: bool = False) -> pd.DataFrame:

"""

Load the Card & Krueger (1994) minimum wage dataset.

This classic dataset examines the effect of New Jersey's 1992 minimum wage

increase on employment in fast-food restaurants, using Pennsylvania as

a control group.

The study is a canonical example of the Difference-in-Differences method.

Parameters

----------

force_download : bool, default=False

If True, re-download the dataset even if cached.

Returns

-------

pd.DataFrame

Dataset with columns:

- store_id : int - Unique store identifier

- state : str - 'NJ' (New Jersey, treated) or 'PA' (Pennsylvania, control)

- chain : str - Fast food chain ('bk', 'kfc', 'roys', 'wendys')

- emp_pre : float - Full-time equivalent employment before (Feb 1992)

- emp_post : float - Full-time equivalent employment after (Nov 1992)

- wage_pre : float - Starting wage before

- wage_post : float - Starting wage after

- treated : int - 1 if NJ, 0 if PA

- emp_change : float - Change in employment (emp_post - emp_pre)

Notes

-----

The minimum wage in New Jersey increased from $4.25 to $5.05 on April 1, 1992.

Pennsylvania's minimum wage remained at $4.25.

Original finding: No significant negative effect of minimum wage increase

on employment (ATT ≈ +2.8 FTE employees).

The canonical survey is incomplete: 12 stores lack ``emp_pre``, 14 lack

``emp_post``, and ``wage_pre``/``wage_post`` are missing for 20 and 21

stores. Drop the missing outcome rows before fitting, as the example below

does; estimators reject missing outcomes rather than dropping them silently.

The synthetic fallback frame is complete, so code that skips this step will

work offline and fail once the canonical source is reachable.

The canonical data are checksum-verified and returned with

``df.attrs["source"] == "card_krueger_public_data"``. A download failure

falls back to a checksum-valid cache entry when one exists, returning that

canonical data; a pin mismatch additionally warns that the upstream file has

changed. Only when neither a verified cache entry nor a verified fresh

download is available - or when the source fails to parse or validate - does

the loader emit one ``UserWarning`` containing ``SYNTHETIC``

and return ``df.attrs["source"] == "synthetic_fallback"``.

References

----------

Card, D., & Krueger, A. B. (1994). Minimum Wages and Employment: A Case Study

of the Fast-Food Industry in New Jersey and Pennsylvania. *American Economic

Review*, 84(4), 772-793.

Examples

--------

>>> from diff_diff.datasets import load_card_krueger

>>> from diff_diff import DifferenceInDifferences

>>>

>>> # Load and prepare data

>>> ck = load_card_krueger()

>>> ck_long = ck.melt(

... id_vars=['store_id', 'state', 'treated'],

... value_vars=['emp_pre', 'emp_post'],

... var_name='period', value_name='employment'

... )

>>> ck_long['post'] = (ck_long['period'] == 'emp_post').astype(int)

>>>

>>> # 26 store-waves have no employment reading in the source survey

>>> ck_long = ck_long.dropna(subset=['employment'])

>>>

>>> # Estimate DiD

>>> did = DifferenceInDifferences()

>>> results = did.fit(ck_long, outcome='employment', treatment='treated', post='post')

"""

return _load_verified_dataset(

cache_name="card_krueger",

source="card_krueger_public_data",

force_download=force_download,

load_source=_load_card_krueger_source,

prepare=_prepare_card_krueger,

validate_source=_validate_card_krueger_source,

validate_fallback=_validate_card_krueger,

fallback=_construct_card_krueger_data,

)

def _construct_card_krueger_data() -> pd.DataFrame:

"""

Construct Card-Krueger dataset from summary statistics.

This is a fallback when the online source is unavailable.

Uses aggregated data that preserves the key DiD estimates.

"""

# Representative sample based on published summary statistics

np.random.seed(1994) # Card-Krueger publication year, for reproducibility

stores = []

store_id = 1

# New Jersey stores (treated) - summary stats from paper

# Mean emp before: 20.44, after: 21.03

# Mean wage before: 4.61, after: 5.08

for chain in ["bk", "kfc", "roys", "wendys"]:

n_stores = {"bk": 85, "kfc": 62, "roys": 48, "wendys": 36}[chain]

for _ in range(n_stores):

emp_pre = np.random.normal(20.44, 8.5)

emp_post = emp_pre + np.random.normal(0.59, 7.0) # Change ≈ 0.59

emp_pre = max(0, emp_pre)

emp_post = max(0, emp_post)

stores.append(

{

"store_id": store_id,

"state": "NJ",

"chain": chain,

"emp_pre": round(emp_pre, 1),

"emp_post": round(emp_post, 1),

"wage_pre": round(np.random.normal(4.61, 0.35), 2),

"wage_post": round(np.random.normal(5.08, 0.12), 2),

}

)

store_id += 1

# Pennsylvania stores (control) - summary stats from paper

# Mean emp before: 23.33, after: 21.17

# Mean wage before: 4.63, after: 4.62

for chain in ["bk", "kfc", "roys", "wendys"]:

n_stores = {"bk": 30, "kfc": 20, "roys": 14, "wendys": 15}[chain]

for _ in range(n_stores):

emp_pre = np.random.normal(23.33, 8.2)

emp_post = emp_pre + np.random.normal(-2.16, 7.0) # Change ≈ -2.16

emp_pre = max(0, emp_pre)

emp_post = max(0, emp_post)

stores.append(

{

"store_id": store_id,

"state": "PA",

"chain": chain,

"emp_pre": round(emp_pre, 1),

"emp_post": round(emp_post, 1),

"wage_pre": round(np.random.normal(4.63, 0.35), 2),

"wage_post": round(np.random.normal(4.62, 0.35), 2),

}

)

store_id += 1

df = pd.DataFrame(stores)

df["treated"] = (df["state"] == "NJ").astype(int)

df["emp_change"] = df["emp_post"] - df["emp_pre"]

return df

def load_castle_doctrine(force_download: bool = False) -> pd.DataFrame:

"""

Load Castle Doctrine / Stand Your Ground laws dataset.

This dataset tracks the staggered adoption of Castle Doctrine (Stand Your

Ground) laws across U.S. states, which expanded self-defense rights.

It's commonly used to demonstrate heterogeneous treatment timing methods

like Callaway-Sant'Anna or Sun-Abraham.

Parameters

----------

force_download : bool, default=False

If True, re-download the dataset even if cached.

Returns

-------

pd.DataFrame

Panel dataset with columns:

- state : str - State abbreviation

- year : int - Year (2000-2010)

- first_treat : int - Year of law adoption (0 = never adopted)

- homicide_rate : float - Homicides per 100,000 population

- population : int - State population

- income : float - State median income

- treated : int - 1 if law in effect, 0 otherwise

- treatment_exposure : float - Fraction of the year the law was in effect

- cohort : int - Alias for first_treat

Notes

-----

Castle Doctrine laws remove the duty to retreat before using deadly force

in self-defense. States adopted these laws at different times between

2005 and 2009, creating a staggered treatment design.

The canonical data are checksum-verified and returned with

``df.attrs["source"] == "cheng_hoekstra_castle_data"``. A download failure

falls back to a checksum-valid cache entry when one exists, returning that

canonical data; a pin mismatch additionally warns that the upstream file has

changed. Only when neither a verified cache entry nor a verified fresh

download is available - or when the source fails to parse or validate - does

the loader emit one ``UserWarning`` containing ``SYNTHETIC``

and mark the returned frame as ``"synthetic_fallback"``.

``treatment_exposure`` is fractional only on canonical frames; synthetic

fallback frames set it to a binary 0/1 copy of ``treated`` and therefore

carry no partial-year information.

Replicating Cheng-Hoekstra (2013) requires the paper's regressor and outcome:

regress ``log(homicide_rate)`` on ``treatment_exposure`` (their ``CDL_it``,

the proportion of the year the law was in effect), not the binary ``treated``.

See "Castle Doctrine treatment coding" in ``docs/methodology/REGISTRY.md``.

References

----------

Cheng, C., & Hoekstra, M. (2013). Does Strengthening Self-Defense Law Deter

Crime or Escalate Violence? Evidence from Expansions to Castle Doctrine.

*Journal of Human Resources*, 48(3), 821-854.

Examples

--------

>>> from diff_diff.datasets import load_castle_doctrine

>>> from diff_diff import CallawaySantAnna

>>>

>>> castle = load_castle_doctrine()

>>> cs = CallawaySantAnna(control_group="never_treated")

>>> results = cs.fit(

... castle,

... outcome="homicide_rate",

... unit="state",

... time="year",

... first_treat="first_treat"

... )

"""