diff --git a/docs/api/business_report.rst b/docs/api/business_report.rst index ae17bffa..0439a499 100644 --- a/docs/api/business_report.rst +++ b/docs/api/business_report.rst @@ -70,11 +70,17 @@ Example .. code-block:: python - from diff_diff import CallawaySantAnna, BusinessReport + from diff_diff import CallawaySantAnna, BusinessReport, generate_staggered_data + + # Staggered-rollout loyalty program across stores + df = generate_staggered_data( + n_units=60, n_periods=10, cohort_periods=[4, 7], + never_treated_frac=0.3, treatment_effect=5.0, seed=42, + ).rename(columns={"unit": "store", "outcome": "revenue"}) cs = CallawaySantAnna(base_period="universal").fit( df, outcome="revenue", unit="store", time="period", - first_treat="first_treat", aggregate="event_study", + first_treat="first_treat", ) report = BusinessReport( cs, diff --git a/docs/api/mmm.rst b/docs/api/mmm.rst index a7caa384..fb4cd0be 100644 --- a/docs/api/mmm.rst +++ b/docs/api/mmm.rst @@ -45,16 +45,27 @@ Example .. code-block:: python - from diff_diff import SyntheticDiD, to_pymc_marketing_lift_test + from diff_diff import SyntheticDiD, generate_synthetic_control_data + from diff_diff import to_pymc_marketing_lift_test # Single-treated-geo experiment (US-CA): TV spend raised there vs control - # geos. With exactly ONE treated geo, the SDID ATT is that geo's per-week - # lift, so labelling the row with its coordinate is sound. (For MULTIPLE - # treated geos the pooled ATT is an average - do not assign it to one geo; - # export each geo's own effect, or omit the geo dim and match aggregate - # spend to an aggregate lift.) + # geos. Unit 0 is the treated geo; the donors are control geos. SyntheticDiD + # needs a block (ever-treated) indicator, so we use the unit-level flag. + panel = generate_synthetic_control_data( + n_donors=20, n_pre=24, n_post=6, treatment_effect=5.0, + effect_type="constant", seed=42, + ).rename(columns={"unit": "geo", "period": "week", "outcome": "revenue", + "treat": "treated"}) + + # Weeks 0..23 are pre-campaign; weeks 24..29 are the campaign window. + post_weeks = [w for w in sorted(panel["week"].unique()) if w >= 24] + # With exactly ONE treated geo, the SDID ATT is that geo's per-week lift, + # so labelling the row with its coordinate is sound. (For MULTIPLE treated + # geos the pooled ATT is an average - do not assign it to one geo; export + # each geo's own effect, or omit the geo dim and match aggregate spend to + # an aggregate lift.) result = SyntheticDiD().fit(panel, outcome='revenue', treatment='treated', - unit='geo', time='week') + unit='geo', time='week', post_periods=post_weeks) df_lift = to_pymc_marketing_lift_test( channel='tv', @@ -79,7 +90,15 @@ Example .. code-block:: python - from diff_diff import DifferenceInDifferences, to_meridian_roi_prior + from diff_diff import DifferenceInDifferences, generate_did_data + from diff_diff import to_meridian_roi_prior + + # Simulated geo experiment: 100 markets, 4 quarters, campaign in quarter 2+. + # Rename the outcome column to 'revenue' to match the MMM's input schema. + panel = generate_did_data( + n_units=100, n_periods=4, treatment_effect=5.0, + treatment_period=2, seed=42, + ).rename(columns={"outcome": "revenue"}) result = DifferenceInDifferences().fit(panel, outcome='revenue', treatment='treated', post='post') diff --git a/docs/api/triple_diff.rst b/docs/api/triple_diff.rst index 8d4a5ee8..4c91540e 100644 --- a/docs/api/triple_diff.rst +++ b/docs/api/triple_diff.rst @@ -38,21 +38,23 @@ engine, so the staggered numbers are identical to the deprecated class's. .. code-block:: python - from diff_diff import TripleDifference + from diff_diff import TripleDifference, generate_ddd_data, generate_staggered_ddd_data - # 2x2x2 design (unchanged) + # 2x2x2 design (unchanged) - columns: outcome, group, partition, time + df = generate_ddd_data(n_per_cell=100, treatment_effect=2.0, seed=42) ddd = TripleDifference(estimation_method="dr") - res = ddd.fit(df, outcome="y", group="state", partition="eligible", post="post") + res = ddd.fit(df, outcome="outcome", group="group", partition="partition", post="time") # staggered adoption - the staggered params are keyword-only + sdf = generate_staggered_ddd_data(n_units=120, n_periods=8, seed=42) sddd = TripleDifference(estimation_method="dr", control_group="not_yet_treated") res = sddd.fit( - df, - outcome="y", - partition="eligible", - unit="id", + sdf, + outcome="outcome", + partition="eligibility", + unit="unit", time="period", - first_treat="enacted", + first_treat="first_treat", aggregate="event_study", ) @@ -129,38 +131,51 @@ Example Usage Basic usage:: - from diff_diff import TripleDifference + from diff_diff import TripleDifference, generate_ddd_data + + # Synthetic DDD panel: group (0/1), partition (0/1), time (0=pre/1=post) + data = generate_ddd_data(n_per_cell=100, treatment_effect=2.0, seed=42) ddd = TripleDifference(estimation_method='dr') results = ddd.fit( data, - outcome='wages', - group='policy_state', # 1=state enacted policy, 0=control state - partition='female', # 1=women (affected by policy), 0=men - post='post' # 1=post-policy, 0=pre-policy + outcome='outcome', + group='group', # 1=state enacted policy, 0=control state + partition='partition', # 1=women (affected by policy), 0=men + post='time' # 1=post-policy, 0=pre-policy ) results.print_summary() With covariates:: + from diff_diff import TripleDifference, generate_ddd_data + + data = generate_ddd_data( + n_per_cell=100, treatment_effect=2.0, + add_covariates=True, seed=42, + ) + + ddd = TripleDifference(estimation_method='dr') results = ddd.fit( data, - outcome='wages', - group='policy_state', - partition='female', - post='post', - covariates=['age', 'education', 'experience'] + outcome='outcome', + group='group', + partition='partition', + post='time', + covariates=['age', 'education'] ) Quick one-call estimation (the ``triple_difference()`` wrapper is deprecated since 3.9 and removed in 4.0):: - from diff_diff import TripleDifference + from diff_diff import TripleDifference, generate_ddd_data + + data = generate_ddd_data(n_per_cell=100, treatment_effect=2.0, seed=42) results = TripleDifference(estimation_method='dr').fit( data, - outcome='wages', - group='policy_state', - partition='female', - post='post', + outcome='outcome', + group='group', + partition='partition', + post='time', ) diff --git a/docs/index.rst b/docs/index.rst index 9294a49b..c9e4cce5 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -10,11 +10,21 @@ It provides sklearn-like estimators with statsmodels-style output for econometri .. code-block:: python - from diff_diff import DifferenceInDifferences + from diff_diff import DifferenceInDifferences, generate_did_data + + # Simulate a panel: 100 units, 10 periods, true treatment lift of 5.0 + data = generate_did_data( + n_units=100, + n_periods=10, + treatment_effect=5.0, + treatment_period=5, + treatment_fraction=0.5, + seed=42, + ) # Fit a basic DiD model did = DifferenceInDifferences() - results = did.fit(data, outcome='y', treatment='treated', post='post') + results = did.fit(data, outcome='outcome', treatment='treated', post='post') print(results.summary()) Key Features diff --git a/docs/practitioner_decision_tree.rst b/docs/practitioner_decision_tree.rst index ca617105..2f8b84bf 100644 --- a/docs/practitioner_decision_tree.rst +++ b/docs/practitioner_decision_tree.rst @@ -299,20 +299,38 @@ identification rests on stronger structural assumptions (Design 1). .. code-block:: python + import numpy as np + import pandas as pd from diff_diff import HeterogeneousAdoptionDiD, did_had_pretest_workflow + # Universal rollout: every market advertises (no holdout), at varying spend. + # Spend = 0 in the pre-period (baseline); each market keeps a fixed spend + # level in the post-period (with a low-spend group as the anchor). + rng = np.random.default_rng(42) + n_markets = 60 + spend = np.where(rng.random(n_markets) < 0.3, 1.0, rng.uniform(1.0, 10.0, n_markets)) + base = rng.normal(100, 10, n_markets) + data = pd.DataFrame( + { + "unit": np.repeat(np.arange(n_markets), 2), + "period": np.tile([1, 2], n_markets), + "dose": np.column_stack([np.zeros(n_markets), spend]).ravel(), + "outcome": np.column_stack([base, base + 3.0 * spend]).ravel(), + } + ) + # Run the pretest battery first - it surfaces violations of the HAD # identification assumptions (it does NOT pick the design path; the # estimator does that internally from the dose support). pretests = did_had_pretest_workflow( - data, outcome="y", unit="unit", + data, outcome="outcome", unit="unit", time="period", dose="dose", ) print(pretests) est = HeterogeneousAdoptionDiD() results = est.fit( - data, outcome="y", unit="unit", + data, outcome="outcome", unit="unit", time="period", dose="dose", ) print(f"Resolved estimand: {results.target_parameter}") @@ -416,10 +434,26 @@ See :doc:`practitioner_getting_started` for an end-to-end example. .. code-block:: python + import pandas as pd from diff_diff import DifferenceInDifferences, SurveyDesign + # Toy survey panel: 4 markets measured before/after the campaign, + # with sampling weights, strata, and geographic clusters (PSUs). + survey = pd.DataFrame( + { + "unit": [0, 0, 1, 1, 2, 2, 3, 3], + "period": [1, 2, 1, 2, 1, 2, 1, 2], + "treated": [1, 1, 1, 1, 0, 0, 0, 0], + "post": [0, 1, 0, 1, 0, 1, 0, 1], + "outcome": [100, 106, 102, 108, 95, 96, 97, 98], + "sample_weight": [1.2, 1.2, 0.9, 0.9, 1.1, 1.1, 1.0, 1.0], + "stratum": ["A", "A", "A", "A", "B", "B", "B", "B"], + "cluster_id": [0, 0, 1, 1, 2, 2, 3, 3], + } + ) + # Reference column names in your data; SurveyDesign resolves them at fit time. - survey = SurveyDesign( + survey_design = SurveyDesign( weights="sample_weight", # observation-level sampling weight strata="stratum", # stratification variable psu="cluster_id", # primary sampling unit (e.g., geography) @@ -427,9 +461,10 @@ See :doc:`practitioner_getting_started` for an end-to-end example. did = DifferenceInDifferences() results = did.fit( - data, outcome="outcome", treatment="treated", - post="post", survey_design=survey, + survey, outcome="outcome", treatment="treated", + post="post", survey_design=survey_design, ) + print(results.summary()) .. tip:: diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 0554c1e9..57e2c8a3 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -43,6 +43,7 @@ The simplest DiD design has two groups (treated/control) and two periods (pre/po treatment_effect=5.0, treatment_period=5, treatment_fraction=0.5, + seed=42, ) # Fit the model @@ -61,13 +62,27 @@ Output: .. code-block:: text - Difference-in-Differences Results - ================================== - ATT: 5.123 - Std. Error: 0.456 - t-statistic: 11.23 - p-value: 0.000 - 95% CI: [4.229, 6.017] + ====================================================================== + Difference-in-Differences Estimation Results + ====================================================================== + + Observations: 1000 + Treated: 500 + Control: 500 + R-squared: 0.7332 + Variance: HC1 heteroskedasticity-robust + + ---------------------------------------------------------------------- + Parameter Estimate Std. Err. t-stat P>|t| + ---------------------------------------------------------------------- + ATT 5.1216 0.2455 20.863 0.0000 *** + ---------------------------------------------------------------------- + + 95% Confidence Interval: [4.6399, 5.6034] + CV (SE/abs(ATT)): 0.0479 + + Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1 + ====================================================================== Using Formula Interface ----------------------- @@ -86,6 +101,13 @@ Control for confounders with the ``covariates`` parameter: .. code-block:: python + import numpy as np + + # Add two confounders to the simulated panel + rng = np.random.default_rng(0) + data["age"] = rng.integers(20, 70, size=len(data)) + data["income"] = rng.normal(50_000, 15_000, size=len(data)).round(0) + results = did.fit( data, outcome='outcome', @@ -94,6 +116,8 @@ Control for confounders with the ``covariates`` parameter: covariates=['age', 'income'] ) + print(f"Covariate-adjusted ATT: {results.att:.4f}") + Cluster-Robust Standard Errors ------------------------------ @@ -101,8 +125,8 @@ For panel data, cluster standard errors at the unit level: .. code-block:: python - did = DifferenceInDifferences(cluster='unit_id') - results = did.fit(data, outcome='y', treatment='treated', post='post') + did = DifferenceInDifferences(cluster='unit') + results = did.fit(data, outcome='outcome', treatment='treated', post='post') Two-Way Fixed Effects --------------------- @@ -118,7 +142,7 @@ For panel data with multiple periods: data, outcome='outcome', treatment='treated', - unit='unit_id', + unit='unit', post='post' ) @@ -135,11 +159,11 @@ effects): from diff_diff import TwoWayFixedEffects event = TwoWayFixedEffects() - results = event.fit( + event_study_results = event.fit( data, outcome='outcome', treatment='treated', - unit='unit_id', + unit='unit', event_study=True, time='period', post_periods=[5, 6, 7, 8, 9], @@ -148,7 +172,7 @@ effects): # Plot the event study from diff_diff.visualization import plot_event_study - ax = plot_event_study(results) + ax = plot_event_study(event_study_results) Staggered Adoption ------------------ @@ -157,13 +181,22 @@ When treatment is adopted at different times across units: .. code-block:: python - from diff_diff import CallawaySantAnna + from diff_diff import CallawaySantAnna, generate_staggered_data + + # Staggered data carries a ``first_treat`` column (0 for never treated) + staggered = generate_staggered_data( + n_units=100, + n_periods=10, + cohort_periods=[4, 7], + never_treated_frac=0.3, + seed=42, + ) cs = CallawaySantAnna() results = cs.fit( - data, + staggered, outcome='outcome', - unit='unit_id', + unit='unit', time='period', first_treat='first_treat' ) diff --git a/tests/test_doc_snippets.py b/tests/test_doc_snippets.py index 4e41ce6e..5d82c5a0 100644 --- a/tests/test_doc_snippets.py +++ b/tests/test_doc_snippets.py @@ -39,6 +39,13 @@ "api/pretrends.rst", "api/power.rst", "api/changes_in_changes.rst", + "api/business_report.rst", + "api/diagnostic_report.rst", + "api/estimators.rst", + "api/mmm.rst", + "api/triple_diff.rst", + "practitioner_decision_tree.rst", + "practitioner_getting_started.rst", "python_comparison.rst", "r_comparison.rst", ] @@ -418,6 +425,7 @@ def _restore_datasets_module(): "r_comparison:block4", "r_comparison:block7", "troubleshooting:block8", + "practitioner_getting_started:block5", }