Skip to content

docs: make all RST code examples runnable in Python #760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wenddymacro wants to merge 1 commit into
base: main
Choose a base branch
from
Open

docs: make all RST code examples runnable in Python #760

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions docs/api/business_report.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
35 changes: 27 additions & 8 deletions docs/api/mmm.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand All @@ -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')
Expand Down
61 changes: 38 additions & 23 deletions docs/api/triple_diff.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)

Expand Down Expand Up @@ -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',
)
14 changes: 12 additions & 2 deletions docs/index.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
45 changes: 40 additions & 5 deletions docs/practitioner_decision_tree.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -416,20 +434,37 @@ 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)
)

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::

Expand Down
Loading