pyfector.plotting

  1"""
  2Plotting functions for pyfector results.
  3
  4Produces matplotlib/plotly figures equivalent to R fect's plot types:
  5- gap: dynamic treatment effects with CI bands
  6- status: treatment status heatmap
  7- factors: latent factor time series
  8- counterfactual: actual vs counterfactual for selected units
  9- equiv: equivalence test visualization
 10"""
 11
 12from __future__ import annotations
 13
 14from typing import Literal
 15
 16import numpy as np
 17
 18
 19def plot(
 20    result,
 21    kind: Literal["gap", "status", "factors", "counterfactual", "equiv", "calendar"] = "gap",
 22    units: list | None = None,
 23    show_ci: bool = True,
 24    title: str | None = None,
 25    figsize: tuple[float, float] = (10, 6),
 26    ax=None,
 27    **kwargs,
 28):
 29    """Plot fect results.
 30
 31    Parameters
 32    ----------
 33    result : FectResult
 34        Output from ``pyfector.fect()``.
 35    kind : str
 36        Plot type.
 37    units : list, optional
 38        Unit IDs for counterfactual plot.
 39    show_ci : bool
 40        Show confidence intervals (requires ``se=True`` in estimation).
 41    """
 42    import matplotlib.pyplot as plt
 43
 44    if ax is None:
 45        fig, ax = plt.subplots(figsize=figsize)
 46    else:
 47        fig = ax.get_figure()
 48
 49    if kind == "gap":
 50        _plot_gap(result, ax, show_ci, **kwargs)
 51    elif kind == "status":
 52        _plot_status(result, ax, **kwargs)
 53    elif kind == "factors":
 54        _plot_factors(result, ax, **kwargs)
 55    elif kind == "counterfactual":
 56        _plot_counterfactual(result, ax, units, **kwargs)
 57    elif kind == "equiv":
 58        _plot_equiv(result, ax, **kwargs)
 59    elif kind == "calendar":
 60        _plot_calendar(result, ax, show_ci, **kwargs)
 61    else:
 62        raise ValueError(f"Unknown plot kind: {kind}")
 63
 64    if title:
 65        ax.set_title(title)
 66
 67    fig.tight_layout()
 68    return fig
 69
 70
 71def _plot_gap(result, ax, show_ci, count_bar: bool = True, **kwargs):
 72    """Dynamic treatment effects (gap) plot."""
 73    time_on = result.time_on
 74    att_on = result.att_on
 75
 76    if time_on is None or att_on is None:
 77        ax.text(0.5, 0.5, "No dynamic effects available", ha="center", va="center",
 78                transform=ax.transAxes)
 79        return
 80
 81    # Main line
 82    ax.plot(time_on, att_on, "o-", color="#2166ac", markersize=4, linewidth=1.5,
 83            label="ATT", zorder=3)
 84
 85    # CI bands
 86    if show_ci and result.inference is not None:
 87        inf = result.inference
 88        ax.fill_between(
 89            time_on, inf.att_on_ci_lower, inf.att_on_ci_upper,
 90            alpha=0.2, color="#2166ac", label="95% CI",
 91        )
 92
 93    # Zero line
 94    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--", zorder=1)
 95
 96    # Treatment onset line
 97    ax.axvline(0, color="#b2182b", linewidth=1, linestyle=":", alpha=0.7,
 98               label="Treatment onset", zorder=2)
 99
100    # Count bars at bottom
101    if count_bar and result.count_on is not None:
102        ax2 = ax.twinx()
103        ax2.bar(time_on, result.count_on, alpha=0.15, color="gray", width=0.8)
104        ax2.set_ylabel("Count", color="gray", fontsize=9)
105        ax2.tick_params(axis="y", labelcolor="gray", labelsize=8)
106        max_count = max(result.count_on) if len(result.count_on) > 0 else 1
107        ax2.set_ylim(0, max_count * 4)
108
109    ax.set_xlabel("Relative time to treatment")
110    ax.set_ylabel("Treatment effect")
111    ax.legend(loc="best", fontsize=9)
112
113
114def _plot_status(result, ax, **kwargs):
115    """Treatment status heatmap."""
116    import matplotlib.colors as mcolors
117
118    panel = result.panel
119    if panel is None:
120        return
121
122    status = panel.obs_status
123
124    # Custom colormap: 1=treated(red), 2=control(blue), 3=missing(gray), 4=removed(white)
125    cmap = mcolors.ListedColormap(["#b2182b", "#2166ac", "#d9d9d9", "#ffffff"])
126    bounds = [0.5, 1.5, 2.5, 3.5, 4.5]
127    norm = mcolors.BoundaryNorm(bounds, cmap.N)
128
129    # Sort units: controls first, then treated by treatment timing
130    order = np.argsort(panel.unit_type)
131
132    im = ax.imshow(
133        status[:, order].T, aspect="auto", cmap=cmap, norm=norm,
134        interpolation="nearest",
135    )
136
137    ax.set_xlabel("Time period")
138    ax.set_ylabel("Unit")
139
140    # Legend
141    from matplotlib.patches import Patch
142    legend_elements = [
143        Patch(facecolor="#b2182b", label="Treated"),
144        Patch(facecolor="#2166ac", label="Control"),
145        Patch(facecolor="#d9d9d9", label="Missing"),
146    ]
147    ax.legend(handles=legend_elements, loc="upper right", fontsize=8)
148
149
150def _plot_factors(result, ax, **kwargs):
151    """Plot estimated latent factors over time."""
152    if result.factors is None:
153        ax.text(0.5, 0.5, "No factors estimated", ha="center", va="center",
154                transform=ax.transAxes)
155        return
156
157    T, r = result.factors.shape
158    time_ids = result.panel.time_ids if result.panel is not None else np.arange(T)
159
160    for k in range(r):
161        ax.plot(time_ids, result.factors[:, k], label=f"Factor {k+1}", linewidth=1.2)
162
163    ax.set_xlabel("Time")
164    ax.set_ylabel("Factor value")
165    ax.legend(fontsize=9)
166
167
168def _plot_counterfactual(result, ax, units=None, **kwargs):
169    """Actual vs counterfactual outcomes for selected units."""
170    panel = result.panel
171    if panel is None or result.Y_ct is None:
172        return
173
174    if units is None:
175        # Pick first 3 treated units
176        treated_idx = np.where(panel.unit_type == 2)[0][:3]
177    else:
178        treated_idx = [np.where(panel.unit_ids == u)[0][0] for u in units
179                       if u in panel.unit_ids]
180
181    time_ids = panel.time_ids
182    colors = ["#2166ac", "#b2182b", "#4daf4a", "#984ea3", "#ff7f00"]
183
184    for i, j in enumerate(treated_idx):
185        obs_mask = panel.I[:, j] > 0
186        color = colors[i % len(colors)]
187        uid = panel.unit_ids[j]
188
189        ax.plot(time_ids[obs_mask], panel.Y[obs_mask, j], "o-",
190                color=color, markersize=3, linewidth=1,
191                label=f"Unit {uid} (actual)")
192        ax.plot(time_ids, result.Y_ct[:, j], "--",
193                color=color, linewidth=1, alpha=0.7,
194                label=f"Unit {uid} (counterfactual)")
195
196    ax.set_xlabel("Time")
197    ax.set_ylabel("Outcome")
198    ax.legend(fontsize=8)
199
200
201def _plot_equiv(result, ax, threshold: float = 0.36, **kwargs):
202    """Equivalence test visualization."""
203    if result.att_on is None or result.inference is None:
204        ax.text(0.5, 0.5, "Inference required", ha="center", va="center",
205                transform=ax.transAxes)
206        return
207
208    inf = result.inference
209    # Only pre-treatment periods
210    pre_mask = result.time_on < 0
211    if not np.any(pre_mask):
212        return
213
214    t_pre = result.time_on[pre_mask]
215    att_pre = result.att_on[pre_mask]
216    ci_lo = inf.att_on_ci_lower[pre_mask]
217    ci_hi = inf.att_on_ci_upper[pre_mask]
218
219    ax.errorbar(t_pre, att_pre, yerr=[att_pre - ci_lo, ci_hi - att_pre],
220                fmt="o", color="#2166ac", capsize=3, markersize=4)
221
222    # Equivalence bounds
223    ax.axhspan(-threshold, threshold, alpha=0.1, color="green",
224               label=f"Equiv. bounds (+-{threshold})")
225    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--")
226
227    ax.set_xlabel("Relative time (pre-treatment)")
228    ax.set_ylabel("ATT")
229    ax.legend(fontsize=9)
230
231
232def _plot_calendar(result, ax, show_ci, **kwargs):
233    """Calendar-time ATT plot."""
234    if result.eff is None or result.panel is None:
235        return
236
237    panel = result.panel
238    D = panel.D
239    eff = result.eff
240
241    # ATT by calendar time
242    time_ids = panel.time_ids
243    cal_att = []
244    for t in range(panel.T):
245        treated_t = D[t, :] > 0
246        if np.any(treated_t):
247            cal_att.append(float(np.mean(eff[t, treated_t])))
248        else:
249            cal_att.append(np.nan)
250
251    cal_att = np.array(cal_att)
252    valid = ~np.isnan(cal_att)
253
254    ax.plot(time_ids[valid], cal_att[valid], "o-", color="#2166ac",
255            markersize=3, linewidth=1)
256    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--")
257    ax.set_xlabel("Calendar time")
258    ax.set_ylabel("ATT")