diff-diff 3.3.3__tar.gz → 3.4.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {diff_diff-3.3.3 → diff_diff-3.4.1}/PKG-INFO +3 -2
  2. {diff_diff-3.3.3 → diff_diff-3.4.1}/README.md +1 -0
  3. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/__init__.py +93 -17
  4. diff_diff-3.4.1/diff_diff/agent_workflow.py +270 -0
  5. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/bacon.py +348 -134
  6. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/business_report.py +18 -3
  7. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/chaisemartin_dhaultfoeuille.py +161 -14
  8. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/chaisemartin_dhaultfoeuille_results.py +35 -19
  9. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/conley.py +92 -34
  10. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/diagnostic_report.py +164 -18
  11. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/efficient_did.py +8 -4
  12. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/estimators.py +153 -91
  13. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/guides/llms-autonomous.txt +6 -0
  14. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/guides/llms-full.txt +68 -4
  15. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/guides/llms.txt +2 -1
  16. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/had.py +63 -7
  17. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/had_pretests.py +68 -0
  18. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/honest_did.py +36 -1
  19. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/linalg.py +648 -131
  20. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/prep.py +1 -0
  21. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/prep_dgp.py +254 -3
  22. diff_diff-3.4.1/diff_diff/pretrends.py +1879 -0
  23. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/results.py +193 -0
  24. diff_diff-3.4.1/diff_diff/spillover.py +3424 -0
  25. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/sun_abraham.py +504 -45
  26. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/survey.py +195 -4
  27. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/trop.py +22 -93
  28. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/trop_global.py +18 -75
  29. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/trop_local.py +123 -0
  30. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/twfe.py +200 -95
  31. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/two_stage.py +1110 -6
  32. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/utils.py +113 -50
  33. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_event_study.py +21 -11
  34. {diff_diff-3.3.3 → diff_diff-3.4.1}/pyproject.toml +7 -2
  35. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/Cargo.lock +5 -5
  36. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/Cargo.toml +1 -1
  37. diff_diff-3.3.3/diff_diff/pretrends.py +0 -1133
  38. {diff_diff-3.3.3 → diff_diff-3.4.1}/LICENSE +0 -0
  39. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/_backend.py +0 -0
  40. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/_guides_api.py +0 -0
  41. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/_nprobust_port.py +0 -0
  42. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/_reporting_helpers.py +0 -0
  43. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/bootstrap_utils.py +0 -0
  44. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
  45. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/continuous_did.py +0 -0
  46. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/continuous_did_bspline.py +0 -0
  47. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/continuous_did_results.py +0 -0
  48. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/datasets.py +0 -0
  49. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/diagnostics.py +0 -0
  50. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/efficient_did_bootstrap.py +0 -0
  51. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/efficient_did_covariates.py +0 -0
  52. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/efficient_did_results.py +0 -0
  53. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/efficient_did_weights.py +0 -0
  54. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/guides/__init__.py +0 -0
  55. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/guides/llms-practitioner.txt +0 -0
  56. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/imputation.py +0 -0
  57. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/imputation_bootstrap.py +0 -0
  58. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/imputation_results.py +0 -0
  59. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/local_linear.py +0 -0
  60. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/power.py +0 -0
  61. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/practitioner.py +0 -0
  62. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/profile.py +0 -0
  63. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/stacked_did.py +0 -0
  64. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/stacked_did_results.py +0 -0
  65. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered.py +0 -0
  66. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered_aggregation.py +0 -0
  67. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered_bootstrap.py +0 -0
  68. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered_results.py +0 -0
  69. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered_triple_diff.py +0 -0
  70. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/staggered_triple_diff_results.py +0 -0
  71. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/synthetic_did.py +0 -0
  72. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/triple_diff.py +0 -0
  73. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/trop_results.py +0 -0
  74. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/two_stage_bootstrap.py +0 -0
  75. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/two_stage_results.py +0 -0
  76. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/__init__.py +0 -0
  77. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_common.py +0 -0
  78. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_continuous.py +0 -0
  79. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_diagnostic.py +0 -0
  80. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_power.py +0 -0
  81. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_staggered.py +0 -0
  82. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/visualization/_synthetic.py +0 -0
  83. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/wooldridge.py +0 -0
  84. {diff_diff-3.3.3 → diff_diff-3.4.1}/diff_diff/wooldridge_results.py +0 -0
  85. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/build.rs +0 -0
  86. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/src/bootstrap.rs +0 -0
  87. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/src/lib.rs +0 -0
  88. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/src/linalg.rs +0 -0
  89. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/src/trop.rs +0 -0
  90. {diff_diff-3.3.3 → diff_diff-3.4.1}/rust/src/weights.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diff-diff
3
- Version: 3.3.3
3
+ Version: 3.4.1
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: Operating System :: OS Independent
@@ -16,7 +16,7 @@ Classifier: Topic :: Scientific/Engineering :: Information Analysis
16
16
  Classifier: Topic :: Scientific/Engineering
17
17
  Requires-Dist: numpy>=1.20.0
18
18
  Requires-Dist: pandas>=1.3.0
19
- Requires-Dist: scipy>=1.7.0
19
+ Requires-Dist: scipy>=1.10
20
20
  Requires-Dist: pytest>=7.0 ; extra == 'dev'
21
21
  Requires-Dist: pytest-xdist>=3.0 ; extra == 'dev'
22
22
  Requires-Dist: pytest-cov>=4.0 ; extra == 'dev'
@@ -159,6 +159,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
159
159
  - [SunAbraham](https://diff-diff.readthedocs.io/en/stable/api/staggered.html) - Sun & Abraham (2021) interaction-weighted estimator for heterogeneity-robust event studies
160
160
  - [ImputationDiD](https://diff-diff.readthedocs.io/en/stable/api/imputation.html) - Borusyak, Jaravel & Spiess (2024) imputation estimator, most efficient under homogeneous effects
161
161
  - [TwoStageDiD](https://diff-diff.readthedocs.io/en/stable/api/two_stage.html) - Gardner (2022) two-stage estimator with GMM sandwich variance
162
+ - [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html) - Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover on near-control units; handles non-staggered and staggered timing; supports survey-design variance under `survey_design=` for HC1 / CR1 (Wave E.1 Binder TSL) and Conley (Wave E.2 panel-aware stratified-Conley sandwich on per-period PSU totals; extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC — `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=<col>`)
162
163
  - [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html) - Synthetic DiD combining standard DiD and synthetic control for few treated units
163
164
  - [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
164
165
  - [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html) - Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
@@ -106,6 +106,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
106
106
  - [SunAbraham](https://diff-diff.readthedocs.io/en/stable/api/staggered.html) - Sun & Abraham (2021) interaction-weighted estimator for heterogeneity-robust event studies
107
107
  - [ImputationDiD](https://diff-diff.readthedocs.io/en/stable/api/imputation.html) - Borusyak, Jaravel & Spiess (2024) imputation estimator, most efficient under homogeneous effects
108
108
  - [TwoStageDiD](https://diff-diff.readthedocs.io/en/stable/api/two_stage.html) - Gardner (2022) two-stage estimator with GMM sandwich variance
109
+ - [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html) - Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover on near-control units; handles non-staggered and staggered timing; supports survey-design variance under `survey_design=` for HC1 / CR1 (Wave E.1 Binder TSL) and Conley (Wave E.2 panel-aware stratified-Conley sandwich on per-period PSU totals; extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC — `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=<col>`)
109
110
  - [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html) - Synthetic DiD combining standard DiD and synthetic control for few treated units
110
111
  - [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
111
112
  - [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html) - Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
@@ -1,23 +1,23 @@
1
- """
2
- diff-diff: A library for Difference-in-Differences analysis.
3
-
4
- This library provides sklearn-like estimators for causal inference
5
- using the difference-in-differences methodology.
1
+ """diff-diff: Difference-in-Differences causal inference with sklearn-like API.
2
+ Recommended starting call for LLM agents:
3
+ ``diff_diff.agent_workflow(df, unit=..., time=..., treatment=..., outcome=...)``
4
+ prints a copy-pasteable workflow with your column names wired in.
6
5
 
7
- For AI agents:
6
+ The orchestrator names the full sequence:
8
7
 
9
- 1. Describe your data: ``diff_diff.profile_panel(df, unit=..., time=...,
10
- treatment=..., outcome=...)``
11
- 2. Consult the reference: ``diff_diff.get_llm_guide("autonomous")``
8
+ 1. Describe the panel: diff_diff.profile_panel(df, ...)
9
+ 2. Choose an estimator: diff_diff.get_llm_guide("autonomous")
12
10
  (estimator-support matrix + reasoning)
13
- 3. Follow the workflow: ``diff_diff.get_llm_guide("practitioner")``
14
- (Baker et al. (2025) 8-step recipe)
15
- 4. Report results: ``diff_diff.BusinessReport(results)``
16
- (structured agent-legible output)
11
+ 3. Fit: <Estimator>(...).fit(df, ...)
12
+ 4. Validate: diff_diff.practitioner_next_steps(result)
13
+ 5. Report: diff_diff.BusinessReport(result)
14
+
15
+ For a comprehensive API reference call ``diff_diff.get_llm_guide("full")``.
16
+ For the Baker et al. (2025) 8-step practitioner recipe call
17
+ ``diff_diff.get_llm_guide("practitioner")``.
17
18
 
18
- For a comprehensive API reference call ``diff_diff.get_llm_guide("full")``;
19
- ``practitioner_next_steps(results)`` returns context-aware guidance after
20
- any estimator's ``fit()``.
19
+ This library provides sklearn-like estimators for causal inference using
20
+ the difference-in-differences methodology.
21
21
  """
22
22
 
23
23
  # Import backend detection from dedicated module (avoids circular imports)
@@ -125,6 +125,7 @@ from diff_diff.prep import (
125
125
  generate_continuous_did_data,
126
126
  generate_did_data,
127
127
  generate_ddd_data,
128
+ generate_ddd_panel_data,
128
129
  generate_event_study_data,
129
130
  generate_factor_data,
130
131
  generate_panel_data,
@@ -171,6 +172,10 @@ from diff_diff.two_stage import (
171
172
  TwoStageDiDResults,
172
173
  two_stage_did,
173
174
  )
175
+ from diff_diff.spillover import (
176
+ SpilloverDiD,
177
+ )
178
+ from diff_diff.results import SpilloverDiDResults # re-export
174
179
  from diff_diff.stacked_did import (
175
180
  StackedDiD,
176
181
  StackedDiDResults,
@@ -251,6 +256,7 @@ from diff_diff.diagnostic_report import (
251
256
  DiagnosticReportResults,
252
257
  )
253
258
  from diff_diff._guides_api import get_llm_guide
259
+ from diff_diff.agent_workflow import agent_workflow
254
260
  from diff_diff.profile import (
255
261
  Alert,
256
262
  OutcomeShape,
@@ -287,7 +293,7 @@ ETWFE = WooldridgeDiD
287
293
  DCDH = ChaisemartinDHaultfoeuille
288
294
  HAD = HeterogeneousAdoptionDiD
289
295
 
290
- __version__ = "3.3.3"
296
+ __version__ = "3.4.1"
291
297
  __all__ = [
292
298
  # Estimators
293
299
  "DifferenceInDifferences",
@@ -300,6 +306,7 @@ __all__ = [
300
306
  "SunAbraham",
301
307
  "ImputationDiD",
302
308
  "TwoStageDiD",
309
+ "SpilloverDiD",
303
310
  "TripleDifference",
304
311
  "TROP",
305
312
  "StackedDiD",
@@ -341,6 +348,7 @@ __all__ = [
341
348
  "TwoStageDiDResults",
342
349
  "TwoStageBootstrapResults",
343
350
  "two_stage_did",
351
+ "SpilloverDiDResults",
344
352
  "TripleDifferenceResults",
345
353
  "triple_difference",
346
354
  "StaggeredTripleDifference",
@@ -403,6 +411,7 @@ __all__ = [
403
411
  "generate_staggered_data",
404
412
  "generate_factor_data",
405
413
  "generate_ddd_data",
414
+ "generate_ddd_panel_data",
406
415
  "generate_panel_data",
407
416
  "generate_event_study_data",
408
417
  "generate_staggered_ddd_data",
@@ -495,6 +504,7 @@ __all__ = [
495
504
  "list_datasets",
496
505
  "clear_cache",
497
506
  # Practitioner guidance
507
+ "agent_workflow",
498
508
  "practitioner_next_steps",
499
509
  "BusinessReport",
500
510
  "BusinessContext",
@@ -511,3 +521,69 @@ __all__ = [
511
521
  # LLM guide accessor
512
522
  "get_llm_guide",
513
523
  ]
524
+
525
+ # Agent-facing entrypoints surface first in dir(diff_diff). LLM agents
526
+ # follow a `dir -> help -> docstring -> use` discovery loop; surfacing
527
+ # these names first measurably improves discoverability vs the default
528
+ # alphabetic ordering. Internal — read by tests/test_agent_discoverability.py.
529
+ _AGENT_FACING_ORDER = (
530
+ "agent_workflow",
531
+ "profile_panel",
532
+ "get_llm_guide",
533
+ "practitioner_next_steps",
534
+ "BusinessReport",
535
+ "DiagnosticReport",
536
+ )
537
+
538
+
539
+ class _OrderedName(str):
540
+ """str subclass that sorts by _AGENT_FACING_ORDER priority.
541
+
542
+ Python's built-in dir() always sorts the result of __dir__()
543
+ alphabetically (CPython Objects/object.c::_dir_object unconditionally
544
+ calls PyList_Sort), so returning a list in our preferred order is
545
+ not enough. But PyList_Sort uses __lt__ for comparisons, so a str
546
+ subclass with a custom __lt__ can subvert the alphabetic default
547
+ while remaining a fully usable str for every other operation.
548
+
549
+ ALL names returned by __dir__() must be _OrderedName, not just the
550
+ priority head: when Python compares an _OrderedName against a plain
551
+ str, the reflected-method protocol prefers str's inherited __gt__
552
+ (because _OrderedName is a subclass of str), which sorts purely
553
+ alphabetically and breaks the ordering. With every element wrapped,
554
+ all comparisons go through this __lt__: priority head sorts to
555
+ front, tail (default priority 1<<30) falls through to alphabetic
556
+ via str.__lt__.
557
+ """
558
+
559
+ _ORDER = {n: i for i, n in enumerate(_AGENT_FACING_ORDER)}
560
+
561
+ def __lt__(self, other):
562
+ sp = self._ORDER.get(str(self), 1 << 30)
563
+ op = self._ORDER.get(str(other), 1 << 30)
564
+ if sp != op:
565
+ return sp < op
566
+ return str.__lt__(self, other)
567
+
568
+
569
+ def __dir__():
570
+ """Surface agent-facing entrypoints first; remainder alphabetic.
571
+
572
+ Returns the full module namespace (matching default `dir(module)`
573
+ membership — keeps `__doc__`, `__name__`, etc. accessible via
574
+ `inspect.getmembers`) with priority names re-ordered to the head
575
+ via `_OrderedName`'s custom `__lt__`.
576
+
577
+ `__all__` order does not affect `dir(module)`. CPython sorts the
578
+ result of `__dir__()` alphabetically, so we return `_OrderedName`
579
+ instances (str subclass with custom `__lt__`) for every name; the
580
+ custom comparison routes head names to the top and falls back to
581
+ alphabetic for everyone else. See `_OrderedName` docstring for
582
+ why ALL names must be wrapped (mixing plain `str` with the
583
+ subclass triggers Python's reflected-method comparison protocol
584
+ and breaks the ordering).
585
+
586
+ `from diff_diff import *` semantics are unaffected (driven by
587
+ `__all__`, not by `dir()`).
588
+ """
589
+ return [_OrderedName(n) for n in globals()]
@@ -0,0 +1,270 @@
1
+ """Stateless orchestrator: print the recommended diff-diff workflow with
2
+ the caller's column names wired in.
3
+
4
+ This module exists to give LLM agents a single, recognizable entrypoint
5
+ that names the rest of the agent-facing workflow (`profile_panel`,
6
+ `get_llm_guide`, `practitioner_next_steps`, `BusinessReport`). The
7
+ function does not fit, inspect, or recommend — it templates a copy-
8
+ pasteable script.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from typing import Any, Dict, List, Optional, Tuple
14
+
15
+ # Pattern → df-callable estimator class names. Flat union below is the
16
+ # `fit_candidates` field of the returned dict; each name must remain a
17
+ # valid `hasattr(diff_diff, name)` (locked by the contract test in
18
+ # tests/test_agent_discoverability.py and tests/test_agent_workflow.py).
19
+ # Patterns intentionally exclude post-fit and pre-fit diagnostics
20
+ # (PreTrendsPower takes pre-treatment coefficients, HonestDiD takes a
21
+ # fitted results object); those are mentioned separately in the
22
+ # templated Step 4 of the script.
23
+ _WORKFLOW_PATTERNS: Tuple[Tuple[str, Tuple[str, ...]], ...] = (
24
+ (
25
+ "Staggered adoption + binary treatment + has_never_treated control",
26
+ ("CallawaySantAnna", "SunAbraham", "ImputationDiD"),
27
+ ),
28
+ (
29
+ "Continuous treatment dose (non-binary numeric intensity)",
30
+ ("ContinuousDiD",),
31
+ ),
32
+ (
33
+ "Heterogeneous adoption intensity across treated units",
34
+ ("HeterogeneousAdoptionDiD",),
35
+ ),
36
+ (
37
+ "Simple 2x2 DiD (binary treatment, two periods, no staggering)",
38
+ ("DifferenceInDifferences",),
39
+ ),
40
+ )
41
+
42
+
43
+ def _safe_kwarg(name: str, value: Optional[str]) -> Optional[str]:
44
+ """Render ``name=<python-literal>`` using repr() for source-safety.
45
+
46
+ Column labels containing quotes, backslashes, or other special
47
+ characters must not break the emitted "copy-pasteable" script.
48
+ Python's built-in ``repr()`` produces a valid string literal for
49
+ any str input (including embedded quotes / backslashes /
50
+ newlines), so ``f"{name}={value!r}"`` is injection-safe by
51
+ construction. ``None`` returns ``None`` so the caller can drop
52
+ the kwarg.
53
+ """
54
+ if value is None:
55
+ return None
56
+ return f"{name}={value!r}"
57
+
58
+
59
+ def _join_kwargs(**kwargs: Optional[str]) -> str:
60
+ parts = [_safe_kwarg(k, v) for k, v in kwargs.items()]
61
+ return ", ".join(p for p in parts if p is not None)
62
+
63
+
64
+ def agent_workflow(
65
+ df: Any,
66
+ *,
67
+ unit: str,
68
+ time: str,
69
+ treatment: str,
70
+ outcome: str,
71
+ first_treat: Optional[str] = None,
72
+ df_name: str = "df",
73
+ verbose: bool = True,
74
+ ) -> Dict[str, Any]:
75
+ """Print the recommended diff-diff workflow with your column names wired in.
76
+
77
+ Stateless orchestrator. Calls nothing internally. Returns a dict;
78
+ optionally prints a copy-pasteable script (``verbose=True``, the
79
+ default). ``df`` is not inspected — column names are templated
80
+ verbatim into the output.
81
+
82
+ Parameters
83
+ ----------
84
+ df : pandas.DataFrame
85
+ Long-format panel data. Not inspected; included so the agent
86
+ can pass the same handle along to the next call.
87
+ unit : str
88
+ Column identifying the cross-sectional unit.
89
+ time : str
90
+ Column identifying the time period.
91
+ treatment : str
92
+ Column holding the treatment indicator or dose.
93
+ outcome : str
94
+ Column holding the outcome variable.
95
+ first_treat : str, optional
96
+ Column with each unit's first-treatment period (or NaN for
97
+ never-treated controls). When supplied, the templated Step 3
98
+ switches from a ``DifferenceInDifferences.fit(treatment=...)``
99
+ example to a ``CallawaySantAnna().fit(first_treat=...)``
100
+ example, matching the actual fit signatures (passing
101
+ ``treatment=`` to CallawaySantAnna's ``.fit()`` would raise
102
+ TypeError).
103
+ df_name : str, default ``"df"``
104
+ Identifier under which the caller's dataframe is bound in
105
+ their namespace. Templated verbatim into the emitted script
106
+ as the first positional argument of every call
107
+ (``profile_panel({df_name}, ...)``,
108
+ ``<Estimator>().fit({df_name}, ...)``) so the script is
109
+ directly executable when the caller's local variable matches.
110
+ If the caller has ``panel = pd.read_parquet(...)``, passing
111
+ ``df_name="panel"`` produces a script that references
112
+ ``panel`` instead of ``df``. Must be a valid Python identifier
113
+ (not enforced; non-identifier values produce a script that
114
+ won't parse).
115
+ verbose : bool, default True
116
+ If True, print the script to stdout. The dict is always
117
+ returned regardless.
118
+
119
+ Returns
120
+ -------
121
+ dict
122
+ Keys:
123
+
124
+ - ``"profile_call"`` (str): call signature for
125
+ :func:`diff_diff.profile_panel`.
126
+ - ``"guide_call"`` (str): call signature for
127
+ :func:`diff_diff.get_llm_guide`.
128
+ - ``"fit_candidates"`` (list of str): flat union of estimator /
129
+ diagnostic class names referenced in the workflow patterns.
130
+ Every name resolves on the top-level ``diff_diff`` namespace.
131
+ - ``"validation_calls"`` (list of str): call signatures for the
132
+ post-fit validation step.
133
+ - ``"reporting_call"`` (str): call signature for
134
+ :class:`diff_diff.BusinessReport`.
135
+ - ``"script"`` (str): printable multi-line workflow.
136
+
137
+ Examples
138
+ --------
139
+ >>> import pandas as pd
140
+ >>> import diff_diff
141
+ >>> df = pd.DataFrame({
142
+ ... "firm_id": [1, 1, 2, 2],
143
+ ... "year": [0, 1, 0, 1],
144
+ ... "treated": [0, 0, 1, 1],
145
+ ... "logwage": [0.1, 0.2, 0.1, 0.9],
146
+ ... })
147
+ >>> out = diff_diff.agent_workflow(df, unit="firm_id", time="year",
148
+ ... treatment="treated", outcome="logwage",
149
+ ... verbose=False)
150
+ >>> "profile_panel" in out["script"]
151
+ True
152
+ """
153
+ del df # intentionally unused: orchestrator templates from column names only
154
+
155
+ profile_call = (
156
+ f"diff_diff.profile_panel({df_name}, "
157
+ f"{_join_kwargs(unit=unit, time=time, treatment=treatment, outcome=outcome)})"
158
+ )
159
+ guide_call = 'diff_diff.get_llm_guide("autonomous")'
160
+
161
+ # Step 3 example: branch on first_treat presence.
162
+ # - With first_treat: a staggered structure is strongly implied, BUT
163
+ # `first_treat` does not by itself identify which estimator to use:
164
+ # CallawaySantAnna (binary staggered), ContinuousDiD (continuous-
165
+ # dose with first_treat), and HeterogeneousAdoptionDiD event-study
166
+ # (heterogeneous intensity with first_treat_col) all accept it.
167
+ # Show CallawaySantAnna as the binary-staggered canonical example
168
+ # and list the alternatives for continuous / heterogeneous designs
169
+ # so an agent isn't steered to the wrong estimator.
170
+ # - Without first_treat: the orchestrator does not inspect df, so it
171
+ # CANNOT infer whether the panel is 2x2 binary vs continuous-dose
172
+ # vs heterogeneous-adoption. Show a DifferenceInDifferences call
173
+ # as the "simple 2x2" example and label it explicitly conditional
174
+ # on that shape.
175
+ if first_treat is not None:
176
+ fit_example_kwargs = _join_kwargs(
177
+ outcome=outcome, unit=unit, time=time, first_treat=first_treat
178
+ )
179
+ fit_example_call = f"diff_diff.CallawaySantAnna().fit({df_name}, {fit_example_kwargs})"
180
+ step3_label_lines = [
181
+ "Step 3 - Fit. Your data has `first_treat` -> staggered structure.",
182
+ "`first_treat` alone does NOT identify a single estimator; pick by",
183
+ "treatment shape:",
184
+ " - Binary staggered : CallawaySantAnna (shown) / SunAbraham / ImputationDiD",
185
+ " - Continuous dose : ContinuousDiD (also takes first_treat=)",
186
+ " - Heterogeneous adoption intensity:",
187
+ " HeterogeneousAdoptionDiD (event study,",
188
+ " takes first_treat_col=, NOT first_treat=)",
189
+ ]
190
+ else:
191
+ fit_example_kwargs = _join_kwargs(
192
+ outcome=outcome, unit=unit, time=time, treatment=treatment
193
+ )
194
+ fit_example_call = (
195
+ f"diff_diff.DifferenceInDifferences().fit({df_name}, {fit_example_kwargs})"
196
+ )
197
+ step3_label_lines = [
198
+ "Step 3 - Fit. Pick a candidate from Step 2's patterns based on your",
199
+ "treatment/time shape. The example below shows the simple 2x2 case",
200
+ "(binary treatment + binary time); substitute ContinuousDiD /",
201
+ "HeterogeneousAdoptionDiD / etc. when your design is not 2x2",
202
+ "(DifferenceInDifferences.fit() validates and rejects non-binary",
203
+ "treatment or time).",
204
+ ]
205
+ step3_comment_block = "\n".join(f"# {line}" for line in step3_label_lines)
206
+
207
+ validation_calls = [
208
+ "diff_diff.practitioner_next_steps(result)",
209
+ ]
210
+ reporting_call = "diff_diff.BusinessReport(result).full_report()"
211
+
212
+ fit_candidates: List[str] = []
213
+ pattern_lines: List[str] = []
214
+ for label, names in _WORKFLOW_PATTERNS:
215
+ pattern_lines.append(f"# - {label}")
216
+ pattern_lines.append(f"# candidates: {', '.join(names)}")
217
+ for n in names:
218
+ if n not in fit_candidates:
219
+ fit_candidates.append(n)
220
+ pattern_block = "\n".join(pattern_lines)
221
+
222
+ diagnostics_block = (
223
+ "# Parallel-trends sensitivity / power (take a fitted result or\n"
224
+ "# pre-trend coefficients, NOT df+columns):\n"
225
+ "# diff_diff.PreTrendsPower / diff_diff.HonestDiD"
226
+ )
227
+
228
+ # Templated output is a valid Python script: every prose line is a
229
+ # `#` comment, every code line stands at column 0 and runs as-is.
230
+ # Step 5 wraps full_report() in print() so end-to-end execution
231
+ # actually produces the stakeholder narrative.
232
+ script = f"""# diff_diff workflow for your data
233
+ # =================================
234
+ #
235
+ # Step 1 - Describe the panel:
236
+ profile = {profile_call}
237
+ print(profile)
238
+
239
+ # Step 2 - Choose an estimator. Consult the routing matrix:
240
+ print({guide_call})
241
+
242
+ # Routing patterns (df-callable estimators):
243
+ {pattern_block}
244
+ #
245
+ {diagnostics_block}
246
+
247
+ {step3_comment_block}
248
+ result = {fit_example_call}
249
+
250
+ # Step 4 - Validate:
251
+ {validation_calls[0]}
252
+
253
+ # Step 5 - Report:
254
+ print({reporting_call})
255
+
256
+ # Full reference: diff_diff.get_llm_guide("full")
257
+ # Practitioner recipe: diff_diff.get_llm_guide("practitioner")
258
+ """
259
+
260
+ if verbose:
261
+ print(script)
262
+
263
+ return {
264
+ "profile_call": profile_call,
265
+ "guide_call": guide_call,
266
+ "fit_candidates": fit_candidates,
267
+ "validation_calls": validation_calls,
268
+ "reporting_call": reporting_call,
269
+ "script": script,
270
+ }