pyfector.fect

  1"""
  2Public entry point for pyfector.
  3
  4This module exposes :func:`fect`, the single function most users call,
  5and the :class:`FectResult` dataclass that packages its output.
  6
  7:func:`fect` orchestrates the full pipeline:
  8
  91. :func:`pyfector.panel.prepare_panel` converts the input DataFrame to
 10   dense ``(T, N)`` matrices with treatment timing and unit
 11   classification.
 122. :func:`pyfector.panel.initial_fit` computes the starting fit ``Y0``
 13   and initial covariate coefficients on the control subsample.
 143. If the user passes a range for ``r`` (IFE) or leaves ``lam=None``
 15   (MC), :mod:`pyfector.cv` chooses the hyperparameter by
 16   cross-validation.
 174. :mod:`pyfector.estimators` iterates the EM loop for the chosen
 18   method (``fe``, ``ife``, ``mc``, or ``cfe``).
 195. Treatment effects are computed from ``eff = Y - Y_ct``; the overall
 20   ATT and dynamic ATT by relative event time are derived in
 21   :func:`_compute_effects` via :func:`numpy.bincount`.  Raw missing
 22   outcomes remain excluded from effect averages; model imputations are
 23   counterfactual predictions, not substitutes for unobserved treated
 24   outcomes.
 256. When ``se=True`` the requested inference routine
 26   (``bootstrap`` or ``jackknife``) from :mod:`pyfector.inference` is
 27   called with a closure that re-runs the EM loop on resampled units.
 28
 29Example
 30-------
 31```python
 32
 33    import pyfector
 34
 35    result = pyfector.fect(
 36        data=df,
 37        Y="outcome", D="treat",
 38        index=("unit", "year"),
 39        X=["gdp", "pop"],
 40        method="ife",
 41        r=(0, 5),            # CV over 0..5 factors
 42        se=True,
 43        nboots=200,
 44        device="cpu",         # or "gpu"
 45        n_jobs=4,
 46        seed=42,
 47    )
 48    result.summary()
 49    result.plot(kind="gap")
 50```
 51"""
 52
 53from __future__ import annotations
 54
 55from dataclasses import dataclass, field
 56from typing import Any, Literal
 57import os
 58
 59import numpy as np
 60
 61from .backend import set_device, get_backend, to_numpy, to_device, make_rng
 62from .panel import PanelData, prepare_panel, initial_fit
 63from .estimators import estimate_ife, estimate_mc, estimate_cfe, EstimationResult
 64from .cv import cv_ife, cv_mc, CVResult
 65from .inference import bootstrap, jackknife, InferenceResult
 66from .diagnostics import (
 67    Diagnostics,
 68    run_diagnostics as _run_diagnostics,
 69    validate_diagnostics_request,
 70)
 71
 72
 73@dataclass
 74class FectResult:
 75    """Container for all fect estimation results."""
 76    # Method info
 77    method: str
 78    """Estimator used for the point fit."""
 79    r_cv: int | None = None
 80    """Selected IFE rank, or the fixed rank supplied by the caller."""
 81    lambda_cv: float | None = None
 82    """Selected or supplied MC penalty."""
 83
 84    # Point estimates
 85    att_avg: float = 0.0
 86    """ATT averaged over observed treated outcome cells."""
 87    att_avg_unit: float = 0.0
 88    """Mean of unit-specific ATTs over units with observed treated outcomes."""
 89
 90    # Dynamic effects
 91    att_on: np.ndarray | None = None
 92    """Mean counterfactual gap at each relative event time, including pre-treatment periods."""
 93    time_on: np.ndarray | None = None
 94    """Relative event-time grid; treatment onset is zero."""
 95    count_on: np.ndarray | None = None
 96    """Number of observed ever-treated units contributing at each event time."""
 97
 98    # Exit effects (treatment reversal)
 99    att_off: np.ndarray | None = None
100    """Post-reversal event-study effects when populated."""
101    time_off: np.ndarray | None = None
102    """Post-reversal event-time grid when populated."""
103
104    # Coefficients
105    beta: np.ndarray | None = None
106    """Estimated covariate coefficients."""
107    covariate_names: list[str] = field(default_factory=list)
108    """Column names corresponding to beta."""
109
110    # Fixed effects
111    mu: float = 0.0
112    """Estimated intercept."""
113    alpha: np.ndarray | None = None   # unit FE
114    """Estimated unit fixed effects."""
115    xi: np.ndarray | None = None      # time FE
116    """Estimated time fixed effects."""
117    factors: np.ndarray | None = None
118    """Estimated latent time factors."""
119    loadings: np.ndarray | None = None
120    """Estimated unit loadings on the latent factors."""
121
122    # Counterfactual and effects matrices
123    Y_ct: np.ndarray | None = None    # T×N counterfactual
124    """Fitted untreated counterfactual outcomes, shaped (T, N)."""
125    eff: np.ndarray | None = None     # T×N treatment effects
126    """Outcome minus fitted counterfactual, shaped (T, N)."""
127    residuals: np.ndarray | None = None
128    """Masked residuals on the untreated fitting sample."""
129
130    # Model fit
131    sigma2: float = 0.0
132    """Chosen estimator residual variance."""
133    sigma2_fect: float = 0.0   # additive-FE baseline residual variance
134    """Additive fixed-effect baseline residual variance used for the TOST bound."""
135    IC: float = 0.0
136    """Estimator information criterion."""
137    PC: float = 0.0
138    """Estimator panel criterion."""
139    rmse: float = 0.0
140    """Reserved field; defaults to 0.0."""
141    niter: int = 0
142    """Number of EM iterations."""
143    converged: bool = False
144    """Whether the estimator reached its convergence tolerance."""
145
146    # Inference
147    inference: InferenceResult | None = None
148    """Bootstrap or jackknife inference results; see pyfector.inference.InferenceResult."""
149
150    # Diagnostics (populated when fect(..., diagnostics="full" | list))
151    diagnostics: Diagnostics | None = None
152    """Requested diagnostic tests and their resolved options."""
153
154    # CV
155    cv_result: CVResult | None = None
156    """Cross-validation candidates and scores; see pyfector.cv.CVResult."""
157
158    # Panel metadata
159    panel: PanelData | None = None
160    """Filtered panel matrices and unit metadata; see pyfector.panel.PanelData."""
161    fit_options: dict[str, Any] = field(default_factory=dict)
162    """Fit settings used by diagnostic refits."""
163
164    # Reproducibility
165    seed: int | None = None
166    """Seed used for cross-validation and bootstrap draws."""
167
168    def summary(self) -> str:
169        """Print summary table of results."""
170        lines = []
171        lines.append(f"pyfector estimation results")
172        lines.append(f"{'='*60}")
173        lines.append(f"Method: {self.method}")
174        if self.r_cv is not None:
175            lines.append(f"Number of factors (CV): {self.r_cv}")
176        if self.lambda_cv is not None:
177            lines.append(f"Lambda (CV): {self.lambda_cv:.6f}")
178        lines.append(f"Converged: {self.converged} (iter={self.niter})")
179        lines.append(f"Sigma^2: {self.sigma2:.6f}")
180        lines.append(f"Sigma^2_fect (FE baseline): {self.sigma2_fect:.6f}")
181        lines.append(f"")
182        lines.append(f"ATT (average): {self.att_avg:.6f}")
183        if self.inference is not None:
184            inf = self.inference
185            lines.append(f"  SE:     {inf.att_avg_se:.6f}")
186            lines.append(f"  CI:     [{inf.att_avg_ci[0]:.6f}, {inf.att_avg_ci[1]:.6f}]")
187            lines.append(f"  p-val:  {inf.att_avg_pval:.4f}")
188
189        if self.beta is not None and len(self.beta) > 0:
190            lines.append(f"")
191            lines.append(f"Coefficients:")
192            for i, name in enumerate(self.covariate_names):
193                lines.append(f"  {name}: {self.beta[i]:.6f}")
194
195        if self.att_on is not None and self.time_on is not None:
196            lines.append(f"")
197            lines.append(f"Dynamic effects (ATT by relative time):")
198            lines.append(f"  {'Time':>6s}  {'ATT':>10s}  {'Count':>6s}", )
199            for i, t in enumerate(self.time_on):
200                count = self.count_on[i] if self.count_on is not None else ""
201                att = self.att_on[i]
202                if self.inference is not None:
203                    se = self.inference.att_on_se[i]
204                    lines.append(f"  {t:>6.0f}  {att:>10.4f}  ({se:.4f})  {count}")
205                else:
206                    lines.append(f"  {t:>6.0f}  {att:>10.4f}  {count}")
207
208        lines.append(f"{'='*60}")
209        if self.panel is not None:
210            lines.append(f"N={self.panel.N}, T={self.panel.T}")
211        if self.seed is not None:
212            lines.append(f"Seed: {self.seed}")
213        if self.diagnostics is not None:
214            lines.append("")
215            lines.append(self.diagnostics.summary())
216        return "\n".join(lines)
217
218    def __repr__(self):
219        return self.summary()
220
221    def plot(self, kind="gap", **kwargs):
222        """Plot results. Shortcut for ``pyfector.plot(self, kind, ...)``."""
223        from .plotting import plot as _plot
224        return _plot(self, kind=kind, **kwargs)
225
226    def diagnose(self, **kwargs):
227        """Run diagnostic tests. Shortcut for ``pyfector.run_diagnostics(self, ...)``."""
228        from .diagnostics import run_diagnostics
229        return run_diagnostics(self, **kwargs)
230
231
232def fect(
233    data,
234    Y: str,
235    D: str,
236    index: tuple[str, str],
237    X: list[str] | None = None,
238    W: str | None = None,
239    group: str | None = None,
240    method: Literal["fe", "ife", "mc", "cfe", "both"] = "ife",
241    force: Literal["none", "unit", "time", "two-way"] = "two-way",
242    r: int | tuple[int, int] = 0,
243    lam: float | None = None,
244    nlambda: int = 10,
245    lambda_candidates: list[float] | np.ndarray | None = None,
246    CV: bool = True,
247    k: int = 10,
248    cv_prop: float = 0.1,
249    cv_nobs: int = 3,
250    cv_treat: bool = True,
251    cv_donut: int = 0,
252    criterion: str = "mspe",
253    cv_rule: Literal["min", "onepct"] = "min",
254    se: bool = False,
255    vartype: Literal["bootstrap", "jackknife"] = "bootstrap",
256    nboots: int = 200,
257    alpha: float = 0.05,
258    tol: float = 1e-7,
259    max_iter: int = 5000,
260    min_T0: int = 1,
261    min_T0_strict: bool = False,
262    max_missing: float = 1.0,
263    normalize: bool = False,
264    # CFE-specific
265    Z: list[str] | None = None,
266    Q: list[str] | None = None,
267    # Performance
268    device: Literal["cpu", "gpu"] = "cpu",
269    n_jobs: int | None = -1,
270    seed: int | None = None,
271    # Diagnostics (optional; run at fit time and attached to result)
272    diagnostics: Literal["none", "full"] | list[str] = "none",
273    diagnostics_options: dict | None = None,
274    *,
275    cluster: str | None = None,
276) -> FectResult:
277    """Estimate counterfactual treatment effects for panel data.
278
279    This is the main Python entry point for the counterfactual estimator
280    workflow.  Where the paper and the historical R package differ,
281    pyfector defaults to the paper's statistical definition and exposes
282    R-package-style behavior through explicit options.
283
284    Missing outcome policy
285    ----------------------
286    pyfector distinguishes raw missing outcomes from counterfactual
287    missingness caused by treatment.  Observed untreated cells
288    (``D == 0`` and non-missing ``Y``) fit the response surface.  Observed
289    treated cells (``D == 1`` and non-missing ``Y``) contribute to ATT as
290    ``Y - Y_ct``.  If a treated outcome is missing in the input data, the
291    model can still produce a counterfactual ``Y_ct`` for that cell, but
292    the cell is not counted in ``att_avg`` or ``att_on`` because the
293    treated potential outcome was not observed.
294
295    By default, ``min_T0`` is enforced only for treated and reversal
296    units.  Sparse controls are retained if they have at least one
297    observed outcome, because they may still inform the low-rank response
298    surface.  Set ``min_T0_strict=True`` to require controls to satisfy
299    ``min_T0`` too, matching the more conservative R fect sparse-panel
300    behavior.
301
302    Parameters
303    ----------
304    data : polars.DataFrame, pandas.DataFrame
305        Long-format panel data.
306    Y, D : str
307        Column names for outcome and binary treatment indicator.
308    index : (str, str)
309        Column names for (unit_id, time_period).
310    X : list of str, optional
311        Time-varying covariates.
312    W : str, optional
313        Observation weight column.
314    group : str, optional
315        Reserved for grouped estimation. Currently raises
316        ``NotImplementedError`` when supplied.
317    method : {"fe", "ife", "mc", "cfe", "both"}
318        Estimation method.
319    force : {"none", "unit", "time", "two-way"}
320        Fixed effects specification.
321    r : int or (int, int)
322        Number of factors.  If tuple, CV selects from range.
323    lam : float, optional
324        Nuclear norm penalty for MC.  If None with CV=True, auto-selected.
325    nlambda : int
326        Number of automatically generated lambda candidates for MC CV.
327    lambda_candidates : array-like, optional
328        Explicit non-negative lambda candidates for MC CV. When supplied,
329        ``nlambda`` is ignored.
330    CV : bool
331        If True, cross-validate over ``r`` for IFE when ``r`` is a tuple,
332        or over ``lam`` for MC when ``lam`` is None.
333    k : int
334        Number of CV folds.
335    cv_prop : float
336        Fraction of eligible observed control cells masked per CV fold.
337    cv_nobs : int
338        Number of consecutive within-unit observations to mask as a block.
339    cv_treat : bool
340        If True, restrict CV masks to pre-treatment cells of ever-treated
341        units. If False, use all observed control cells.
342    cv_donut : int
343        Exclude this many periods around treatment onset from CV evaluation.
344    criterion : {"mspe", "gmspe", "mad"}
345        Cross-validation loss.
346    cv_rule : {"min", "onepct"}
347        CV selection rule. ``"min"`` chooses the strict minimum-score
348        candidate and is the paper-faithful default. ``"onepct"`` chooses
349        the simplest candidate within 1% of the best score (lower ``r`` for
350        IFE, higher ``lam`` for MC).
351    se : bool
352        Compute standard errors via bootstrap/jackknife.
353    vartype : {"bootstrap", "jackknife"}
354        Inference method when ``se=True``.
355    cluster : str, optional
356        Column of cluster labels, constant within each unit. Bootstrap
357        inference resamples whole clusters in both the main fit and the
358        placebo refit. Requires at least two retained clusters.
359    nboots : int
360        Number of bootstrap replications, at least 2. Ignored for jackknife.
361    alpha : float
362        Significance level for confidence intervals and tests.
363    tol : float
364        EM convergence tolerance for final point estimation.
365    max_iter : int
366        Maximum EM iterations.
367    min_T0 : int
368        Minimum untreated/pre-treatment observed periods. By default this is
369        enforced only for treated and treatment-reversal units.
370    min_T0_strict : bool
371        If True, enforce ``min_T0`` on all units, including controls. This
372        matches R fect's conservative handling of sparse control rows.
373    max_missing : float
374        Maximum missing-outcome fraction per unit, in ``[0, 1]``. Units with
375        no observed outcomes are always dropped, regardless of this threshold,
376        because they provide neither fitting information nor observed treated
377        effects.
378    normalize : bool
379        If True, estimate on an outcome standardized by its observed standard
380        deviation, then transform effects back to the original scale.
381    Z, Q : list of str, optional
382        Reserved CFE interaction arguments. Currently raise
383        ``NotImplementedError`` when supplied.
384    device : {"cpu", "gpu"}
385        Compute device.
386    n_jobs : int, optional
387        Parallel workers for CV and bootstrap. ``-1`` or ``None`` uses
388        all available CPUs.
389    seed : int, optional
390        Random seed for full reproducibility.
391
392    Returns
393    -------
394    FectResult
395        Treatment effects, fitted counterfactuals, and requested inference
396        and diagnostics.
397    """
398    # Set device
399    set_device(device)
400    xp = get_backend()
401    n_jobs = _resolve_n_jobs(n_jobs)
402    if device == "gpu":
403        n_jobs = 1
404
405    if group is not None:
406        raise NotImplementedError("The `group` argument is not implemented yet.")
407    if Z is not None or Q is not None:
408        raise NotImplementedError("The `Z` and `Q` CFE interaction arguments are not implemented yet.")
409    if criterion not in {"mspe", "gmspe", "mad"}:
410        raise ValueError("criterion must be 'mspe', 'gmspe', or 'mad'")
411    if cv_rule not in {"min", "onepct"}:
412        raise ValueError("cv_rule must be 'min' or 'onepct'")
413    if min_T0 < 0:
414        raise ValueError("min_T0 must be non-negative")
415    if cluster is not None and vartype != "bootstrap":
416        raise ValueError("cluster is supported only with vartype='bootstrap'")
417    if se and vartype == "bootstrap" and (
418        isinstance(nboots, (bool, np.bool_))
419        or not isinstance(nboots, (int, np.integer)) or nboots < 2
420    ):
421        raise ValueError("nboots must be an integer of at least 2")
422    if not 0.0 <= max_missing <= 1.0:
423        raise ValueError("max_missing must be between 0 and 1")
424
425    # Validate diagnostics request before doing any expensive estimation
426    # so users don't wait for a 30-min MC fit only to find their config
427    # is wrong.
428    requested_diag = validate_diagnostics_request(
429        diagnostics, diagnostics_options, se,
430    )
431
432    # Map force string to int
433    force_map = {"none": 0, "unit": 1, "time": 2, "two-way": 3}
434    force_int = force_map[force]
435
436    # Prepare panel data
437    panel = prepare_panel(
438        data, Y=Y, D=D, index=index, X=X, W=W,
439        group=group, min_T0=min_T0, min_T0_strict=min_T0_strict,
440        max_missing=max_missing,
441        cluster=cluster,
442    )
443
444    if not np.any((panel.D > 0) & (panel.I > 0)):
445        raise ValueError(
446            "No observed treated units remain after filtering. "
447            "Check treatment, min_T0, and missing-outcome filters."
448        )
449    if se and panel.cluster is not None and len(np.unique(panel.cluster)) < 2:
450        raise ValueError("cluster bootstrap requires at least two clusters after filtering")
451
452    # Move to device
453    Y_mat = to_device(panel.Y)
454    D_mat = to_device(panel.D)
455    I_mat = to_device(panel.I)
456    II_mat = to_device(panel.II)
457    X_mat = to_device(panel.X) if panel.X is not None else None
458    W_mat = to_device(panel.W) if panel.W is not None else None
459
460    # Normalize
461    norm_factor = 1.0
462    if normalize:
463        sd_y = float(xp.std(Y_mat[I_mat > 0]))
464        if sd_y > 0:
465            Y_mat = Y_mat / sd_y
466            norm_factor = sd_y
467
468    # Initial fit
469    Y0, beta0 = initial_fit(Y_mat, X_mat, II_mat, force_int)
470
471    # Determine r and lambda
472    r_cv = None
473    lambda_cv = None
474    cv_result = None
475
476    if method == "ife":
477        if isinstance(r, tuple) and CV:
478            cv_result = cv_ife(
479                Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0,
480                force=force_int, r_range=r, k=k, cv_prop=cv_prop,
481                cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut,
482                criterion=criterion, cv_rule=cv_rule,
483                tol=tol, max_iter=max_iter,
484                n_jobs=n_jobs, seed=seed,
485            )
486            r_cv = cv_result.best_r
487        else:
488            r_cv = r if isinstance(r, int) else r[0]
489
490    elif method == "mc":
491        if lam is None and CV:
492            cv_result = cv_mc(
493                Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0,
494                force=force_int, lambda_candidates=lambda_candidates,
495                nlambda=nlambda, k=k, cv_prop=cv_prop,
496                cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut,
497                criterion=criterion, cv_rule=cv_rule,
498                tol=tol, max_iter=max_iter,
499                n_jobs=n_jobs, seed=seed,
500            )
501            lambda_cv = cv_result.best_lambda
502        else:
503            lambda_cv = lam if lam is not None else 0.0
504
505    elif method == "fe":
506        r_cv = 0
507
508    elif method == "cfe":
509        r_cv = r if isinstance(r, int) else r[0]
510
511    # Point estimation
512    if method in ("fe", "ife"):
513        est = estimate_ife(
514            Y_mat, Y0, X_mat, II_mat, W_mat, beta0,
515            r=r_cv, force=force_int, tol=tol, max_iter=max_iter,
516        )
517    elif method == "mc":
518        est = estimate_mc(
519            Y_mat, Y0, X_mat, II_mat, W_mat, beta0,
520            lam=lambda_cv, force=force_int, tol=tol, max_iter=max_iter,
521        )
522    elif method == "cfe":
523        est = estimate_cfe(
524            Y_mat, Y0, X_mat, II_mat, W_mat, beta0,
525            r=r_cv, force=force_int, tol=tol, max_iter=max_iter,
526        )
527    elif method == "both":
528        # Run both IFE and MC, return IFE results with MC comparison
529        if isinstance(r, tuple) and CV:
530            cv_result = cv_ife(
531                Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0,
532                force=force_int, r_range=r, k=k, cv_prop=cv_prop,
533                cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut,
534                criterion=criterion, cv_rule=cv_rule,
535                tol=tol, max_iter=max_iter,
536                n_jobs=n_jobs, seed=seed,
537            )
538            r_cv = cv_result.best_r
539        else:
540            r_cv = r if isinstance(r, int) else r[0]
541        est = estimate_ife(
542            Y_mat, Y0, X_mat, II_mat, W_mat, beta0,
543            r=r_cv, force=force_int, tol=tol, max_iter=max_iter,
544        )
545    else:
546        raise ValueError(f"Unknown method: {method}")
547
548    # Compute effects
549    eff = Y_mat - est.fit
550    Y_ct = est.fit
551
552    # Additive-FE baseline residual variance (Liu et al. 2024 sigma2.fect).
553    # For method="fe" the main estimator IS the additive-FE pass, so reuse
554    # est.sigma2. For ife/mc/cfe/both, run an extra r=0 IFE pass on the
555    # same panel with the user's requested FE structure.
556    if method == "fe":
557        sigma2_fect_value = float(est.sigma2)
558    else:
559        est_fect = estimate_ife(
560            Y_mat, Y0, X_mat, II_mat, W_mat, beta0,
561            r=0, force=force_int, tol=tol, max_iter=max_iter,
562        )
563        sigma2_fect_value = float(est_fect.sigma2)
564
565    # Denormalize
566    if normalize and norm_factor != 1.0:
567        eff = eff * norm_factor
568        Y_ct = Y_ct * norm_factor
569        Y_mat = Y_mat * norm_factor
570        if est.beta is not None:
571            est = est._replace(beta=est.beta * norm_factor)
572        sigma2_fect_value *= norm_factor ** 2
573
574    # ATT computation
575    T_on = to_device(panel.T_on)
576    att_avg, att_on, time_on, count_on, att_avg_unit = _compute_effects(
577        to_numpy(eff), to_numpy(D_mat), to_numpy(panel.T_on), to_numpy(I_mat),
578    )
579
580    # Build result
581    result = FectResult(
582        method=method,
583        r_cv=r_cv,
584        lambda_cv=lambda_cv,
585        att_avg=att_avg,
586        att_avg_unit=att_avg_unit,
587        att_on=att_on,
588        time_on=time_on,
589        count_on=count_on,
590        beta=to_numpy(est.beta) if est.beta is not None else None,
591        covariate_names=panel.covariate_names,
592        mu=est.mu,
593        alpha=to_numpy(est.alpha) if est.alpha is not None else None,
594        xi=to_numpy(est.xi) if est.xi is not None else None,
595        factors=to_numpy(est.factors) if est.factors is not None else None,
596        loadings=to_numpy(est.loadings) if est.loadings is not None else None,
597        Y_ct=to_numpy(Y_ct),
598        eff=to_numpy(eff),
599        residuals=to_numpy(est.residuals),
600        sigma2=est.sigma2,
601        sigma2_fect=sigma2_fect_value,
602        IC=est.IC,
603        PC=est.PC,
604        niter=est.niter,
605        converged=est.converged,
606        cv_result=cv_result,
607        panel=panel,
608        fit_options={
609            "force": force,
610            "force_int": force_int,
611            "tol": tol,
612            "max_iter": max_iter,
613            "normalize": normalize,
614            "norm_factor": norm_factor,
615            "vartype": vartype,
616            "nboots": nboots,
617            "n_jobs": n_jobs,
618            "min_T0": min_T0,
619            "min_T0_strict": min_T0_strict,
620            "cluster": cluster,
621        },
622        seed=seed,
623    )
624
625    # Inference
626    if se:
627        result.inference = _run_inference(
628            result, panel, Y_mat, X_mat, W_mat, beta0, Y0,
629            method=method, r_cv=r_cv, lambda_cv=lambda_cv,
630            force_int=force_int, tol=tol, max_iter=max_iter,
631            vartype=vartype, nboots=nboots, alpha=alpha,
632            n_jobs=n_jobs, seed=seed, normalize=normalize,
633            norm_factor=norm_factor,
634        )
635
636    # Run requested diagnostics at fit time. requested_diag is None when
637    # diagnostics="none". Validation already enforced se=True and
638    # required-config presence.
639    if requested_diag is not None:
640        opts = dict(diagnostics_options or {})
641        if "loo" in requested_diag:
642            opts["loo"] = True
643        else:
644            opts.setdefault("loo", False)
645        result.diagnostics = _run_diagnostics(
646            result, _requested=requested_diag,
647            _strict=isinstance(diagnostics, list), **opts,
648        )
649
650    return result
651
652
653def _compute_effects(eff, D, T_on, I):
654    """Compute overall ATT, per-unit ATT, and dynamic ATT by event time.
655
656    ATT averages are defined only over observed outcome cells.  A missing
657    raw treated outcome has no observed ``Y(1)``, so it is excluded even
658    though the estimator may have produced a counterfactual ``Y_ct`` for
659    that matrix position.
660
661    Dynamic ATT grouping is done with :func:`numpy.bincount` so the cost
662    is O(n_observed_ever_treated_cells) regardless of the number of
663    distinct relative-time values.  Per-unit ATT uses a column-wise
664    masked mean via ``np.add.reduce`` rather than a Python loop.
665    """
666    treated = (D > 0) & (I > 0)
667    n_treated = int(np.sum(treated))
668
669    att_avg = float(np.sum(eff[treated]) / max(n_treated, 1))
670
671    # Per-unit ATT (post-treatment only) — sum then divide by count,
672    # keeping only units that have at least one treated observation.
673    col_counts = treated.sum(axis=0)                                 # (N,)
674    col_sums = (eff * treated).sum(axis=0)                           # (N,)
675    has_any = col_counts > 0
676    if np.any(has_any):
677        unit_atts = col_sums[has_any] / col_counts[has_any]
678        att_avg_unit = float(unit_atts.mean())
679    else:
680        att_avg_unit = 0.0
681
682    # Dynamic ATT by relative event time.  Include pre-treatment periods
683    # for ever-treated units (counterfactual gaps before onset).
684    ever_treated = np.any(D > 0, axis=0)                             # (N,)
685    all_periods = (I > 0) & ever_treated[np.newaxis, :]              # (T, N)
686
687    T_on_flat = T_on[all_periods]
688    eff_flat = eff[all_periods]
689    valid = ~np.isnan(T_on_flat)
690    T_on_flat = T_on_flat[valid].astype(np.int64)
691    eff_flat = eff_flat[valid]
692
693    if T_on_flat.size == 0:
694        return att_avg, np.array([]), np.array([]), np.array([], dtype=np.int64), att_avg_unit
695
696    # Bincount over the shifted integer indices.
697    offset = int(T_on_flat.min())
698    idx = T_on_flat - offset
699    minlength = int(T_on_flat.max() - offset + 1)
700    sums = np.bincount(idx, weights=eff_flat, minlength=minlength)
701    counts = np.bincount(idx, minlength=minlength)
702    keep = counts > 0
703    time_on = (offset + np.arange(minlength))[keep].astype(np.float64)
704    att_on = (sums[keep] / counts[keep]).astype(np.float64)
705    count_on = counts[keep].astype(np.int64)
706
707    return att_avg, att_on, time_on, count_on, att_avg_unit
708
709
710def _run_inference(
711    result, panel, Y_mat, X_mat, W_mat, beta0, Y0,
712    method, r_cv, lambda_cv, force_int, tol, max_iter,
713    vartype, nboots, alpha, n_jobs, seed, normalize, norm_factor,
714):
715    """Run bootstrap or jackknife inference.
716
717    Bootstrap replicates use a relaxed convergence tolerance
718    (``max(tol, 1e-3)``) because bootstrap SEs converge to 3–4
719    significant digits regardless of inner precision.  The full-sample
720    fit is used as the warm-start initialiser for each replicate,
721    which cuts EM iterations by ~30-60 %.
722    """
723    xp = get_backend()
724
725    # Relaxed tolerance: bootstrap SEs don't benefit from tight inner tol.
726    boot_tol = max(tol, 1e-3)
727
728    # Pre-move panel data to device once (avoids CPU→GPU copy per bootstrap rep)
729    II_dev = to_device(panel.II)
730    D_dev = to_device(panel.D)
731    I_dev = to_device(panel.I)
732    T_on_np = panel.T_on  # keep on CPU for ATT computation
733
734    # Warm-start: use full-sample fitted values as initial imputation
735    # for each replicate instead of recomputing initial_fit from scratch.
736    Y0_full = to_device(result.Y_ct)
737
738    def _estimate_fn(unit_idx):
739        """Re-estimate on a subset of units."""
740        Y_sub = Y_mat[:, unit_idx]
741        II_sub = II_dev[:, unit_idx]
742        D_sub = D_dev[:, unit_idx]
743        I_sub = I_dev[:, unit_idx]
744        X_sub = X_mat[:, unit_idx, :] if X_mat is not None else None
745        W_sub = W_mat[:, unit_idx] if W_mat is not None else None
746        T_on_sub = T_on_np[:, unit_idx]
747        beta0_sub = beta0
748
749        # Warm-start from full-sample fit (columns for resampled units).
750        # This is much closer to the replicate's solution than a cold
751        # initial_fit, cutting EM iterations significantly.
752        Y0_sub = Y0_full[:, unit_idx].copy()
753
754        if method in ("fe", "ife"):
755            est = estimate_ife(
756                Y_sub, Y0_sub, X_sub, II_sub, W_sub, beta0_sub,
757                r=r_cv, force=force_int, tol=boot_tol, max_iter=max_iter,
758            )
759        elif method == "mc":
760            est = estimate_mc(
761                Y_sub, Y0_sub, X_sub, II_sub, W_sub, beta0_sub,
762                lam=lambda_cv, force=force_int, tol=boot_tol, max_iter=max_iter,
763            )
764        else:
765            est = estimate_ife(
766                Y_sub, Y0_sub, X_sub, II_sub, W_sub, beta0_sub,
767                r=r_cv or 0, force=force_int, tol=boot_tol, max_iter=max_iter,
768            )
769
770        eff = to_numpy(Y_sub - est.fit)
771        if normalize and norm_factor != 1.0:
772            eff = eff * norm_factor
773
774        return eff, to_numpy(D_sub), T_on_sub, to_numpy(I_sub)
775
776    if vartype == "bootstrap":
777        return bootstrap(
778            _estimate_fn, to_numpy(Y_mat), to_numpy(panel.D),
779            to_numpy(panel.I), panel.T_on, panel.unit_type,
780            nboots=nboots, alpha=alpha, n_jobs=n_jobs, seed=seed,
781            cluster=panel.cluster,
782            point_estimate=(result.eff, panel.D, panel.T_on, panel.I),
783        )
784    else:
785        return jackknife(
786            _estimate_fn, to_numpy(Y_mat), to_numpy(panel.D),
787            to_numpy(panel.I), panel.T_on, panel.unit_type,
788            alpha=alpha, n_jobs=n_jobs,
789        )
790
791
792def _resolve_n_jobs(n_jobs: int | None) -> int:
793    """Normalize public n_jobs values before passing them to joblib."""
794    if n_jobs is None or n_jobs == -1:
795        return os.cpu_count() or 1
796    if n_jobs == 0 or n_jobs < -1:
797        raise ValueError("n_jobs must be a positive integer, -1, or None")
798    return int(n_jobs)