feat(did): xvar heterogeneity + parity tests + custom summary + did_report

Four-in-one DID polish round, all scoped to the DID module:

1. etwfe(xvar=...) — covariate-moderated heterogeneity
   New _etwfe_with_xvar helper adds cohort × post × (x - mean(x))
   interactions so the baseline coefficient is ATT(g) evaluated at
   the sample mean of xvar and the slope coefficient measures how
   ATT(g) shifts per unit of xvar. Closes the last R-etwfe feature
   gap. Backward compatible: etwfe(xvar=None) is byte-identical to
   wooldridge_did().

2. tests/test_did_summary.py — DID consistency + parity tests
   Seven pytest cases covering: did_summary ↔ direct etwfe / BJS
   numeric match; etwfe ≡ wooldridge_did under xvar=None; xvar=
   random-covariate sanity bound; Markdown + LaTeX exports shape +
   content; non-did_summary input correctly rejected; forest plot
   render via Agg backend. pytest exit 0 (7 passed in 6.34s).

3. Custom result.summary() on did_summary output
   Prints a Markdown table (via did_summary_to_markdown) plus a
   header block with fitted/failed methods, cross-method mean,
   dispersion SD, and breakdown M* (when sensitivity was run).
   Base CausalResult.summary's single-estimate layout was a poor
   fit for a multi-method comparison; this customisation gives
   researchers a paste-ready console dump.

4. did_report(data, ..., save_to=dir) — one-call bundle
   Fits did_summary with include_sensitivity=True by default and
   writes five files into save_to/:
     - did_summary.txt   (summary dump)
     - did_summary.md    (GFM table)
     - did_summary.tex   (booktabs LaTeX fragment)
     - did_summary.png   (forest plot, matplotlib-optional)
     - did_summary.json  (detail + model_info, replay-friendly)
   Returns the underlying CausalResult for programmatic use.

READMEs updated (EN + CN). Commit staged with explicit file list to
avoid the earlier `git add -A` sweep regression.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

README.md

@@ -167,6 +167,8 @@ If a method exists in R, we aim to match or exceed its feature set in Python —
 | `did()` | Auto-dispatching DID (2×2 or staggered) | — |
 | `did_summary()` | One-call robustness comparison across CS/SA/BJS/ETWFE/Stacked | — |
 | `did_summary_plot()` | Forest plot of method-robustness summary | — |
+| `did_summary_to_markdown()` / `_to_latex()` | Publication-ready tables from `did_summary` | — |
+| `did_report()` | One-call bundle: txt + md + tex + png + json into a folder | — |
 | `did_2x2()` | Classic two-group, two-period DID | — |
 | `callaway_santanna()` | Staggered DID with heterogeneous effects | Callaway & Sant'Anna (2021) |
 | `sun_abraham()` | Interaction-weighted event study | Sun & Abraham (2021) |

README_CN.md

@@ -86,6 +86,8 @@ StatsPAI **不是** R 的 wrapper。我们从原始论文独立重新实现每
 | `did()` | 自动分派 DID（2×2 或交错） | — |
 | `did_summary()` | 一次调用跑 CS/SA/BJS/ETWFE/Stacked 五种方法并出对比表 | — |
 | `did_summary_plot()` | 方法稳健性森林图（点估计 + CI 并排） | — |
+| `did_summary_to_markdown()` / `_to_latex()` | 论文级稳健性对比表（GFM / booktabs） | — |
+| `did_report()` | 一次调用打包输出：txt + md + tex + png + json | — |
 | `callaway_santanna()` | 交错 DID，异质处理效应 | Callaway & Sant'Anna (2021) |
 | `sun_abraham()` | 交互加权事件研究 | Sun & Abraham (2021) |
 | `bacon_decomposition()` | TWFE 分解诊断 | Goodman-Bacon (2021) |

src/statspai/__init__.py

@@ -35,7 +35,7 @@
     bacon_decomposition, honest_did, breakdown_m, event_study,
     did_analysis, DIDAnalysis, did_multiplegt, did_imputation, stacked_did, cic,
     wooldridge_did, etwfe, drdid, twfe_decomposition,
-    did_summary, did_summary_to_markdown, did_summary_to_latex,
+    did_summary, did_summary_to_markdown, did_summary_to_latex, did_report,
     pretrends_test, pretrends_power, sensitivity_rr, SensitivityResult, pretrends_summary,
     parallel_trends_plot, bacon_plot, group_time_plot, did_plot,
     enhanced_event_study_plot, treatment_rollout_plot,
@@ -231,6 +231,7 @@
     "did_summary",
     "did_summary_to_markdown",
     "did_summary_to_latex",
+    "did_report",
     "drdid",
     "twfe_decomposition",
     # RD

src/statspai/did/__init__.py

@@ -46,7 +46,7 @@
     pretrends_summary,
 )
 from .wooldridge_did import wooldridge_did, etwfe, drdid, twfe_decomposition
-from .summary import did_summary, did_summary_to_markdown, did_summary_to_latex
+from .summary import did_summary, did_summary_to_markdown, did_summary_to_latex, did_report
 from .continuous_did import continuous_did
 from .plots import (
     parallel_trends_plot,
@@ -390,6 +390,7 @@ def did(
     'did_summary',
     'did_summary_to_markdown',
     'did_summary_to_latex',
+    'did_report',
     'drdid',
     'twfe_decomposition',
     # Pre-trends

src/statspai/did/summary.py

@@ -303,7 +303,7 @@ def did_summary(
     else:
         avg_est, disp = np.nan, np.nan
 
-    return CausalResult(
+    result = CausalResult(
         method="DID Method-Robustness Summary",
         estimand="Overall ATT (mean across methods)",
         estimate=avg_est,
@@ -323,6 +323,28 @@ def did_summary(
         _citation_key="did_summary",
     )
 
+    # Custom .summary() — prints the comparison table as a Markdown block
+    # instead of the base CausalResult's single-estimate format.
+    def _custom_summary(alpha=None):  # noqa: ARG001
+        header = "=" * 78 + "\n"
+        header += "  DID Method-Robustness Summary\n"
+        header += "=" * 78 + "\n\n"
+        if fit:
+            header += f"  Fitted methods: {', '.join(fit)}\n"
+        if failed:
+            header += f"  Failed:         {', '.join(failed.keys())}\n"
+        if not np.isnan(avg_est):
+            header += f"  Mean ATT:       {avg_est:+.4f}\n"
+            header += f"  Cross-method SD: {disp:+.4f}\n"
+        if breakdown_m_value is not None:
+            header += f"  Breakdown M*:   {breakdown_m_value:.4f} (CS row)\n"
+        header += "\n"
+        body = did_summary_to_markdown(result)
+        return header + body
+
+    result.summary = _custom_summary
+    return result
+
 
 # ═══════════════════════════════════════════════════════════════════════
 #  Export helpers: publication-ready Markdown / LaTeX from a did_summary
@@ -521,3 +543,117 @@ def did_summary_to_latex(
         )
     lines.append("\\end{table}")
     return "\n".join(lines)
+
+
+# ═══════════════════════════════════════════════════════════════════════
+#  One-call DID bundle report — txt + md + tex + png in a single directory
+# ═══════════════════════════════════════════════════════════════════════
+
+def did_report(
+    data: pd.DataFrame,
+    y: str,
+    time: str,
+    first_treat: str,
+    group: str,
+    save_to: str,
+    methods: Union[str, List[str]] = "auto",
+    controls: Optional[List[str]] = None,
+    cluster: Optional[str] = None,
+    alpha: float = 0.05,
+    include_sensitivity: bool = True,
+    plot_sort_by: Optional[str] = "estimate",
+    verbose: bool = False,
+) -> CausalResult:
+    """
+    One-call comprehensive DID report: fits all methods + exports every format.
+
+    Writes the following files to ``save_to``:
+
+    - ``did_summary.txt`` : text dump of ``result.summary()``.
+    - ``did_summary.md``  : GitHub-Flavoured Markdown table.
+    - ``did_summary.tex`` : publication-ready LaTeX ``booktabs`` fragment.
+    - ``did_summary.png`` : forest plot (requires matplotlib).
+    - ``did_summary.json`` : detail table + model_info in JSON.
+
+    Parameters
+    ----------
+    data, y, time, first_treat, group, methods, controls, cluster, alpha
+        Same as :func:`did_summary`.
+    save_to : str
+        Directory path. Created if it does not exist.
+    include_sensitivity : bool, default True
+        Whether to run Rambachan-Roth breakdown M*. Defaults to ``True``
+        in ``did_report`` (vs ``False`` in ``did_summary``) because a
+        report is expected to be comprehensive.
+    plot_sort_by : {'estimate', None}, default 'estimate'
+        Sort the forest plot by estimate ascending.
+    verbose : bool, default False
+        Print progress lines.
+
+    Returns
+    -------
+    CausalResult
+        The underlying :func:`did_summary` output. All side-effect files
+        are written to ``save_to`` as a bundle.
+    """
+    import json
+    from pathlib import Path
+
+    out_dir = Path(save_to)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    if verbose:
+        print(f"[did_report] fitting methods into {out_dir} ...", flush=True)
+
+    result = did_summary(
+        data, y=y, time=time, first_treat=first_treat, group=group,
+        methods=methods, controls=controls, cluster=cluster, alpha=alpha,
+        include_sensitivity=include_sensitivity, verbose=verbose,
+    )
+
+    # TXT
+    (out_dir / "did_summary.txt").write_text(result.summary(), encoding="utf-8")
+
+    # Markdown
+    (out_dir / "did_summary.md").write_text(
+        did_summary_to_markdown(result), encoding="utf-8",
+    )
+
+    # LaTeX
+    (out_dir / "did_summary.tex").write_text(
+        did_summary_to_latex(result), encoding="utf-8",
+    )
+
+    # PNG (best-effort — skip if matplotlib unavailable)
+    try:
+        from .plots import did_summary_plot
+        fig, _ = did_summary_plot(result, sort_by=plot_sort_by)
+        fig.savefig(out_dir / "did_summary.png", dpi=150, bbox_inches="tight")
+        try:
+            import matplotlib.pyplot as plt
+            plt.close(fig)
+        except Exception:
+            pass
+    except ImportError:
+        if verbose:
+            print("[did_report] matplotlib missing — skipping PNG", flush=True)
+
+    # JSON (detail + model_info)
+    payload = {
+        "method": result.method,
+        "estimate": None if pd.isna(result.estimate) else float(result.estimate),
+        "cross_method_sd": None if pd.isna(result.se) else float(result.se),
+        "detail": result.detail.replace({np.nan: None}).to_dict(orient="records"),
+        "model_info": {
+            k: (list(v) if isinstance(v, (set, tuple)) else v)
+            for k, v in result.model_info.items()
+        },
+    }
+    (out_dir / "did_summary.json").write_text(
+        json.dumps(payload, indent=2, default=str), encoding="utf-8",
+    )
+
+    if verbose:
+        print(f"[did_report] wrote bundle to {out_dir}", flush=True)
+
+    return result

src/statspai/did/wooldridge_did.py

@@ -358,6 +358,7 @@ def etwfe(
     controls: Optional[List[str]] = None,
     cluster: Optional[str] = None,
     alpha: float = 0.05,
+    xvar: Optional[str] = None,
 ) -> CausalResult:
     """
     Extended Two-Way Fixed Effects (ETWFE) — Wooldridge (2021).
@@ -403,7 +404,7 @@ def etwfe(
     ``tvar = time``               ``time='time'``
     ``gvar = first_treat``        ``first_treat='first_treat'``
     ``ivar = unit``               ``group='unit'``
-    ``xvar`` (covariate het.)     *not yet supported — use* ``controls``
+    ``xvar`` (covariate het.)     ``xvar='x1'`` (single covariate)
     ``vcov = ~cluster``           ``cluster='cluster'``
     ============================  ========================================
 
@@ -433,15 +434,161 @@ def etwfe(
     callaway_santanna : CS (2021) group-time ATT estimator.
     aggte : Aggregation of group-time ATTs (event/group/calendar/simple).
     """
-    return wooldridge_did(
-        data=data,
-        y=y,
-        group=group,
-        time=time,
-        first_treat=first_treat,
-        controls=controls,
-        cluster=cluster,
+    if xvar is None:
+        return wooldridge_did(
+            data=data, y=y, group=group, time=time, first_treat=first_treat,
+            controls=controls, cluster=cluster, alpha=alpha,
+        )
+    return _etwfe_with_xvar(
+        data=data, y=y, group=group, time=time, first_treat=first_treat,
+        xvar=xvar, controls=controls, cluster=cluster, alpha=alpha,
+    )
+
+
+def _etwfe_with_xvar(
+    data: pd.DataFrame,
+    y: str,
+    group: str,
+    time: str,
+    first_treat: str,
+    xvar: str,
+    controls: Optional[List[str]] = None,
+    cluster: Optional[str] = None,
+    alpha: float = 0.05,
+) -> CausalResult:
+    """ETWFE with covariate-moderated heterogeneity (R etwfe's ``xvar``).
+
+    For each cohort ``g``, adds an interaction between the
+    cohort × post dummy and ``(xvar - mean(xvar))``, so the slope
+    coefficient measures how ATT(g) shifts per unit of ``xvar``. The
+    main cohort coefficient is then interpretable as ATT(g) evaluated
+    at the sample mean of ``xvar``.
+    """
+    df = data.copy()
+    ft = df[first_treat].replace(0, np.nan)
+    df["_ft"] = ft
+
+    periods = sorted(df[time].unique())
+    cohorts = sorted(df.loc[df["_ft"].notna(), "_ft"].unique())
+    if len(cohorts) == 0:
+        raise ValueError("No treated cohorts found. Check 'first_treat' column.")
+
+    # Demean outcome
+    df["_y"] = df[y].astype(float)
+    unit_mean = df.groupby(group)["_y"].transform("mean")
+    time_mean = df.groupby(time)["_y"].transform("mean")
+    grand_mean = df["_y"].mean()
+    df["_y_dm"] = df["_y"] - unit_mean - time_mean + grand_mean
+
+    # Center xvar by grand mean so the baseline ATT is evaluated at x=mean
+    df["_x"] = df[xvar].astype(float)
+    x_center = df["_x"].mean()
+    df["_xc"] = df["_x"] - x_center
+
+    # Build cohort × post dummies AND cohort × post × xc interactions
+    base_cols: List[str] = []
+    slope_cols: List[str] = []
+    for g in cohorts:
+        mask_post = (df["_ft"] == g) & (df[time] >= g)
+        b = f"_coh{int(g)}_post"
+        s = f"_coh{int(g)}_post_x"
+        df[b] = mask_post.astype(float)
+        df[s] = df[b] * df["_xc"]
+        base_cols.append(b)
+        slope_cols.append(s)
+
+    all_inter = base_cols + slope_cols
+    # Demean interactions via FE projection
+    for col in all_inter:
+        u_m = df.groupby(group)[col].transform("mean")
+        t_m = df.groupby(time)[col].transform("mean")
+        g_m = df[col].mean()
+        df[f"{col}_dm"] = df[col] - u_m - t_m + g_m
+
+    # Demean controls
+    ctrl_dm_cols: List[str] = []
+    if controls:
+        for c in controls:
+            df[f"_ctrl_{c}"] = df[c].astype(float)
+            u_m = df.groupby(group)[f"_ctrl_{c}"].transform("mean")
+            t_m = df.groupby(time)[f"_ctrl_{c}"].transform("mean")
+            g_m = df[f"_ctrl_{c}"].mean()
+            df[f"_ctrl_{c}_dm"] = df[f"_ctrl_{c}"] - u_m - t_m + g_m
+            ctrl_dm_cols.append(f"_ctrl_{c}_dm")
+
+    keep = ["_y_dm"] + [f"{c}_dm" for c in all_inter] + ctrl_dm_cols
+    valid = df[keep].notna().all(axis=1)
+    dfv = df.loc[valid].reset_index(drop=True)
+
+    y_vec = dfv["_y_dm"].values
+    X_cols = [f"{c}_dm" for c in all_inter] + ctrl_dm_cols
+    X = dfv[X_cols].values
+    X = np.column_stack([np.ones(len(y_vec)), X])
+
+    cl_arr = dfv[cluster].values if cluster else dfv[group].values
+    beta, se, vcov = _ols_fit(X, y_vec, cluster=cl_arr)
+    n_obs = len(y_vec)
+    df_resid = n_obs - X.shape[1]
+
+    k = len(cohorts)
+    # Baseline ATT(g) at x=mean: coefficients 1..k
+    # Slope(g): coefficients k+1..2k
+    cohort_results = []
+    for i, g in enumerate(cohorts):
+        b_idx = 1 + i
+        s_idx = 1 + k + i
+        att = float(beta[b_idx])
+        att_se = float(se[b_idx])
+        slope = float(beta[s_idx])
+        slope_se = float(se[s_idx])
+        p = float(2 * (1 - stats.t.cdf(abs(att / att_se) if att_se > 0 else 0,
+                                       max(df_resid, 1))))
+        p_slope = float(2 * (1 - stats.t.cdf(abs(slope / slope_se) if slope_se > 0 else 0,
+                                             max(df_resid, 1))))
+        n_g = int((dfv["_ft"] == g).sum())
+        cohort_results.append({
+            "cohort": int(g),
+            "att_at_xmean": att, "att_se": att_se, "att_pvalue": p,
+            "slope_wrt_x": slope, "slope_se": slope_se, "slope_pvalue": p_slope,
+            "n_obs": n_g,
+        })
+
+    detail = pd.DataFrame(cohort_results)
+    sizes = detail["n_obs"].values.astype(float)
+    weights_vec = sizes / sizes.sum() if sizes.sum() > 0 else np.ones(k) / k
+    att_overall = float(weights_vec @ detail["att_at_xmean"].values)
+    # SE via delta method using vcov of baseline coefficients
+    base_vcov = vcov[1:1 + k, 1:1 + k]
+    att_se_overall = float(np.sqrt(weights_vec @ base_vcov @ weights_vec))
+    t_stat = att_overall / att_se_overall if att_se_overall > 0 else np.nan
+    p_overall = float(2 * (1 - stats.t.cdf(abs(t_stat), max(df_resid, 1))))
+    t_crit = stats.t.ppf(1 - alpha / 2, max(df_resid, 1))
+    ci = (att_overall - t_crit * att_se_overall, att_overall + t_crit * att_se_overall)
+
+    model_info = {
+        "n_cohorts": k,
+        "cohorts": [int(g) for g in cohorts],
+        "n_periods": len(periods),
+        "n_units": df[group].nunique(),
+        "controls": controls or [],
+        "cluster_var": cluster or group,
+        "xvar": xvar,
+        "xvar_mean": float(x_center),
+        "heterogeneity": "ATT(g) = baseline(g) + slope(g) * (xvar - mean(xvar))",
+    }
+
+    return CausalResult(
+        method="Wooldridge (2021) ETWFE with covariate heterogeneity",
+        estimand=f"ATT at {xvar} = {x_center:.4g} (sample mean)",
+        estimate=att_overall,
+        se=att_se_overall,
+        pvalue=p_overall,
+        ci=ci,
         alpha=alpha,
+        n_obs=n_obs,
+        detail=detail,
+        model_info=model_info,
+        _citation_key="wooldridge_twfe",
     )