diff-diff/diff_diff/power.py at main · UtilityTools/diff-diff · GitHub

History

3196 lines (2858 loc) · 110 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

658

659

660

661

662

663

664

665

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

682

683

684

685

686

687

688

689

690

691

692

693

694

695

696

697

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

799

800

801

802

803

804

805

806

807

808

809

810

811

812

813

814

815

816

817

818

819

820

821

822

823

824

825

826

827

828

829

830

831

832

833

834

835

836

837

838

839

840

841

842

843

844

845

846

847

848

849

850

851

852

853

854

855

856

857

858

859

860

861

862

863

864

865

866

867

868

869

870

871

872

873

874

875

876

877

878

879

880

881

882

883

884

885

886

887

888

889

890

891

892

893

894

895

896

897

898

899

900

901

902

903

904

905

906

907

908

909

910

911

912

913

914

915

916

917

918

919

920

921

922

923

924

925

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

975

976

977

978

979

980

981

982

983

984

985

986

987

988

989

990

991

992

993

994

995

996

997

998

999

1000

"""

Power analysis tools for difference-in-differences study design.

This module provides power calculations and simulation-based power analysis

for DiD study design, helping practitioners answer questions like:

- "How many units do I need to detect an effect of size X?"

- "What is the minimum detectable effect given my sample size?"

- "What power do I have to detect a given effect?"

References

----------

Bloom, H. S. (1995). "Minimum Detectable Effects: A Simple Way to Report the

Statistical Power of Experimental Designs." Evaluation Review, 19(5), 547-556.

Burlig, F., Preonas, L., & Woerman, M. (2020). "Panel Data and Experimental Design."

Journal of Development Economics, 144, 102458.

Djimeu, E. W., & Houndolo, D.-G. (2016). "Power Calculation for Causal Inference

in Social Science: Sample Size and Minimum Detectable Effect Determination."

Journal of Development Effectiveness, 8(4), 508-527.

"""

import warnings

from dataclasses import dataclass, field

from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np

import pandas as pd

from scipy import stats

# Maximum sample size returned when effect is too small to detect

# (e.g., zero effect or extremely small relative to noise)

MAX_SAMPLE_SIZE = 2**31 - 1

# ---------------------------------------------------------------------------

# Estimator registry — maps estimator class names to DGP/fit/extract profiles

# ---------------------------------------------------------------------------

@dataclass

class _EstimatorProfile:

"""Internal profile describing how to run power simulations for an estimator."""

default_dgp: Callable

dgp_kwargs_builder: Callable

fit_kwargs_builder: Callable

result_extractor: Callable

min_n: int = 20

# ---------------------------------------------------------------------------

# SurveyPowerConfig — carries DGP survey params for simulation power

# ---------------------------------------------------------------------------

@dataclass

class SurveyPowerConfig:

"""Configuration for survey-aware power simulations.

When passed to :func:`simulate_power`, :func:`simulate_mde`, or

:func:`simulate_sample_size`, the simulation loop generates data with

:func:`~diff_diff.prep.generate_survey_did_data` and automatically

injects a ``SurveyDesign`` into the estimator's ``fit()`` call.

Parameters

----------

n_strata : int, default=5

Number of geographic strata.

psu_per_stratum : int, default=8

Number of primary sampling units (PSUs) per stratum. Must be >= 2

for Taylor Series Linearization variance estimation.

fpc_per_stratum : float, default=200.0

Finite population correction (total PSUs per stratum).

weight_variation : str, default="moderate"

Sampling weight dispersion: ``"none"`` (all equal), ``"moderate"``

(range ~1-2), ``"high"`` (range ~1-4).

psu_re_sd : float, default=2.0

Standard deviation of PSU random effects. Controls intra-cluster

correlation and drives DEFF > 1.

psu_period_factor : float, default=0.5

Multiplier for PSU-period interaction shocks.

icc : float, optional

Target intra-class correlation (0 < icc < 1). Overrides

``psu_re_sd`` via variance decomposition.

weight_cv : float, optional

Target coefficient of variation for weights. Overrides

``weight_variation``.

informative_sampling : bool, default=False

If True, weights correlate with Y(0).

heterogeneous_te_by_strata : bool, default=False

If True, treatment effect varies by stratum.

include_replicate_weights : bool, default=False

If True, add JK1 delete-one-PSU replicate weight columns.

survey_design : SurveyDesign, optional

Override the auto-built SurveyDesign. When None, a default

``SurveyDesign(weights="weight", strata="stratum", psu="psu",

fpc="fpc")`` is used, matching ``generate_survey_did_data`` output.

Examples

--------

>>> from diff_diff import CallawaySantAnna, simulate_power, SurveyPowerConfig

>>> config = SurveyPowerConfig(n_strata=5, psu_per_stratum=8, icc=0.05)

>>> results = simulate_power(

... CallawaySantAnna(),

... n_units=200,

... treatment_effect=2.0,

... survey_config=config,

... n_simulations=100,

... seed=42,

... )

"""

n_strata: int = 5

psu_per_stratum: int = 8

fpc_per_stratum: float = 200.0

weight_variation: str = "moderate"

psu_re_sd: float = 2.0

psu_period_factor: float = 0.5

icc: Optional[float] = None

weight_cv: Optional[float] = None

informative_sampling: bool = False

heterogeneous_te_by_strata: bool = False

include_replicate_weights: bool = False

survey_design: Optional[Any] = None

def __post_init__(self) -> None:

if self.n_strata < 1:

raise ValueError(f"n_strata must be >= 1, got {self.n_strata}")

if self.psu_per_stratum < 2:

raise ValueError(

f"psu_per_stratum must be >= 2 for TSL variance estimation, "

f"got {self.psu_per_stratum}"

)

if self.weight_variation not in ("none", "moderate", "high"):

raise ValueError(

f"weight_variation must be 'none', 'moderate', or 'high', "

f"got '{self.weight_variation}'"

)

if not np.isfinite(self.psu_re_sd) or self.psu_re_sd < 0:

raise ValueError(f"psu_re_sd must be finite and >= 0, got {self.psu_re_sd}")

if not np.isfinite(self.fpc_per_stratum):

raise ValueError(f"fpc_per_stratum must be finite, got {self.fpc_per_stratum}")

if self.icc is not None and not (0 < self.icc < 1):

raise ValueError(f"icc must be between 0 and 1 (exclusive), got {self.icc}")

if self.icc is not None and self.psu_re_sd != 2.0:

raise ValueError(

"Cannot specify both icc and a non-default psu_re_sd. "

"icc overrides psu_re_sd via the ICC formula."

)

if self.weight_cv is not None:

if not np.isfinite(self.weight_cv) or self.weight_cv <= 0:

raise ValueError(f"weight_cv must be finite and > 0, got {self.weight_cv}")

if self.weight_variation != "moderate":

raise ValueError(

"Cannot specify both weight_cv and a non-default "

"weight_variation. weight_cv overrides weight_variation."

)

if not np.isfinite(self.psu_period_factor) or self.psu_period_factor < 0:

raise ValueError(

f"psu_period_factor must be finite and >= 0, got {self.psu_period_factor}"

)

if self.fpc_per_stratum < self.psu_per_stratum:

raise ValueError(

f"fpc_per_stratum ({self.fpc_per_stratum}) must be >= "

f"psu_per_stratum ({self.psu_per_stratum})"

)

def _build_survey_design(self) -> Any:

"""Return a SurveyDesign for this config.

Reflects the live ``self.survey_design`` value every call (no

caching). Finding #28 (axis J, silent-failures audit): the

previous ``_cached_survey_design`` was populated on first call

and never invalidated on mutation, so ``config.survey_design =

other_design`` silently kept returning the original. Since the

default ``SurveyDesign(...)`` construction is microseconds and

user-provided designs are just reference copies, there's no cache

cost worth keeping.

"""

if self.survey_design is not None:

return self.survey_design

from diff_diff.survey import SurveyDesign

return SurveyDesign(

weights="weight", strata="stratum", psu="psu", fpc="fpc"

)

@property

def min_viable_n(self) -> int:

"""Minimum n_units for a viable survey design (>= 2 units per PSU)."""

return self.n_strata * self.psu_per_stratum * 2

# -- DGP kwargs adapters -----------------------------------------------------

def _basic_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

) -> Dict[str, Any]:

return dict(

n_units=n_units,

n_periods=n_periods,

treatment_effect=treatment_effect,

treatment_fraction=treatment_fraction,

treatment_period=treatment_period,

noise_sd=sigma,

)

def _staggered_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

) -> Dict[str, Any]:

return dict(

n_units=n_units,

n_periods=n_periods,

treatment_effect=treatment_effect,

never_treated_frac=1 - treatment_fraction,

cohort_periods=[treatment_period],

dynamic_effects=False,

noise_sd=sigma,

)

def _factor_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

) -> Dict[str, Any]:

n_pre = treatment_period

n_post = n_periods - treatment_period

return dict(

n_units=n_units,

n_pre=n_pre,

n_post=n_post,

n_treated=max(1, int(n_units * treatment_fraction)),

treatment_effect=treatment_effect,

noise_sd=sigma,

)

def _ddd_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

) -> Dict[str, Any]:

return dict(

n_per_cell=max(2, n_units // 8),

treatment_effect=treatment_effect,

noise_sd=sigma,

)

# -- Fit kwargs builders ------------------------------------------------------

def _basic_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(outcome="outcome", treatment="treated", time="post")

def _twfe_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(outcome="outcome", treatment="treated", time="post", unit="unit")

def _multiperiod_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(

outcome="outcome",

treatment="treated",

time="period",

post_periods=list(range(treatment_period, n_periods)),

)

def _staggered_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat")

def _ddd_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(outcome="outcome", group="group", partition="partition", time="time")

def _trop_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

return dict(outcome="outcome", treatment="treated", unit="unit", time="period")

def _sdid_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

periods = sorted(data["period"].unique())

post_periods = [p for p in periods if p >= treatment_period]

return dict(

outcome="outcome",

treatment="treat",

unit="unit",

time="period",

post_periods=post_periods,

)

# -- Survey-aware DGP kwargs adapter ------------------------------------------

def _survey_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

survey_config: SurveyPowerConfig,

) -> Dict[str, Any]:

"""Build kwargs for generate_survey_did_data from simulate_power params."""

return dict(

n_units=n_units,

n_periods=n_periods,

treatment_effect=treatment_effect,

never_treated_frac=1 - treatment_fraction,

# 0-indexed treatment_period → 1-indexed cohort_periods

cohort_periods=[treatment_period + 1],

noise_sd=sigma,

dynamic_effects=False,

n_strata=survey_config.n_strata,

psu_per_stratum=survey_config.psu_per_stratum,

fpc_per_stratum=survey_config.fpc_per_stratum,

weight_variation=survey_config.weight_variation,

psu_re_sd=survey_config.psu_re_sd,

psu_period_factor=survey_config.psu_period_factor,

icc=survey_config.icc,

weight_cv=survey_config.weight_cv,

informative_sampling=survey_config.informative_sampling,

heterogeneous_te_by_strata=survey_config.heterogeneous_te_by_strata,

include_replicate_weights=survey_config.include_replicate_weights,

return_true_population_att=True,

)

# -- Survey-aware fit kwargs builders -----------------------------------------

def _survey_basic_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

survey_config: SurveyPowerConfig,

) -> Dict[str, Any]:

"""Fit kwargs for DifferenceInDifferences with survey design.

Uses ``ever_treated`` (time-invariant group indicator) rather than the

survey DGP's ``treated`` column (which is post-only: 1{g>0, t>=g}).

DifferenceInDifferences internally constructs ``treatment * time``,

so passing the post-only flag would make that interaction rank-deficient.

"""

return dict(

outcome="outcome",

treatment="ever_treated",

time="post",

survey_design=survey_config._build_survey_design(),

)

def _survey_twfe_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

survey_config: SurveyPowerConfig,

) -> Dict[str, Any]:

"""Fit kwargs for TwoWayFixedEffects with survey design."""

return dict(

outcome="outcome",

treatment="ever_treated",

time="post",

unit="unit",

survey_design=survey_config._build_survey_design(),

)

def _survey_multiperiod_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

survey_config: SurveyPowerConfig,

) -> Dict[str, Any]:

"""Fit kwargs for MultiPeriodDiD with survey design (1-indexed periods)."""

return dict(

outcome="outcome",

treatment="ever_treated",

unit="unit",

time="period",

# 1-indexed: post periods run from treatment_period+1 to n_periods

post_periods=list(range(treatment_period + 1, n_periods + 1)),

survey_design=survey_config._build_survey_design(),

)

def _survey_staggered_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

survey_config: SurveyPowerConfig,

) -> Dict[str, Any]:

"""Fit kwargs for staggered estimators (CS, SA, etc.) with survey design."""

return dict(

outcome="outcome",

unit="unit",

time="period",

first_treat="first_treat",

survey_design=survey_config._build_survey_design(),

)

# -- Result extractors --------------------------------------------------------

def _extract_simple(result: Any) -> Tuple[float, float, float, Tuple[float, float]]:

return (result.att, result.se, result.p_value, result.conf_int)

def _extract_multiperiod(

result: Any,

) -> Tuple[float, float, float, Tuple[float, float]]:

return (result.avg_att, result.avg_se, result.avg_p_value, result.avg_conf_int)

def _extract_staggered(

result: Any,

) -> Tuple[float, float, float, Tuple[float, float]]:

_nan = float("nan")

_nan_ci = (_nan, _nan)

def _first(r: Any, *attrs: str, default: Any = _nan) -> Any:

for a in attrs:

v = getattr(r, a, None)

if v is not None:

return v

return default

return (

result.overall_att,

_first(result, "overall_se", "overall_att_se"),

_first(result, "overall_p_value", "overall_att_p_value"),

_first(result, "overall_conf_int", "overall_att_ci", default=_nan_ci),

)

# Keys derived from simulate_power() public params — overriding these

# via data_generator_kwargs would desync the DGP from the result object.

_PROTECTED_DGP_KEYS = frozenset(

{

"treatment_effect", # → true_effect in results / MDE search variable

"noise_sd", # → sigma param

"n_units", # → sample-size search variable

"n_periods", # → n_periods param

"treatment_fraction", # → treatment_fraction param

"treatment_period", # → treatment_period param

"n_pre", # → derived from treatment_period in factor-model DGPs

"n_post", # → derived from n_periods - treatment_period in factor-model DGPs

}

)

# Keys managed by SurveyPowerConfig — block in data_generator_kwargs when

# survey_config is active to prevent silent conflicts.

_SURVEY_CONFIG_KEYS = frozenset(

{

"n_strata",

"psu_per_stratum",

"fpc_per_stratum",

"weight_variation",

"psu_re_sd",

"psu_period_factor",

"icc",

"weight_cv",

"informative_sampling",

"heterogeneous_te_by_strata",

"include_replicate_weights",

"return_true_population_att",

"dynamic_effects",

"cohort_periods",

"never_treated_frac",

}

)

# -- Staggered DGP compatibility check ----------------------------------------

_STAGGERED_ESTIMATORS = frozenset(

{

"CallawaySantAnna",

"SunAbraham",

"ImputationDiD",

"TwoStageDiD",

"StackedDiD",

"EfficientDiD",

}

)

# Estimators that need a derived `post` column when using survey DGP

# (survey DGP produces `period`/`first_treat` but not `post`).

_SURVEY_POST_ESTIMATORS = frozenset({"DifferenceInDifferences", "TwoWayFixedEffects"})

# Survey fit kwargs builder lookup — maps estimator name to builder function.

_SURVEY_FIT_BUILDERS: Dict[str, Callable] = {

"DifferenceInDifferences": _survey_basic_fit_kwargs,

"TwoWayFixedEffects": _survey_twfe_fit_kwargs,

"MultiPeriodDiD": _survey_multiperiod_fit_kwargs,

**{name: _survey_staggered_fit_kwargs for name in _STAGGERED_ESTIMATORS},

}

# Unsupported: factor-model and triple-diff estimators (survey DGP produces

# staggered cohort data, not factor-model or 2x2x2 data).

_SURVEY_UNSUPPORTED = frozenset({"TROP", "SyntheticDiD", "TripleDifference"})

def _check_staggered_dgp_compat(

estimator: Any,

data_generator_kwargs: Optional[Dict[str, Any]],

) -> None:

"""Warn if a staggered estimator's settings don't match the default DGP."""

name = type(estimator).__name__

if name not in _STAGGERED_ESTIMATORS:

return

dgp_overrides = data_generator_kwargs or {}

cohort_periods = dgp_overrides.get("cohort_periods")

has_multi_cohort = cohort_periods is not None and len(set(cohort_periods)) >= 2

issues: List[str] = []

# Check control_group="not_yet_treated" (CS, SA)

cg = getattr(estimator, "control_group", "never_treated")

if cg == "not_yet_treated" and not has_multi_cohort:

issues.append(

f' - {name} has control_group="not_yet_treated" but the default '

f"DGP generates a single treatment cohort with never-treated "

f"controls. Power may not reflect the intended not-yet-treated "

f"design.\n"

f" Fix: pass data_generator_kwargs="

f'{{"cohort_periods": [2, 4], "never_treated_frac": 0.0}} '

f"(or a custom data_generator)."

)

# Check anticipation > 0 (all staggered)

antic = getattr(estimator, "anticipation", 0)

if antic > 0:

issues.append(

f" - {name} has anticipation={antic} but the default DGP does "

f"not model anticipatory effects. The estimator will look for "

f"treatment effects {antic} period(s) before the DGP generates "

f"them, biasing power estimates.\n"

f" Fix: supply a custom data_generator that shifts the "

f"effect onset."

)

# Check clean_control on StackedDiD

if name == "StackedDiD":

cc = getattr(estimator, "clean_control", "not_yet_treated")

if cc == "strict" and not has_multi_cohort:

issues.append(

' - StackedDiD has clean_control="strict" but the default '

"single-cohort DGP makes strict controls equivalent to "

"never-treated controls.\n"

" Fix: pass data_generator_kwargs="

'{"cohort_periods": [2, 4]} '

"to test true strict clean-control behavior."

)

if issues:

msg = (

f"Staggered power DGP mismatch for {name}. The default "

f"single-cohort DGP may not match the estimator "

f"configuration:\n" + "\n".join(issues)

)

warnings.warn(msg, UserWarning, stacklevel=2)

def _ddd_effective_n(

n_units: int, data_generator_kwargs: Optional[Dict[str, Any]]

) -> Optional[int]:

"""Return effective DDD sample size, or None if no rounding occurred."""

overrides = data_generator_kwargs or {}

if "n_per_cell" in overrides:

eff = overrides["n_per_cell"] * 8

else:

eff = max(2, n_units // 8) * 8

return eff if eff != n_units else None

def _check_ddd_dgp_compat(

n_units: int,

n_periods: int,

treatment_fraction: float,

treatment_period: int,

data_generator_kwargs: Optional[Dict[str, Any]],

) -> None:

"""Warn when simulation inputs don't match DDD's fixed 2×2×2 design."""

issues: List[str] = []

# DDD is a fixed 2-period factorial; n_periods and treatment_period are ignored

if n_periods != 2:

issues.append(

f"n_periods={n_periods} is ignored (DDD uses a fixed " f"2-period design: pre/post)"

)

if treatment_period != 1:

issues.append(

f"treatment_period={treatment_period} is ignored (DDD "

f"always treats in the second period)"

)

# DDD's 2×2×2 factorial has inherent 50% treatment fraction

if treatment_fraction != 0.5:

issues.append(

f"treatment_fraction={treatment_fraction} is ignored "

f"(DDD uses a balanced 2×2×2 factorial where 50% of "

f"groups are treated)"

)

# n_units rounding: n_per_cell = max(2, n_units // 8)

eff_n = _ddd_effective_n(n_units, data_generator_kwargs)

if eff_n is not None:

eff_n_per_cell = eff_n // 8

issues.append(

f"effective sample size is {eff_n} "

f"(n_per_cell={eff_n_per_cell} × 8 cells), "

f"not the requested n_units={n_units}"

)

if issues:

warnings.warn(

"TripleDifference uses a fixed 2×2×2 factorial DGP "

"(group × partition × time). "

+ "; ".join(issues)

+ ". Pass a custom data_generator for non-standard DDD designs.",

UserWarning,

stacklevel=2,

)

def _check_sdid_placebo_data(

data: pd.DataFrame,

estimator: Any,

est_kwargs: Dict[str, Any],

) -> None:

"""Check SyntheticDiD placebo feasibility on realized data.

This catches infeasible designs on the custom-DGP path where the

pre-generation check (which uses ``n_units * treatment_fraction``)

cannot run because treatment allocation is determined by the DGP.

"""

vm = getattr(estimator, "variance_method", "placebo")

if vm != "placebo":

return

treat_col = est_kwargs.get("treatment", "treat")

unit_col = est_kwargs.get("unit", "unit")

if treat_col not in data.columns or unit_col not in data.columns:

return # fit will fail with a more specific error

unit_treat = data.groupby(unit_col)[treat_col].first()

n_treated = int(unit_treat.sum())

n_control = len(unit_treat) - n_treated

if n_control <= n_treated:

raise ValueError(

f"SyntheticDiD placebo variance requires more control than "

f"treated units, but the generated data has n_control={n_control}, "

f"n_treated={n_treated}. Either adjust your data_generator so that "

f"n_control > n_treated, or use "

f"SyntheticDiD(variance_method='bootstrap') (paper-faithful refit; "

f"~5-30x slower than placebo) or SyntheticDiD(variance_method='jackknife')."

)

# -- Registry construction (deferred to avoid import-time cost) ---------------

_ESTIMATOR_REGISTRY: Optional[Dict[str, _EstimatorProfile]] = None

def _get_registry() -> Dict[str, _EstimatorProfile]:

"""Lazily build and return the estimator registry."""

global _ESTIMATOR_REGISTRY # noqa: PLW0603

if _ESTIMATOR_REGISTRY is not None:

return _ESTIMATOR_REGISTRY

from diff_diff.prep import (

generate_ddd_data,

generate_did_data,

generate_factor_data,

generate_staggered_data,

)

_ESTIMATOR_REGISTRY = {

# --- Basic DiD group ---

"DifferenceInDifferences": _EstimatorProfile(

default_dgp=generate_did_data,

dgp_kwargs_builder=_basic_dgp_kwargs,

fit_kwargs_builder=_basic_fit_kwargs,

result_extractor=_extract_simple,

min_n=20,

"TwoWayFixedEffects": _EstimatorProfile(

default_dgp=generate_did_data,

dgp_kwargs_builder=_basic_dgp_kwargs,

fit_kwargs_builder=_twfe_fit_kwargs,

result_extractor=_extract_simple,

min_n=20,

"MultiPeriodDiD": _EstimatorProfile(

default_dgp=generate_did_data,

dgp_kwargs_builder=_basic_dgp_kwargs,

fit_kwargs_builder=_multiperiod_fit_kwargs,

result_extractor=_extract_multiperiod,

min_n=20,

# --- Staggered group ---

"CallawaySantAnna": _EstimatorProfile(

default_dgp=generate_staggered_data,

dgp_kwargs_builder=_staggered_dgp_kwargs,

fit_kwargs_builder=_staggered_fit_kwargs,