Demos

Two end-to-end walkthroughs ship in the repository root. Each one is a single self-contained script you can run with python demo_<name>.py after installing smmargins.

Williams-style logit walkthrough

demo_margins.py reproduces, on a simulated voting dataset, every core statistic in Richard Williams’ Margins01 notes and then exercises the 0.3 inference surface end-to-end:

1. Adjusted predictions at specific values (APR / margins, at(...))

2. APM vs AAP (margins, atmeans vs margins)

3. MER vs MEM vs AME for a continuous covariate

4. Discrete contrasts for a multi-level categorical variable

5. Discrete change for a 0/1 dummy

6. Williams’ classic interaction example: AME of age by sex

7. Predicted probability over age, by sex (table for plotting)

8. Analytic vs FD parity check

9. Robust covariance via cov_type="HC3"

10. Krinsky–Robb simulation VCE

11. Pairs bootstrap VCE

12. Simultaneous CIs via sup-t

13. Cluster-robust SEs (cov_type="cluster" with cov_kwds=)

14. Multiple-comparison adjustments side-by-side (Bonferroni / Šidák / sup-t)

15. User-supplied parameter covariance (vcov=)

Highlights from the script

Fit a logit with an interaction:

fit = smf.logit(
    "voted ~ age + income + C(educ) + female + age:female",
    data=df,
).fit()
M = Margins(fit)

APR — predictions at policy-relevant ages, averaging everything else over the sample:

M.predict(atexog={"age": [25, 45, 65]})

MER, MEM, and AME for age — these can differ meaningfully in nonlinear models with interactions:

M.dydx("age", atexog={"age": [25, 45, 65]})   # MER
M.dydx("age", at="mean")                      # MEM
M.dydx("age")                                 # AME

Discrete AME for a multi-level factor with an explicit reference level:

M.dydx("educ", reference="college")

Williams’ interaction lesson — same model, AME of age for each sex:

M.dydx("age", atexog={"female": [0, 1]})

Robust SEs and alternative VCEs (sections 9–11):

Margins(fit, cov_type="HC3").dydx("age")          # heteroskedastic-robust
M.dydx("age", vce="simulation",
       n_sims=2000, sim_seed=42)                   # Krinsky–Robb
M.dydx("age", vce="bootstrap",
       n_boot=500, boot_seed=42)                   # pairs bootstrap

Cluster-robust SEs through cov_type="cluster" with cluster IDs passed in cov_kwds (section 13):

Margins(fit, cov_type="cluster",
        cov_kwds={"groups": df["household"]}).dydx("age")

Family-wise CI methods side-by-side at five ages (section 14) — for a correlated family of predictions, sup-t is typically narrower than Bonferroni / Šidák:

common = dict(atexog={"age": [25, 35, 45, 55, 65]},
              vce="simulation", n_sims=4000, sim_seed=123)
M.predict(**common, ci_method="pointwise")
M.predict(**common, ci_method="bonferroni")
M.predict(**common, ci_method="sidak")
M.predict(**common, ci_method="sup-t")

User-supplied parameter covariance (section 15) — drop in any \((k, k)\) matrix and smmargins sandwiches it through the Jacobian:

Margins(fit, vcov=my_vcov_matrix).dydx("age")

Full source

"""
demo_margins.py
===============

Walkthrough of the core analyses in Richard Williams' *Margins01* notes
(https://academicweb.nd.edu/~rwilliam/stats/Margins01.pdf), implemented
on top of StatsModels + patsy + the ``smmargins`` package.

Sections
--------
  1.  Adjusted predictions at specific values (APR / ``margins, at(...)``)
  2.  APM vs AAP (``margins, atmeans`` vs ``margins``)
  3.  MER vs MEM vs AME for a continuous covariate
  4.  Discrete contrast for a categorical variable
  5.  Discrete change for a 0/1 dummy
  6.  AME by interaction subgroup (Williams' motivating example)
  7.  Predicted probability over age, by sex (table for plotting)
  8.  Analytic vs FD parity check
  9.  Robust covariance (``cov_type="HC3"``)
  10. Krinsky–Robb simulation VCE
  11. Pairs bootstrap VCE
  12. Simultaneous CIs via sup-t
  13. Cluster-robust SEs (``cov_type="cluster"``)
  14. Multiple-comparison adjustments side-by-side (Bonferroni / Sidak / sup-t)
  15. User-supplied parameter covariance (``vcov=``)
  16. Prediction scales (``scale=``) and a custom ``Transform``
  17. Subgroup AMEs via ``over=`` and observation weights
  18. Joint Wald test and pairwise comparisons
"""

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

from smmargins import Margins, Transform

pd.options.display.width = 120
pd.options.display.float_format = "{: .4f}".format

# ---------------------------------------------------------------------------
# Simulate a binary-outcome dataset with structure similar to Williams' notes
# ---------------------------------------------------------------------------
rng = np.random.default_rng(7)
N = 5_000
df = pd.DataFrame(
    {
        "age":    rng.normal(45, 12, N).clip(18, 90),
        "income": rng.lognormal(10.5, 0.4, N),          # ~36k median
        "educ":   rng.choice(["hs", "college", "grad"], N, p=[0.4, 0.4, 0.2]),
        "female": rng.integers(0, 2, N),
    }
)
eta = (
    -4.0
    + 0.05 * df["age"]
    + 0.00001 * df["income"]
    + 0.8 * (df["educ"] == "college")
    + 1.4 * (df["educ"] == "grad")
    + 0.3 * df["female"]
    - 0.0004 * df["age"] * (df["female"])        # interaction
)
df["voted"] = (rng.uniform(0, 1, N) < 1 / (1 + np.exp(-eta))).astype(int)

print("Sample:")
print(df.head(3), "\n")

# ---------------------------------------------------------------------------
# Fit a logit with an interaction, like the Williams example
# ---------------------------------------------------------------------------
fit = smf.logit(
    "voted ~ age + income + C(educ) + female + age:female",
    data=df,
).fit(disp=False)
print("=" * 80)
print("Fitted logit")
print("=" * 80)
print(fit.summary().tables[1])
print()

# `analytic=True` is the default: the outer ∂g/∂β goes through
# `family.link.inverse_deriv` for any GLM (Logit/Probit/Poisson/...) and
# the identity link for OLS/WLS/GLS, falling back to central finite
# differences only when the link derivative isn't available. Set
# `analytic=False` to force FD; you'll get the same answers (see the
# parity check at the bottom of this file) but pay p extra forward
# predict() calls per statistic.
M = Margins(fit)

# ---------------------------------------------------------------------------
# 1. Adjusted predictions at representative values (APR)
#    Stata: margins, at(age=(25 45 65))
# ---------------------------------------------------------------------------
print("=" * 80)
print("1. APR  (predict at age=25,45,65; everything else at sample values)")
print("=" * 80)
print(M.predict(atexog={"age": [25, 45, 65]}))
print()

# ---------------------------------------------------------------------------
# 2. Adjusted prediction at means (APM)  vs  average adjusted prediction (AAP)
# ---------------------------------------------------------------------------
print("=" * 80)
print("2. APM  (margins, atmeans)   vs   AAP  (margins)")
print("=" * 80)
print("APM:"); print(M.predict(at="mean"))
print("\nAAP:"); print(M.predict())
print()

# ---------------------------------------------------------------------------
# 3. Marginal effect: MER vs MEM vs AME for `age`
#    (Williams points out these three can differ meaningfully in nonlinear
#    models with interactions)
# ---------------------------------------------------------------------------
print("=" * 80)
print("3. d Pr(voted)/d age : MER (at age=25,45,65),  MEM, and AME")
print("=" * 80)
print("MER (at age=25,45,65):")
print(M.dydx("age", atexog={"age": [25, 45, 65]}))
print("\nMEM (at means of everything):")
print(M.dydx("age", at="mean"))
print("\nAME (averaged over the sample):")
print(M.dydx("age"))
print()

# ---------------------------------------------------------------------------
# 4. Discrete contrast for the categorical variable `educ`
# ---------------------------------------------------------------------------
print("=" * 80)
print("4. Discrete AME for educ  (each level vs 'college' as reference)")
print("=" * 80)
print(M.dydx("educ", reference="college"))
print()

# ---------------------------------------------------------------------------
# 5. Discrete change for the dummy `female`  (auto-detected as discrete)
# ---------------------------------------------------------------------------
print("=" * 80)
print("5. AME for female (0/1 dummy):  Pr(voted|female=1) - Pr(voted|female=0)")
print("=" * 80)
print(M.dydx("female"))
print()

# ---------------------------------------------------------------------------
# 6. Interaction-sensitivity: marginal effect of age, separately for men/women
#    This is Williams' classic motivating example: the interaction coefficient
#    alone tells you little about what the marginal effect actually is for any
#    given subpopulation.
# ---------------------------------------------------------------------------
print("=" * 80)
print("6. AME of age, separately by sex  (Williams' interaction illustration)")
print("=" * 80)
print(M.dydx("age", atexog={"female": [0, 1]}))
print()

# ---------------------------------------------------------------------------
# 7. Adjusted predictions, age by sex — table suitable for plotting
# ---------------------------------------------------------------------------
print("=" * 80)
print("7. Predicted Pr(voted) over age, for each sex")
print("=" * 80)
tbl = M.predict(atexog={"age": list(range(20, 91, 10)), "female": [0, 1]})
print(tbl)

# ---------------------------------------------------------------------------
# 8. Analytic vs FD: same answers, faster path
#    Logit exposes `family.link.inverse_deriv`, so the analytic outer
#    Jacobian is used by default. Toggling `analytic=False` reroutes
#    every statistic through central finite differences — useful as a
#    sanity check or when working with a custom Link subclass that
#    doesn't implement inverse_deriv.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("8. Analytic vs FD — same numbers, taken via different paths")
print("=" * 80)
M_fd = Margins(fit, analytic=False)
ame_an = M.dydx("age")
ame_fd = M_fd.dydx("age")
print(f"AME(age) analytic : est={ame_an.estimate[0]: .8f}  se={ame_an.se[0]: .8f}")
print(f"AME(age) FD       : est={ame_fd.estimate[0]: .8f}  se={ame_fd.se[0]: .8f}")
print(f"max abs diff      : "
      f"est {abs(ame_an.estimate[0] - ame_fd.estimate[0]): .2e}, "
      f"se {abs(ame_an.se[0] - ame_fd.se[0]): .2e}")

# ---------------------------------------------------------------------------
# 9. Robust covariance (Feature 1)
#    Recompute SEs with HC3 heteroskedasticity-consistent covariance.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("9. Robust covariance — HC3")
print("=" * 80)
M_hc3 = Margins(fit, cov_type="HC3")
print(M_hc3.dydx("age"))

# ---------------------------------------------------------------------------
# 10. Krinsky–Robb simulation VCE (Feature 2)
#     Draw parameters from their sampling distribution and evaluate margins.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("10. Krinsky–Robb simulation VCE")
print("=" * 80)
print(M.dydx("age", vce="simulation", n_sims=2000, sim_seed=42))

# ---------------------------------------------------------------------------
# 11. Bootstrap VCE (Feature 3)
#     Pairs bootstrap with 500 replications.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("11. Bootstrap VCE")
print("=" * 80)
print(M.dydx("age", vce="bootstrap", n_boot=500, boot_seed=42))

# ---------------------------------------------------------------------------
# 12. Simultaneous CIs — sup-t (Feature 4)
#     Use simulation draws to compute simultaneous CIs for a family of margins.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("12. Simultaneous CIs (sup-t)")
print("=" * 80)
print(M.predict(atexog={"age": [25, 45, 65]},
                vce="simulation", n_sims=2000, sim_seed=42,
                ci_method="sup-t"))

# ---------------------------------------------------------------------------
# 13. Cluster-robust SEs
#     Synthesize a clustering structure (e.g., households of ~10 voters who
#     share unobserved local effects). Cluster-robust SEs propagate that
#     correlation through the Jacobian to the AME.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("13. Cluster-robust SEs vs nonrobust  (synthetic household clusters)")
print("=" * 80)
df_c = df.copy()
df_c["household"] = rng.integers(0, N // 10, N)  # ~10 obs per cluster
fit_c = smf.logit(
    "voted ~ age + income + C(educ) + female + age:female",
    data=df_c,
).fit(disp=False)
M_nonrobust = Margins(fit_c)
M_cluster = Margins(fit_c, cov_type="cluster",
                    cov_kwds={"groups": df_c["household"]})
ame_nr = M_nonrobust.dydx("age").se[0]
ame_cl = M_cluster.dydx("age").se[0]
print(f"AME(age) SE — nonrobust : {ame_nr: .6f}")
print(f"AME(age) SE — cluster    : {ame_cl: .6f}   (ratio {ame_cl / ame_nr:.2f}x)")

# ---------------------------------------------------------------------------
# 14. Multiple-comparison adjustments
#     A family of 5 marginal effects at different ages. Pointwise CIs
#     under-cover the joint event "all 5 contain the truth"; Bonferroni
#     and Sidak inflate the critical value uniformly; sup-t uses the
#     simulation draws to exploit correlation across the family.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("14. Family-wise CI methods at age=25,35,45,55,65")
print("=" * 80)
ages = [25, 35, 45, 55, 65]
common = dict(atexog={"age": ages}, vce="simulation",
              n_sims=4000, sim_seed=123)
pw   = M.predict(**common, ci_method="pointwise")
bonf = M.predict(**common, ci_method="bonferroni")
sidk = M.predict(**common, ci_method="sidak")
supt = M.predict(**common, ci_method="sup-t")

widths = pd.DataFrame({
    "age":        ages,
    "pointwise":  pw.ci_upper   - pw.ci_lower,
    "bonferroni": bonf.ci_upper - bonf.ci_lower,
    "sidak":      sidk.ci_upper - sidk.ci_lower,
    "sup-t":      supt.ci_upper - supt.ci_lower,
}).set_index("age")
print("CI widths:")
print(widths)
print("\nBonferroni >= Sidak (always); for correlated margins sup-t is "
      "typically narrower than both.")

# ---------------------------------------------------------------------------
# 15. User-supplied vcov
#     Drop in any (k, k) covariance matrix you trust — e.g. a sandwich
#     computed offline, a Bayesian posterior covariance, or the output
#     of a custom resampling scheme — and smmargins will sandwich it
#     through the Jacobian without recomputing anything else.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("15. User-supplied parameter covariance (vcov=)")
print("=" * 80)
V_default = fit.cov_params().to_numpy()
V_inflated = V_default * 1.5     # toy example: assume 50% wider sampling cov
M_v = Margins(fit, vcov=V_inflated)
ame_default = M.dydx("age").se[0]
ame_user    = M_v.dydx("age").se[0]
print(f"AME(age) SE — default cov_params() : {ame_default: .6f}")
print(f"AME(age) SE — vcov = 1.5 x default : {ame_user: .6f}   "
      f"(ratio {ame_user / ame_default:.3f}, expect ≈ sqrt(1.5)={np.sqrt(1.5):.3f})")


# ---------------------------------------------------------------------------
# 16. Prediction scales (`scale=`) and a custom Transform
#     The default response scale gives Λ(η) for logit. Switch to "linear"
#     to get the linear predictor (Stata 'xb'), to "or" for odds-ratio
#     scale, or pass a Transform for any user-defined λ.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("16. Prediction scales (scale=) and custom Transform")
print("=" * 80)
ame_resp   = M.dydx("age").estimate[0]
ame_linear = M.dydx("age", scale="linear").estimate[0]
ame_or     = M.dydx("age", scale="or").estimate[0]
print(f"AME(age) on response scale      : {ame_resp: .6f}   (probability units)")
print(f"AME(age) on linear scale        : {ame_linear: .6f}   (= beta_age, log-odds)")
print(f"AME(age) on odds-ratio scale    : {ame_or: .6f}")

# Custom Transform: square the linear predictor (silly but illustrative)
square = Transform(
    value=lambda e: e ** 2,
    grad=lambda e: 2 * e,
    hess=lambda e: np.full_like(e, 2.0),
    name="square",
)
ame_sq = M.dydx("age", scale=square).estimate[0]
print(f"AME(age) under square transform : {ame_sq: .6f}")

# ---------------------------------------------------------------------------
# 17. Subgroup AMEs (over=) and observation weights
#     over= partitions the sample and averages within each subgroup,
#     keeping the FULL joint covariance (so cross-subgroup tests are valid).
#     weights= threads through every average; pass at construction time.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("17. Subgroup AMEs (over=) and observation weights")
print("=" * 80)
print("AME(age) by educ:")
print(M.dydx("age", over="educ").summary())

w = rng.uniform(0.5, 2.0, N)
M_w = Margins(fit, weights=w)
ame_unw = M.dydx("age").estimate[0]
ame_w   = M_w.dydx("age").estimate[0]
print()
print(f"AME(age) — unweighted         : {ame_unw: .6f}")
print(f"AME(age) — sampling weighted  : {ame_w: .6f}")

# ---------------------------------------------------------------------------
# 18. Joint Wald test and pairwise comparisons
#     wald() tests linear restrictions on a result; pairwise() builds the
#     all-level-vs-level contrast matrix for a factor and applies whatever
#     ci_method= you ask for (sup-t needs simulation/bootstrap draws).
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("18. Joint Wald and pairwise comparisons")
print("=" * 80)
res = M.dydx(["age", "income", "female"])
joint = res.wald()
print(f"Joint H0: AME(age) = AME(income) = AME(female) = 0")
print(f"  chi^2 = {joint.stat:.4f},  df = {joint.df},  p = {joint.pvalue:.3g}")

res_educ = M.dydx("educ")
print()
print("Pairwise comparisons across educ levels (Bonferroni-adjusted CIs):")
print(res_educ.pairwise(by="educ", ci_method="bonferroni").summary())

# ---------------------------------------------------------------------------
# 19. Per-variable DSL (values=)
#     Fix specific variables without restating every column.
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("19. Per-variable DSL — values= and default_values=")
print("=" * 80)
print("Predictions at income=p25/p50/p75, everything else at mean:")
print(M.predict(values={"income": ["p25", "p50", "p75"]}, default_values="mean"))

print()
print("Counterfactual: income + 10% via Expr:")
from smmargins import Expr
print(M.predict(values={"income": Expr("income * 1.10")}))

# ---------------------------------------------------------------------------
# 20. Plotting
# ---------------------------------------------------------------------------
print()
print("=" * 80)
print("20. Plotting — prediction curve vs age")
print("=" * 80)
from smmargins import plot_predictions
fig, ax = plot_predictions(M, "age")
fig.savefig("demo_plot_predictions.png", dpi=150)
print("Saved demo_plot_predictions.png")

print()
print("Plotting — AME of age vs income, by sex:")
from smmargins import plot_slopes
fig2, ax2 = plot_slopes(M, "age", condition="income", by="female")
fig2.savefig("demo_plot_slopes.png", dpi=150)
print("Saved demo_plot_slopes.png")

Healthcare-style 2x2 difference-in-differences

demo_did.py answers a clinical question:

Is there a rate difference of condition \(X\) between groups A and B, with or without a preexisting condition \(Y\)?

The script fits a logit on simulated patient data and reports, on the probability scale:

• 4 cell predictions \(P(X \mid \text{group}, Y)\)

• 2 simple effects \(P(X \mid B, Y) - P(X \mid A, Y)\) at each \(Y\)

• 1 difference-in-differences (whether the A-vs-B gap depends on \(Y\))

All with delta-method standard errors and confidence intervals. The DiD here is not the coefficient on the group:Y interaction — that coefficient is on the log-odds scale, while the clinical question is about probabilities. This is Ai & Norton (2003) in practice; see Mathematical motivation for the derivation.

Highlights from the script

Fit and call did():

fit = smf.logit(
    "condition_X ~ C(group) + preexist_Y + C(group):preexist_Y "
    "+ age + female",
    data=df,
).fit()
M = Margins(fit)

did = M.did("group", "preexist_Y",
            group_levels=["A", "B"],
            condition_levels=[0, 1])
print(did)              # cells + simple effects + DiD

Same DiD at one specific patient profile (60-year-old male):

M.did("group", "preexist_Y",
      group_levels=["A", "B"], condition_levels=[0, 1],
      atexog={"age": 60, "female": 0})

Plot-ready cell table:

tbl = did.cells.summary()        # estimate / SE / CI per cell

Full source

"""
demo_did.py
===========

Difference-in-differences example directly matching the question:

    "Is there a rate difference of condition X between group A and B,
     with or without preexisting condition Y?"

We fit a logit for P(X=1) on group (A/B), preexisting Y (0/1), their
interaction, and control covariates. Then we use Margins.did() to get,
on the *probability* scale:

  * 4 cell predictions      P(X | group, Y)
  * 2 simple effects        P(X|B,Y) - P(X|A,Y)   at each Y
  * 1 DiD                   (simple effect at Y=1) - (simple effect at Y=0)

All with delta-method standard errors and CIs.

The DiD here is the "does the A-vs-B gap depend on Y?" question.  It is
NOT the coefficient on group×Y (that's on the log-odds scale); on the
probability scale you have to go through the inverse link — which is
exactly what Margins.did() does.
"""
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

from smmargins import Margins

pd.options.display.width = 140
pd.options.display.float_format = "{: .4f}".format

# ---------------------------------------------------------------------------
# Simulate patient-level data
# ---------------------------------------------------------------------------
rng = np.random.default_rng(42)
N = 6_000
df = pd.DataFrame({
    "group":        rng.choice(["A", "B"], N, p=[0.55, 0.45]),
    "preexist_Y":   rng.integers(0, 2, N),             # 0 = no Y, 1 = has Y
    "age":          rng.normal(55, 15, N).clip(18, 95),
    "female":       rng.integers(0, 2, N),
})

# True data-generating process:
#   * baseline rate of X depends on age, sex, and Y
#   * group B has a modest additive bump in the log-odds
#   * the group effect is AMPLIFIED among patients with preexisting Y
#     (this is the thing we want to detect)
eta = (
    -3.5
    + 0.04 * df["age"]
    - 0.3 * df["female"]
    + 0.5 * (df["group"] == "B")
    + 1.1 * df["preexist_Y"]
    + 0.8 * (df["group"] == "B") * df["preexist_Y"]   # interaction
)
df["condition_X"] = (rng.uniform(0, 1, N) < 1 / (1 + np.exp(-eta))).astype(int)

print("Raw sample rates of condition X by cell:")
print(df.groupby(["group", "preexist_Y"])["condition_X"].mean().round(4))
print()

# ---------------------------------------------------------------------------
# Fit the logit with the group × preexist_Y interaction + controls
# ---------------------------------------------------------------------------
fit = smf.logit(
    "condition_X ~ C(group) + preexist_Y + C(group):preexist_Y + age + female",
    data=df,
).fit(disp=False)

print("=" * 84)
print("Logit model (coefficients are on the log-odds scale)")
print("=" * 84)
print(fit.summary().tables[1])
print()

# ---------------------------------------------------------------------------
# DiD on the *probability* (response) scale — what the clinical question asks
# ---------------------------------------------------------------------------
# Margins(fit) uses analytic outer Jacobians via family.link.inverse_deriv
# when available (Logit qualifies), falling back to central finite
# differences otherwise. did() reuses predict()'s machinery, so it
# inherits the analytic path automatically. Set Margins(fit,
# analytic=False) to force FD if you ever want to cross-check.
M = Margins(fit)
did = M.did("group", "preexist_Y",
            group_levels=["A", "B"],
            condition_levels=[0, 1])

print("=" * 84)
print("DiD on the probability scale, averaged over age and sex distribution")
print("=" * 84)
print(did)

# ---------------------------------------------------------------------------
# Interpretation
# ---------------------------------------------------------------------------
pA0 = did.cells.estimate[0]   # group=A, Y=0
pA1 = did.cells.estimate[1]   # group=A, Y=1
pB0 = did.cells.estimate[2]   # group=B, Y=0
pB1 = did.cells.estimate[3]   # group=B, Y=1
se_simple_Y0 = did.simple_effects.se[0]
se_simple_Y1 = did.simple_effects.se[1]
did_est, did_se = did.did.estimate[0], did.did.se[0]

print()
print("=" * 84)
print("Plain-language summary")
print("=" * 84)
print(f"Condition X rate, group A, no preexisting Y : {pA0:.3%}")
print(f"Condition X rate, group A, with Y           : {pA1:.3%}")
print(f"Condition X rate, group B, no preexisting Y : {pB0:.3%}")
print(f"Condition X rate, group B, with Y           : {pB1:.3%}")
print()
print(f"Rate difference (B - A) among NO-Y patients  : "
      f"{(pB0 - pA0):+.3%}  (SE {se_simple_Y0:.3%})")
print(f"Rate difference (B - A) among WITH-Y patients: "
      f"{(pB1 - pA1):+.3%}  (SE {se_simple_Y1:.3%})")
print()
print(f"Difference-in-differences                   : "
      f"{did_est:+.3%}  (SE {did_se:.3%})")
print(f"  -> the B-vs-A gap is {abs(did_est):.3%} larger among patients "
      f"with preexisting Y.")
print(f"  -> 95% CI: ({did.did.ci_lower[0]:+.3%}, {did.did.ci_upper[0]:+.3%})")
print(f"  -> p-value: {did.did.pvalues[0]:.4g}")

# ---------------------------------------------------------------------------
# Sensitivity: DiD at a specific patient profile (e.g. 60-year-old male)
# ---------------------------------------------------------------------------
print()
print("=" * 84)
print("DiD at a specific profile: 60-year-old male")
print("=" * 84)
did_profile = M.did(
    "group", "preexist_Y",
    group_levels=["A", "B"], condition_levels=[0, 1],
    atexog={"age": 60, "female": 0},
)
print(did_profile.did)

# ---------------------------------------------------------------------------
# Bonus: plottable table of cell predictions with CIs
# ---------------------------------------------------------------------------
print()
print("=" * 84)
print("Cells with 95% CIs (suitable for a plot)")
print("=" * 84)
tbl = did.cells.summary().copy()
print(tbl)

# If you wanted to plot:
#   import matplotlib.pyplot as plt
#   fig, ax = plt.subplots()
#   for g in ["A", "B"]:
#       sub = tbl[tbl.index.str.contains(f"group={g}")]
#       ax.errorbar([0, 1],
#                   sub["prediction"].values,
#                   yerr=(sub["prediction"] - sub["[95% CI lo]"]).values,
#                   marker="o", label=f"group {g}", capsize=4)
#   ax.set_xticks([0, 1]); ax.set_xticklabels(["no Y", "with Y"])
#   ax.set_ylabel("P(condition X)"); ax.legend(); plt.show()