pyfector
pyfector — Fast counterfactual estimators for panel data in Python.
A high-performance Python reimplementation of the R fect package,
featuring randomized SVD, GPU acceleration (CuPy), parallel computing
(joblib), Polars data ingestion, and seeded reproducibility.
Usage::
import pyfector
result = pyfector.fect(
data=df,
Y="outcome", D="treat",
index=("unit", "year"),
method="ife",
r=(0, 5),
se=True,
seed=42,
)
result.summary()
result.plot()
1""" 2pyfector — Fast counterfactual estimators for panel data in Python. 3 4A high-performance Python reimplementation of the R ``fect`` package, 5featuring randomized SVD, GPU acceleration (CuPy), parallel computing 6(joblib), Polars data ingestion, and seeded reproducibility. 7 8Usage:: 9 10 import pyfector 11 12 result = pyfector.fect( 13 data=df, 14 Y="outcome", D="treat", 15 index=("unit", "year"), 16 method="ife", 17 r=(0, 5), 18 se=True, 19 seed=42, 20 ) 21 result.summary() 22 result.plot() 23 24""" 25 26__version__ = "0.2.1" 27 28from .fect import fect, FectResult 29from .backend import set_device, get_device 30from .diagnostics import ( 31 run_diagnostics, 32 Diagnostics, 33 DiagnosticResult, 34 TostResult, 35 PretrendFResult, 36 EquivFResult, 37 PlaceboResult, 38 CarryoverResult, 39 LooResult, 40) 41from .plotting import plot 42 43__all__ = [ 44 "fect", 45 "FectResult", 46 "set_device", 47 "get_device", 48 "run_diagnostics", 49 "Diagnostics", 50 "DiagnosticResult", 51 "TostResult", 52 "PretrendFResult", 53 "EquivFResult", 54 "PlaceboResult", 55 "CarryoverResult", 56 "LooResult", 57 "plot", 58]
199def fect( 200 data, 201 Y: str, 202 D: str, 203 index: tuple[str, str], 204 X: list[str] | None = None, 205 W: str | None = None, 206 group: str | None = None, 207 method: Literal["fe", "ife", "mc", "cfe", "both"] = "ife", 208 force: Literal["none", "unit", "time", "two-way"] = "two-way", 209 r: int | tuple[int, int] = 0, 210 lam: float | None = None, 211 nlambda: int = 10, 212 lambda_candidates: list[float] | np.ndarray | None = None, 213 CV: bool = True, 214 k: int = 10, 215 cv_prop: float = 0.1, 216 cv_nobs: int = 3, 217 cv_treat: bool = True, 218 cv_donut: int = 0, 219 criterion: str = "mspe", 220 cv_rule: Literal["min", "onepct"] = "min", 221 se: bool = False, 222 vartype: Literal["bootstrap", "jackknife"] = "bootstrap", 223 nboots: int = 200, 224 alpha: float = 0.05, 225 tol: float = 1e-7, 226 max_iter: int = 5000, 227 min_T0: int = 1, 228 min_T0_strict: bool = False, 229 max_missing: float = 1.0, 230 normalize: bool = False, 231 # CFE-specific 232 Z: list[str] | None = None, 233 Q: list[str] | None = None, 234 # Performance 235 device: Literal["cpu", "gpu"] = "cpu", 236 n_jobs: int | None = -1, 237 seed: int | None = None, 238 # Diagnostics (optional; run at fit time and attached to result) 239 diagnostics: Literal["none", "full"] | list[str] = "none", 240 diagnostics_options: dict | None = None, 241) -> FectResult: 242 """Estimate counterfactual treatment effects for panel data. 243 244 This is the main Python entry point for the counterfactual estimator 245 workflow. Where the paper and the historical R package differ, 246 pyfector defaults to the paper's statistical definition and exposes 247 R-package-style behavior through explicit options. 248 249 Missing outcome policy 250 ---------------------- 251 pyfector distinguishes raw missing outcomes from counterfactual 252 missingness caused by treatment. Observed untreated cells 253 (``D == 0`` and non-missing ``Y``) fit the response surface. Observed 254 treated cells (``D == 1`` and non-missing ``Y``) contribute to ATT as 255 ``Y - Y_ct``. If a treated outcome is missing in the input data, the 256 model can still produce a counterfactual ``Y_ct`` for that cell, but 257 the cell is not counted in ``att_avg`` or ``att_on`` because the 258 treated potential outcome was not observed. 259 260 By default, ``min_T0`` is enforced only for treated and reversal 261 units. Sparse controls are retained if they have at least one 262 observed outcome, because they may still inform the low-rank response 263 surface. Set ``min_T0_strict=True`` to require controls to satisfy 264 ``min_T0`` too, matching the more conservative R fect sparse-panel 265 behavior. 266 267 Parameters 268 ---------- 269 data : polars.DataFrame, pandas.DataFrame 270 Long-format panel data. 271 Y, D : str 272 Column names for outcome and binary treatment indicator. 273 index : (str, str) 274 Column names for (unit_id, time_period). 275 X : list of str, optional 276 Time-varying covariates. 277 W : str, optional 278 Observation weight column. 279 group : str, optional 280 Reserved for grouped estimation. Currently raises 281 ``NotImplementedError`` when supplied. 282 method : {"fe", "ife", "mc", "cfe", "both"} 283 Estimation method. 284 force : {"none", "unit", "time", "two-way"} 285 Fixed effects specification. 286 r : int or (int, int) 287 Number of factors. If tuple, CV selects from range. 288 lam : float, optional 289 Nuclear norm penalty for MC. If None with CV=True, auto-selected. 290 nlambda : int 291 Number of automatically generated lambda candidates for MC CV. 292 lambda_candidates : array-like, optional 293 Explicit non-negative lambda candidates for MC CV. When supplied, 294 ``nlambda`` is ignored. 295 CV : bool 296 If True, cross-validate over ``r`` for IFE when ``r`` is a tuple, 297 or over ``lam`` for MC when ``lam`` is None. 298 k : int 299 Number of CV folds. 300 cv_prop : float 301 Fraction of eligible observed control cells masked per CV fold. 302 cv_nobs : int 303 Number of consecutive within-unit observations to mask as a block. 304 cv_treat : bool 305 If True, restrict CV masks to pre-treatment cells of ever-treated 306 units. If False, use all observed control cells. 307 cv_donut : int 308 Exclude this many periods around treatment onset from CV evaluation. 309 criterion : {"mspe", "gmspe", "mad"} 310 Cross-validation loss. 311 cv_rule : {"min", "onepct"} 312 CV selection rule. ``"min"`` chooses the strict minimum-score 313 candidate and is the paper-faithful default. ``"onepct"`` chooses 314 the simplest candidate within 1% of the best score (lower ``r`` for 315 IFE, higher ``lam`` for MC). 316 se : bool 317 Compute standard errors via bootstrap/jackknife. 318 vartype : {"bootstrap", "jackknife"} 319 Inference method when ``se=True``. 320 nboots : int 321 Number of bootstrap replications. Ignored for jackknife. 322 alpha : float 323 Significance level for confidence intervals and tests. 324 tol : float 325 EM convergence tolerance for final point estimation. 326 max_iter : int 327 Maximum EM iterations. 328 min_T0 : int 329 Minimum untreated/pre-treatment observed periods. By default this is 330 enforced only for treated and treatment-reversal units. 331 min_T0_strict : bool 332 If True, enforce ``min_T0`` on all units, including controls. This 333 matches R fect's conservative handling of sparse control rows. 334 max_missing : float 335 Maximum missing-outcome fraction per unit, in ``[0, 1]``. Units with 336 no observed outcomes are always dropped, regardless of this threshold, 337 because they provide neither fitting information nor observed treated 338 effects. 339 normalize : bool 340 If True, estimate on an outcome standardized by its observed standard 341 deviation, then transform effects back to the original scale. 342 Z, Q : list of str, optional 343 Reserved CFE interaction arguments. Currently raise 344 ``NotImplementedError`` when supplied. 345 device : {"cpu", "gpu"} 346 Compute device. 347 n_jobs : int, optional 348 Parallel workers for CV and bootstrap. ``-1`` or ``None`` uses 349 all available CPUs. 350 seed : int, optional 351 Random seed for full reproducibility. 352 """ 353 # Set device 354 set_device(device) 355 xp = get_backend() 356 n_jobs = _resolve_n_jobs(n_jobs) 357 if device == "gpu": 358 n_jobs = 1 359 360 if group is not None: 361 raise NotImplementedError("The `group` argument is not implemented yet.") 362 if Z is not None or Q is not None: 363 raise NotImplementedError("The `Z` and `Q` CFE interaction arguments are not implemented yet.") 364 if criterion not in {"mspe", "gmspe", "mad"}: 365 raise ValueError("criterion must be 'mspe', 'gmspe', or 'mad'") 366 if cv_rule not in {"min", "onepct"}: 367 raise ValueError("cv_rule must be 'min' or 'onepct'") 368 if min_T0 < 0: 369 raise ValueError("min_T0 must be non-negative") 370 if not 0.0 <= max_missing <= 1.0: 371 raise ValueError("max_missing must be between 0 and 1") 372 373 # Validate diagnostics request before doing any expensive estimation 374 # so users don't wait for a 30-min MC fit only to find their config 375 # is wrong. 376 requested_diag = validate_diagnostics_request( 377 diagnostics, diagnostics_options, se, 378 ) 379 380 # Map force string to int 381 force_map = {"none": 0, "unit": 1, "time": 2, "two-way": 3} 382 force_int = force_map[force] 383 384 # Prepare panel data 385 panel = prepare_panel( 386 data, Y=Y, D=D, index=index, X=X, W=W, 387 group=group, min_T0=min_T0, min_T0_strict=min_T0_strict, 388 max_missing=max_missing, 389 ) 390 391 # Move to device 392 Y_mat = to_device(panel.Y) 393 D_mat = to_device(panel.D) 394 I_mat = to_device(panel.I) 395 II_mat = to_device(panel.II) 396 X_mat = to_device(panel.X) if panel.X is not None else None 397 W_mat = to_device(panel.W) if panel.W is not None else None 398 399 # Normalize 400 norm_factor = 1.0 401 if normalize: 402 sd_y = float(xp.std(Y_mat[I_mat > 0])) 403 if sd_y > 0: 404 Y_mat = Y_mat / sd_y 405 norm_factor = sd_y 406 407 # Initial fit 408 Y0, beta0 = initial_fit(Y_mat, X_mat, II_mat, force_int) 409 410 # Determine r and lambda 411 r_cv = None 412 lambda_cv = None 413 cv_result = None 414 415 if method == "ife": 416 if isinstance(r, tuple) and CV: 417 cv_result = cv_ife( 418 Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0, 419 force=force_int, r_range=r, k=k, cv_prop=cv_prop, 420 cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut, 421 criterion=criterion, cv_rule=cv_rule, 422 tol=tol, max_iter=max_iter, 423 n_jobs=n_jobs, seed=seed, 424 ) 425 r_cv = cv_result.best_r 426 else: 427 r_cv = r if isinstance(r, int) else r[0] 428 429 elif method == "mc": 430 if lam is None and CV: 431 cv_result = cv_mc( 432 Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0, 433 force=force_int, lambda_candidates=lambda_candidates, 434 nlambda=nlambda, k=k, cv_prop=cv_prop, 435 cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut, 436 criterion=criterion, cv_rule=cv_rule, 437 tol=tol, max_iter=max_iter, 438 n_jobs=n_jobs, seed=seed, 439 ) 440 lambda_cv = cv_result.best_lambda 441 else: 442 lambda_cv = lam if lam is not None else 0.0 443 444 elif method == "fe": 445 r_cv = 0 446 447 elif method == "cfe": 448 r_cv = r if isinstance(r, int) else r[0] 449 450 # Point estimation 451 if method in ("fe", "ife"): 452 est = estimate_ife( 453 Y_mat, Y0, X_mat, II_mat, W_mat, beta0, 454 r=r_cv, force=force_int, tol=tol, max_iter=max_iter, 455 ) 456 elif method == "mc": 457 est = estimate_mc( 458 Y_mat, Y0, X_mat, II_mat, W_mat, beta0, 459 lam=lambda_cv, force=force_int, tol=tol, max_iter=max_iter, 460 ) 461 elif method == "cfe": 462 est = estimate_cfe( 463 Y_mat, Y0, X_mat, II_mat, W_mat, beta0, 464 r=r_cv, force=force_int, tol=tol, max_iter=max_iter, 465 ) 466 elif method == "both": 467 # Run both IFE and MC, return IFE results with MC comparison 468 if isinstance(r, tuple) and CV: 469 cv_result = cv_ife( 470 Y_mat, Y0, X_mat, I_mat, II_mat, D_mat, W_mat, beta0, 471 force=force_int, r_range=r, k=k, cv_prop=cv_prop, 472 cv_nobs=cv_nobs, cv_treat=cv_treat, cv_donut=cv_donut, 473 criterion=criterion, cv_rule=cv_rule, 474 tol=tol, max_iter=max_iter, 475 n_jobs=n_jobs, seed=seed, 476 ) 477 r_cv = cv_result.best_r 478 else: 479 r_cv = r if isinstance(r, int) else r[0] 480 est = estimate_ife( 481 Y_mat, Y0, X_mat, II_mat, W_mat, beta0, 482 r=r_cv, force=force_int, tol=tol, max_iter=max_iter, 483 ) 484 else: 485 raise ValueError(f"Unknown method: {method}") 486 487 # Compute effects 488 eff = Y_mat - est.fit 489 Y_ct = est.fit 490 491 # Additive-FE baseline residual variance (Liu et al. 2024 sigma2.fect). 492 # For method="fe" the main estimator IS the additive-FE pass, so reuse 493 # est.sigma2. For ife/mc/cfe/both, run an extra r=0 IFE pass on the 494 # same panel with the user's requested FE structure. 495 if method == "fe": 496 sigma2_fect_value = float(est.sigma2) 497 else: 498 est_fect = estimate_ife( 499 Y_mat, Y0, X_mat, II_mat, W_mat, beta0, 500 r=0, force=force_int, tol=tol, max_iter=max_iter, 501 ) 502 sigma2_fect_value = float(est_fect.sigma2) 503 504 # Denormalize 505 if normalize and norm_factor != 1.0: 506 eff = eff * norm_factor 507 Y_ct = Y_ct * norm_factor 508 Y_mat = Y_mat * norm_factor 509 if est.beta is not None: 510 est = est._replace(beta=est.beta * norm_factor) 511 sigma2_fect_value *= norm_factor ** 2 512 513 # ATT computation 514 T_on = to_device(panel.T_on) 515 att_avg, att_on, time_on, count_on, att_avg_unit = _compute_effects( 516 to_numpy(eff), to_numpy(D_mat), to_numpy(panel.T_on), to_numpy(I_mat), 517 ) 518 519 # Build result 520 result = FectResult( 521 method=method, 522 r_cv=r_cv, 523 lambda_cv=lambda_cv, 524 att_avg=att_avg, 525 att_avg_unit=att_avg_unit, 526 att_on=att_on, 527 time_on=time_on, 528 count_on=count_on, 529 beta=to_numpy(est.beta) if est.beta is not None else None, 530 covariate_names=panel.covariate_names, 531 mu=est.mu, 532 alpha=to_numpy(est.alpha) if est.alpha is not None else None, 533 xi=to_numpy(est.xi) if est.xi is not None else None, 534 factors=to_numpy(est.factors) if est.factors is not None else None, 535 loadings=to_numpy(est.loadings) if est.loadings is not None else None, 536 Y_ct=to_numpy(Y_ct), 537 eff=to_numpy(eff), 538 residuals=to_numpy(est.residuals), 539 sigma2=est.sigma2, 540 sigma2_fect=sigma2_fect_value, 541 IC=est.IC, 542 PC=est.PC, 543 niter=est.niter, 544 converged=est.converged, 545 cv_result=cv_result, 546 panel=panel, 547 fit_options={ 548 "force": force, 549 "force_int": force_int, 550 "tol": tol, 551 "max_iter": max_iter, 552 "normalize": normalize, 553 "norm_factor": norm_factor, 554 "vartype": vartype, 555 "nboots": nboots, 556 "n_jobs": n_jobs, 557 }, 558 seed=seed, 559 ) 560 561 # Inference 562 if se: 563 result.inference = _run_inference( 564 result, panel, Y_mat, X_mat, W_mat, beta0, Y0, 565 method=method, r_cv=r_cv, lambda_cv=lambda_cv, 566 force_int=force_int, tol=tol, max_iter=max_iter, 567 vartype=vartype, nboots=nboots, alpha=alpha, 568 n_jobs=n_jobs, seed=seed, normalize=normalize, 569 norm_factor=norm_factor, 570 ) 571 572 # Run requested diagnostics at fit time. requested_diag is None when 573 # diagnostics="none". Validation already enforced se=True and 574 # required-config presence. 575 if requested_diag is not None: 576 opts = dict(diagnostics_options or {}) 577 if "loo" in requested_diag: 578 opts["loo"] = True 579 else: 580 opts.setdefault("loo", False) 581 result.diagnostics = _run_diagnostics( 582 result, _requested=requested_diag, **opts, 583 ) 584 585 return result
Estimate counterfactual treatment effects for panel data.
This is the main Python entry point for the counterfactual estimator workflow. Where the paper and the historical R package differ, pyfector defaults to the paper's statistical definition and exposes R-package-style behavior through explicit options.
Missing outcome policy
pyfector distinguishes raw missing outcomes from counterfactual
missingness caused by treatment. Observed untreated cells
(D == 0 and non-missing Y) fit the response surface. Observed
treated cells (D == 1 and non-missing Y) contribute to ATT as
Y - Y_ct. If a treated outcome is missing in the input data, the
model can still produce a counterfactual Y_ct for that cell, but
the cell is not counted in att_avg or att_on because the
treated potential outcome was not observed.
By default, min_T0 is enforced only for treated and reversal
units. Sparse controls are retained if they have at least one
observed outcome, because they may still inform the low-rank response
surface. Set min_T0_strict=True to require controls to satisfy
min_T0 too, matching the more conservative R fect sparse-panel
behavior.
Parameters
data : polars.DataFrame, pandas.DataFrame
Long-format panel data.
Y, D : str
Column names for outcome and binary treatment indicator.
index : (str, str)
Column names for (unit_id, time_period).
X : list of str, optional
Time-varying covariates.
W : str, optional
Observation weight column.
group : str, optional
Reserved for grouped estimation. Currently raises
NotImplementedError when supplied.
method : {"fe", "ife", "mc", "cfe", "both"}
Estimation method.
force : {"none", "unit", "time", "two-way"}
Fixed effects specification.
r : int or (int, int)
Number of factors. If tuple, CV selects from range.
lam : float, optional
Nuclear norm penalty for MC. If None with CV=True, auto-selected.
nlambda : int
Number of automatically generated lambda candidates for MC CV.
lambda_candidates : array-like, optional
Explicit non-negative lambda candidates for MC CV. When supplied,
nlambda is ignored.
CV : bool
If True, cross-validate over r for IFE when r is a tuple,
or over lam for MC when lam is None.
k : int
Number of CV folds.
cv_prop : float
Fraction of eligible observed control cells masked per CV fold.
cv_nobs : int
Number of consecutive within-unit observations to mask as a block.
cv_treat : bool
If True, restrict CV masks to pre-treatment cells of ever-treated
units. If False, use all observed control cells.
cv_donut : int
Exclude this many periods around treatment onset from CV evaluation.
criterion : {"mspe", "gmspe", "mad"}
Cross-validation loss.
cv_rule : {"min", "onepct"}
CV selection rule. "min" chooses the strict minimum-score
candidate and is the paper-faithful default. "onepct" chooses
the simplest candidate within 1% of the best score (lower r for
IFE, higher lam for MC).
se : bool
Compute standard errors via bootstrap/jackknife.
vartype : {"bootstrap", "jackknife"}
Inference method when se=True.
nboots : int
Number of bootstrap replications. Ignored for jackknife.
alpha : float
Significance level for confidence intervals and tests.
tol : float
EM convergence tolerance for final point estimation.
max_iter : int
Maximum EM iterations.
min_T0 : int
Minimum untreated/pre-treatment observed periods. By default this is
enforced only for treated and treatment-reversal units.
min_T0_strict : bool
If True, enforce min_T0 on all units, including controls. This
matches R fect's conservative handling of sparse control rows.
max_missing : float
Maximum missing-outcome fraction per unit, in [0, 1]. Units with
no observed outcomes are always dropped, regardless of this threshold,
because they provide neither fitting information nor observed treated
effects.
normalize : bool
If True, estimate on an outcome standardized by its observed standard
deviation, then transform effects back to the original scale.
Z, Q : list of str, optional
Reserved CFE interaction arguments. Currently raise
NotImplementedError when supplied.
device : {"cpu", "gpu"}
Compute device.
n_jobs : int, optional
Parallel workers for CV and bootstrap. -1 or None uses
all available CPUs.
seed : int, optional
Random seed for full reproducibility.
73@dataclass 74class FectResult: 75 """Container for all fect estimation results.""" 76 # Method info 77 method: str 78 r_cv: int | None = None 79 lambda_cv: float | None = None 80 81 # Point estimates 82 att_avg: float = 0.0 83 att_avg_unit: float = 0.0 84 85 # Dynamic effects 86 att_on: np.ndarray | None = None 87 time_on: np.ndarray | None = None 88 count_on: np.ndarray | None = None 89 90 # Exit effects (treatment reversal) 91 att_off: np.ndarray | None = None 92 time_off: np.ndarray | None = None 93 94 # Coefficients 95 beta: np.ndarray | None = None 96 covariate_names: list[str] = field(default_factory=list) 97 98 # Fixed effects 99 mu: float = 0.0 100 alpha: np.ndarray | None = None # unit FE 101 xi: np.ndarray | None = None # time FE 102 factors: np.ndarray | None = None 103 loadings: np.ndarray | None = None 104 105 # Counterfactual and effects matrices 106 Y_ct: np.ndarray | None = None # T×N counterfactual 107 eff: np.ndarray | None = None # T×N treatment effects 108 residuals: np.ndarray | None = None 109 110 # Model fit 111 sigma2: float = 0.0 112 sigma2_fect: float = 0.0 # additive-FE baseline residual variance 113 IC: float = 0.0 114 PC: float = 0.0 115 rmse: float = 0.0 116 niter: int = 0 117 converged: bool = False 118 119 # Inference 120 inference: InferenceResult | None = None 121 122 # Diagnostics (populated when fect(..., diagnostics="full" | list)) 123 diagnostics: Diagnostics | None = None 124 125 # CV 126 cv_result: CVResult | None = None 127 128 # Panel metadata 129 panel: PanelData | None = None 130 fit_options: dict[str, Any] = field(default_factory=dict) 131 132 # Reproducibility 133 seed: int | None = None 134 135 def summary(self) -> str: 136 """Print summary table of results.""" 137 lines = [] 138 lines.append(f"pyfector estimation results") 139 lines.append(f"{'='*60}") 140 lines.append(f"Method: {self.method}") 141 if self.r_cv is not None: 142 lines.append(f"Number of factors (CV): {self.r_cv}") 143 if self.lambda_cv is not None: 144 lines.append(f"Lambda (CV): {self.lambda_cv:.6f}") 145 lines.append(f"Converged: {self.converged} (iter={self.niter})") 146 lines.append(f"Sigma^2: {self.sigma2:.6f}") 147 lines.append(f"Sigma^2_fect (FE baseline): {self.sigma2_fect:.6f}") 148 lines.append(f"") 149 lines.append(f"ATT (average): {self.att_avg:.6f}") 150 if self.inference is not None: 151 inf = self.inference 152 lines.append(f" SE: {inf.att_avg_se:.6f}") 153 lines.append(f" CI: [{inf.att_avg_ci[0]:.6f}, {inf.att_avg_ci[1]:.6f}]") 154 lines.append(f" p-val: {inf.att_avg_pval:.4f}") 155 156 if self.beta is not None and len(self.beta) > 0: 157 lines.append(f"") 158 lines.append(f"Coefficients:") 159 for i, name in enumerate(self.covariate_names): 160 lines.append(f" {name}: {self.beta[i]:.6f}") 161 162 if self.att_on is not None and self.time_on is not None: 163 lines.append(f"") 164 lines.append(f"Dynamic effects (ATT by relative time):") 165 lines.append(f" {'Time':>6s} {'ATT':>10s} {'Count':>6s}", ) 166 for i, t in enumerate(self.time_on): 167 count = self.count_on[i] if self.count_on is not None else "" 168 att = self.att_on[i] 169 if self.inference is not None: 170 se = self.inference.att_on_se[i] 171 lines.append(f" {t:>6.0f} {att:>10.4f} ({se:.4f}) {count}") 172 else: 173 lines.append(f" {t:>6.0f} {att:>10.4f} {count}") 174 175 lines.append(f"{'='*60}") 176 if self.panel is not None: 177 lines.append(f"N={self.panel.N}, T={self.panel.T}") 178 if self.seed is not None: 179 lines.append(f"Seed: {self.seed}") 180 if self.diagnostics is not None: 181 lines.append("") 182 lines.append(self.diagnostics.summary()) 183 return "\n".join(lines) 184 185 def __repr__(self): 186 return self.summary() 187 188 def plot(self, kind="gap", **kwargs): 189 """Plot results. Shortcut for ``pyfector.plot(self, kind, ...)``.""" 190 from .plotting import plot as _plot 191 return _plot(self, kind=kind, **kwargs) 192 193 def diagnose(self, **kwargs): 194 """Run diagnostic tests. Shortcut for ``pyfector.run_diagnostics(self, ...)``.""" 195 from .diagnostics import run_diagnostics 196 return run_diagnostics(self, **kwargs)
Container for all fect estimation results.
135 def summary(self) -> str: 136 """Print summary table of results.""" 137 lines = [] 138 lines.append(f"pyfector estimation results") 139 lines.append(f"{'='*60}") 140 lines.append(f"Method: {self.method}") 141 if self.r_cv is not None: 142 lines.append(f"Number of factors (CV): {self.r_cv}") 143 if self.lambda_cv is not None: 144 lines.append(f"Lambda (CV): {self.lambda_cv:.6f}") 145 lines.append(f"Converged: {self.converged} (iter={self.niter})") 146 lines.append(f"Sigma^2: {self.sigma2:.6f}") 147 lines.append(f"Sigma^2_fect (FE baseline): {self.sigma2_fect:.6f}") 148 lines.append(f"") 149 lines.append(f"ATT (average): {self.att_avg:.6f}") 150 if self.inference is not None: 151 inf = self.inference 152 lines.append(f" SE: {inf.att_avg_se:.6f}") 153 lines.append(f" CI: [{inf.att_avg_ci[0]:.6f}, {inf.att_avg_ci[1]:.6f}]") 154 lines.append(f" p-val: {inf.att_avg_pval:.4f}") 155 156 if self.beta is not None and len(self.beta) > 0: 157 lines.append(f"") 158 lines.append(f"Coefficients:") 159 for i, name in enumerate(self.covariate_names): 160 lines.append(f" {name}: {self.beta[i]:.6f}") 161 162 if self.att_on is not None and self.time_on is not None: 163 lines.append(f"") 164 lines.append(f"Dynamic effects (ATT by relative time):") 165 lines.append(f" {'Time':>6s} {'ATT':>10s} {'Count':>6s}", ) 166 for i, t in enumerate(self.time_on): 167 count = self.count_on[i] if self.count_on is not None else "" 168 att = self.att_on[i] 169 if self.inference is not None: 170 se = self.inference.att_on_se[i] 171 lines.append(f" {t:>6.0f} {att:>10.4f} ({se:.4f}) {count}") 172 else: 173 lines.append(f" {t:>6.0f} {att:>10.4f} {count}") 174 175 lines.append(f"{'='*60}") 176 if self.panel is not None: 177 lines.append(f"N={self.panel.N}, T={self.panel.T}") 178 if self.seed is not None: 179 lines.append(f"Seed: {self.seed}") 180 if self.diagnostics is not None: 181 lines.append("") 182 lines.append(self.diagnostics.summary()) 183 return "\n".join(lines)
Print summary table of results.
188 def plot(self, kind="gap", **kwargs): 189 """Plot results. Shortcut for ``pyfector.plot(self, kind, ...)``.""" 190 from .plotting import plot as _plot 191 return _plot(self, kind=kind, **kwargs)
Plot results. Shortcut for pyfector.plot(self, kind, ...).
193 def diagnose(self, **kwargs): 194 """Run diagnostic tests. Shortcut for ``pyfector.run_diagnostics(self, ...)``.""" 195 from .diagnostics import run_diagnostics 196 return run_diagnostics(self, **kwargs)
Run diagnostic tests. Shortcut for pyfector.run_diagnostics(self, ...).
35def set_device(device: Literal["cpu", "gpu"]) -> None: 36 """Set the compute device globally.""" 37 global _DEVICE 38 if device == "gpu" and not _check_cupy(): 39 raise ImportError( 40 "CuPy is required for GPU support. Install it with: " 41 "pip install cupy-cuda12x (adjust for your CUDA version)" 42 ) 43 _DEVICE = device
Set the compute device globally.
Return the current device.
355def run_diagnostics( 356 result, 357 f_threshold: float = 0.5, 358 tost_threshold: float = 0.36, 359 placebo_period: tuple[int, int] | None = None, 360 carryover_period: tuple[int, int] | None = None, 361 loo: bool = False, 362 alpha: float = 0.05, 363 *, 364 _requested: list[str] | None = None, 365) -> Diagnostics: 366 """Run diagnostic tests on a FectResult. 367 368 Parameters 369 ---------- 370 result : FectResult 371 Must have inference results (``se=True``). 372 f_threshold : float 373 Non-centrality parameter for equivalence F-test. 374 tost_threshold : float 375 Equivalence bound for TOST. The literal ``0.36`` (default or 376 explicit) triggers Liu et al. (2024)'s scale-aware bound, 377 ``0.36 * sqrt(result.sigma2_fect)``. Any other positive float 378 is taken as an absolute outcome-scale bound. ``None`` is 379 invalid. 380 placebo_period : (start, end), optional 381 Relative pre-treatment window for the holdout/refit placebo test. 382 Observed cells with ``start <= time_on <= end`` and ``time_on < 0`` 383 are removed from the fitting mask, the model is refit using the 384 selected rank/lambda configuration, and the placebo ATT is 385 computed from those withheld cells. For example, ``(-3, 0)`` 386 withholds relative periods ``-3, -2, -1``. 387 carryover_period : (start, end), optional 388 Relative time window for carryover test. 389 loo : bool 390 If True, run leave-one-out post-period sensitivity. 391 alpha : float 392 Significance cutoff used for ``TostResult.all_pass``. 393 """ 394 from scipy import stats 395 396 diag = Diagnostics() 397 398 if result.inference is None: 399 return diag 400 401 inf = result.inference 402 time_on = result.time_on 403 att_on = result.att_on 404 405 if time_on is None or att_on is None: 406 return diag 407 408 pre_mask = time_on < 0 409 if not np.any(pre_mask): 410 return diag 411 412 pre_idx = np.where(pre_mask)[0] 413 k = len(pre_idx) 414 att_pre = att_on[pre_idx] 415 416 requested = set(_requested) if _requested is not None else _VALID_DIAG_NAMES.copy() 417 if placebo_period is not None: 418 placebo_period = _validate_placebo_period(placebo_period) 419 420 needs_tost_threshold = bool({"tost", "placebo"} & requested) 421 threshold_abs = float("nan") 422 threshold_source: str | None = None 423 sigma2_fect_used: float | None = None 424 if needs_tost_threshold: 425 threshold_abs, threshold_source, sigma2_fect_used = _resolve_tost_threshold( 426 result, tost_threshold 427 ) 428 429 # Pre-trend / equivalence F-tests share the bootstrap covariance pass. 430 if ( 431 ("pretrend_f" in requested or "equiv_f" in requested) 432 and inf.att_on_boot is not None 433 ): 434 boot_pre_all = inf.att_on_boot[pre_idx, :] 435 valid_boot = np.all(np.isfinite(boot_pre_all), axis=0) 436 boot_pre = boot_pre_all[:, valid_boot] 437 n_boot = boot_pre.shape[1] 438 if n_boot <= k: 439 boot_pre = None 440 441 else: 442 boot_pre = None 443 444 if boot_pre is not None: 445 S = np.cov(boot_pre) 446 if k == 1: 447 S = np.asarray(S).reshape(1, 1) 448 449 try: 450 cond = np.linalg.cond(S) 451 S_inv = np.linalg.pinv(S) if cond > 1e12 else np.linalg.inv(S) 452 F_raw = float(att_pre @ S_inv @ att_pre) 453 scale = (n_boot - k) / ((n_boot - 1) * k) 454 F_stat = F_raw * scale 455 456 if np.isfinite(F_stat) and F_stat >= 0: 457 if "pretrend_f" in requested: 458 diag.pretrend_f = PretrendFResult( 459 f_stat=F_stat, 460 p_value=float(1 - stats.f.cdf(F_stat, k, n_boot - k)), 461 df1=k, 462 df2=n_boot - k, 463 ) 464 if "equiv_f" in requested: 465 ncp = n_boot * f_threshold 466 diag.equiv_f = EquivFResult( 467 p_value=float(stats.ncf.cdf(F_stat, k, n_boot - k, ncp)), 468 f_threshold=f_threshold, 469 ) 470 except np.linalg.LinAlgError: 471 pass 472 473 # Per-period TOST. 474 if "tost" in requested and inf.att_on_se is not None: 475 se_pre = inf.att_on_se[pre_idx] 476 477 tost_pvals = np.full(k, np.nan) 478 for i in range(k): 479 if se_pre[i] > 0: 480 if inf.att_on_boot is not None: 481 n_boot_i = int(np.isfinite(inf.att_on_boot[pre_idx[i], :]).sum()) 482 else: 483 n_boot_i = 200 484 df = max(n_boot_i - 1, 1) 485 t_upper = (att_pre[i] - threshold_abs) / se_pre[i] 486 t_lower = (att_pre[i] + threshold_abs) / se_pre[i] 487 p_upper = float(stats.t.cdf(t_upper, df)) 488 p_lower = float(1 - stats.t.cdf(t_lower, df)) 489 tost_pvals[i] = max(p_upper, p_lower) 490 491 finite = tost_pvals[np.isfinite(tost_pvals)] 492 max_p = float(finite.max()) if finite.size else float("nan") 493 all_pass = bool(finite.size and (finite < alpha).all()) 494 495 diag.tost = TostResult( 496 pvals=tost_pvals, 497 periods=time_on[pre_idx].copy(), 498 threshold=threshold_abs, 499 threshold_source=threshold_source, 500 sigma2_fect=sigma2_fect_used, 501 max_pval=max_p, 502 all_pass=all_pass, 503 ) 504 505 # Holdout/refit placebo test. This follows R fect's placeboTest path: 506 # selected pre-treatment cells are removed from II before fitting. 507 if "placebo" in requested and placebo_period is not None: 508 diag.placebo = _run_placebo_test( 509 result, 510 placebo_period=placebo_period, 511 threshold_abs=threshold_abs, 512 ) 513 514 # Carryover (estimate only; bootstrap p deferred). 515 if ( 516 "carryover" in requested 517 and carryover_period is not None 518 and hasattr(result, "att_off") 519 and result.att_off is not None 520 ): 521 c_start, c_end = carryover_period 522 time_off = getattr(result, "time_off", None) 523 if time_off is not None: 524 off_mask = (time_off >= c_start) & (time_off <= c_end) 525 if np.any(off_mask): 526 diag.carryover = CarryoverResult( 527 estimate=float(np.mean(result.att_off[off_mask])), 528 period=(int(c_start), int(c_end)), 529 ) 530 531 # Leave-one-out post-period. 532 if "loo" in requested and loo: 533 post_mask = time_on >= 0 534 post_idx = np.where(post_mask)[0] 535 if len(post_idx) > 1: 536 full_att = result.att_avg 537 loo_atts: list[float] = [] 538 loo_periods: list[float] = [] 539 for drop_i in post_idx: 540 remaining = np.delete(post_idx, np.where(post_idx == drop_i)) 541 if len(remaining) > 0: 542 loo_atts.append(float(np.mean(att_on[remaining]))) 543 loo_periods.append(float(time_on[drop_i])) 544 atts_arr = np.array(loo_atts) 545 diag.loo = LooResult( 546 atts=atts_arr, 547 periods=np.array(loo_periods), 548 max_change=float(np.max(np.abs(atts_arr - full_att))), 549 ) 550 551 diag.options = { 552 "tost_threshold": threshold_abs if needs_tost_threshold else None, 553 "tost_threshold_source": threshold_source, 554 "f_threshold": f_threshold, 555 "alpha": alpha, 556 "placebo_period": placebo_period, 557 "carryover_period": carryover_period, 558 "loo": loo, 559 "sigma2_fect": sigma2_fect_used if sigma2_fect_used is not None 560 else getattr(result, "sigma2_fect", None), 561 } 562 try: 563 from . import __version__ as _ver 564 diag.options["pyfector_version"] = _ver 565 except Exception: 566 pass 567 568 return diag
Run diagnostic tests on a FectResult.
Parameters
result : FectResult
Must have inference results (se=True).
f_threshold : float
Non-centrality parameter for equivalence F-test.
tost_threshold : float
Equivalence bound for TOST. The literal 0.36 (default or
explicit) triggers Liu et al. (2024)'s scale-aware bound,
0.36 * sqrt(result.sigma2_fect). Any other positive float
is taken as an absolute outcome-scale bound. None is
invalid.
placebo_period : (start, end), optional
Relative pre-treatment window for the holdout/refit placebo test.
Observed cells with start <= time_on <= end and time_on < 0
are removed from the fitting mask, the model is refit using the
selected rank/lambda configuration, and the placebo ATT is
computed from those withheld cells. For example, (-3, 0)
withholds relative periods -3, -2, -1.
carryover_period : (start, end), optional
Relative time window for carryover test.
loo : bool
If True, run leave-one-out post-period sensitivity.
alpha : float
Significance cutoff used for TostResult.all_pass.
158@dataclass 159class Diagnostics: 160 """Slim-safe registry of diagnostic test results. 161 162 ``tests`` is the future-proof extension point. Built-in tests expose 163 convenience properties for stable user code, but the container itself 164 does not need a new field every time pyfector adds a diagnostic. 165 """ 166 options: dict = field(default_factory=dict) 167 tests: dict[str, Any] = field(default_factory=dict) 168 169 @property 170 def available(self) -> tuple[str, ...]: 171 """Names of populated diagnostics, with built-ins first.""" 172 known = [name for name in _DIAGNOSTIC_SUMMARY_ORDER if name in self.tests] 173 extra = sorted(name for name in self.tests if name not in set(known)) 174 return tuple(known + extra) 175 176 def get(self, name: str, default: Any = None) -> Any: 177 """Return a diagnostic by name, or ``default`` if absent.""" 178 return self.tests.get(name, default) 179 180 def set_test(self, name: str, value: Any | None) -> None: 181 """Set or remove a diagnostic result by name.""" 182 if not name or not isinstance(name, str): 183 raise ValueError("diagnostic name must be a non-empty string") 184 if value is None: 185 self.tests.pop(name, None) 186 else: 187 self.tests[name] = value 188 189 def set(self, name: str, value: Any | None) -> None: 190 """Compatibility shortcut for :meth:`set_test`.""" 191 self.set_test(name, value) 192 193 def __contains__(self, name: str) -> bool: 194 return name in self.tests 195 196 def __getitem__(self, name: str) -> Any: 197 return self.tests[name] 198 199 def _typed(self, name: str, typ: type) -> Any | None: 200 value = self.tests.get(name) 201 return value if isinstance(value, typ) else None 202 203 @property 204 def tost(self) -> TostResult | None: 205 return self._typed("tost", TostResult) 206 207 @tost.setter 208 def tost(self, value: TostResult | None) -> None: 209 self.set_test("tost", value) 210 211 @property 212 def pretrend_f(self) -> PretrendFResult | None: 213 return self._typed("pretrend_f", PretrendFResult) 214 215 @pretrend_f.setter 216 def pretrend_f(self, value: PretrendFResult | None) -> None: 217 self.set_test("pretrend_f", value) 218 219 @property 220 def equiv_f(self) -> EquivFResult | None: 221 return self._typed("equiv_f", EquivFResult) 222 223 @equiv_f.setter 224 def equiv_f(self, value: EquivFResult | None) -> None: 225 self.set_test("equiv_f", value) 226 227 @property 228 def placebo(self) -> PlaceboResult | None: 229 return self._typed("placebo", PlaceboResult) 230 231 @placebo.setter 232 def placebo(self, value: PlaceboResult | None) -> None: 233 self.set_test("placebo", value) 234 235 @property 236 def carryover(self) -> CarryoverResult | None: 237 return self._typed("carryover", CarryoverResult) 238 239 @carryover.setter 240 def carryover(self, value: CarryoverResult | None) -> None: 241 self.set_test("carryover", value) 242 243 @property 244 def loo(self) -> LooResult | None: 245 return self._typed("loo", LooResult) 246 247 @loo.setter 248 def loo(self, value: LooResult | None) -> None: 249 self.set_test("loo", value) 250 251 def summary(self) -> str: 252 lines = ["Diagnostic Tests", "=" * 50] 253 for name in self.available: 254 sub = self.tests[name] 255 if sub is not None: 256 if hasattr(sub, "summary"): 257 lines.append(sub.summary()) 258 else: 259 lines.append(f"{name}: {sub!r}") 260 return "\n".join(lines)
Slim-safe registry of diagnostic test results.
tests is the future-proof extension point. Built-in tests expose
convenience properties for stable user code, but the container itself
does not need a new field every time pyfector adds a diagnostic.
169 @property 170 def available(self) -> tuple[str, ...]: 171 """Names of populated diagnostics, with built-ins first.""" 172 known = [name for name in _DIAGNOSTIC_SUMMARY_ORDER if name in self.tests] 173 extra = sorted(name for name in self.tests if name not in set(known)) 174 return tuple(known + extra)
Names of populated diagnostics, with built-ins first.
176 def get(self, name: str, default: Any = None) -> Any: 177 """Return a diagnostic by name, or ``default`` if absent.""" 178 return self.tests.get(name, default)
Return a diagnostic by name, or default if absent.
180 def set_test(self, name: str, value: Any | None) -> None: 181 """Set or remove a diagnostic result by name.""" 182 if not name or not isinstance(name, str): 183 raise ValueError("diagnostic name must be a non-empty string") 184 if value is None: 185 self.tests.pop(name, None) 186 else: 187 self.tests[name] = value
Set or remove a diagnostic result by name.
189 def set(self, name: str, value: Any | None) -> None: 190 """Compatibility shortcut for :meth:`set_test`.""" 191 self.set_test(name, value)
Compatibility shortcut for set_test().
52@dataclass(frozen=True) 53class TostResult: 54 """Per-period TOST equivalence test on pre-treatment coefficients.""" 55 pvals: np.ndarray 56 periods: np.ndarray 57 threshold: float 58 threshold_source: str 59 sigma2_fect: float | None 60 max_pval: float 61 all_pass: bool 62 63 def summary(self) -> str: 64 lines = [ 65 f"TOST per pre-period (threshold={self.threshold:.4f}, " 66 f"source={self.threshold_source}):" 67 ] 68 for i, t in enumerate(self.periods): 69 p = float(self.pvals[i]) 70 status = "PASS" if p < 0.05 else "fail" 71 lines.append(f" t={t:+.0f}: p={p:.4f} [{status}]") 72 lines.append( 73 f" max p={self.max_pval:.4f} all_pass={self.all_pass}" 74 ) 75 return "\n".join(lines)
Per-period TOST equivalence test on pre-treatment coefficients.
63 def summary(self) -> str: 64 lines = [ 65 f"TOST per pre-period (threshold={self.threshold:.4f}, " 66 f"source={self.threshold_source}):" 67 ] 68 for i, t in enumerate(self.periods): 69 p = float(self.pvals[i]) 70 status = "PASS" if p < 0.05 else "fail" 71 lines.append(f" t={t:+.0f}: p={p:.4f} [{status}]") 72 lines.append( 73 f" max p={self.max_pval:.4f} all_pass={self.all_pass}" 74 ) 75 return "\n".join(lines)
78@dataclass(frozen=True) 79class PretrendFResult: 80 """Joint F-test for all pre-treatment ATTs equal to zero.""" 81 f_stat: float 82 p_value: float 83 df1: int 84 df2: int 85 86 def summary(self) -> str: 87 return ( 88 f"Pre-trend F-test:\n" 89 f" F({self.df1},{self.df2}) = {self.f_stat:.4f}, " 90 f"p = {self.p_value:.4f}" 91 )
Joint F-test for all pre-treatment ATTs equal to zero.
94@dataclass(frozen=True) 95class EquivFResult: 96 """Equivalence F-test (non-central F) for pre-trend bound.""" 97 p_value: float 98 f_threshold: float 99 100 def summary(self) -> str: 101 return ( 102 f"Equivalence F-test:\n" 103 f" p = {self.p_value:.4f} (f_threshold={self.f_threshold})" 104 )
Equivalence F-test (non-central F) for pre-trend bound.
107@dataclass(frozen=True) 108class PlaceboResult: 109 """Holdout/refit placebo ATT with bootstrap and TOST p-values.""" 110 estimate: float 111 se: float 112 p_value: float 113 equiv_p_value: float 114 period: tuple[int, int] 115 n_obs: int 116 n_boot: int 117 118 def summary(self) -> str: 119 return ( 120 f"Placebo test (holdout/refit, window {self.period}):\n" 121 f" ATT={self.estimate:.4f} (SE={self.se:.4f}) " 122 f"p={self.p_value:.4f} equiv_p={self.equiv_p_value:.4f}" 123 )
Holdout/refit placebo ATT with bootstrap and TOST p-values.
126@dataclass(frozen=True) 127class CarryoverResult: 128 """Carryover-window mean ATT (post-reversal). SE/p deferred.""" 129 estimate: float 130 period: tuple[int, int] 131 132 def summary(self) -> str: 133 return ( 134 f"Carryover test (window {self.period}):\n" 135 f" ATT={self.estimate:.4f}" 136 )
Carryover-window mean ATT (post-reversal). SE/p deferred.
139@dataclass(frozen=True) 140class LooResult: 141 """Leave-one-out post-period sensitivity.""" 142 atts: np.ndarray 143 periods: np.ndarray 144 max_change: float 145 146 def summary(self) -> str: 147 return ( 148 f"Leave-one-out:\n" 149 f" Max ATT change = {self.max_change:.6f}" 150 )
Leave-one-out post-period sensitivity.
20def plot( 21 result, 22 kind: Literal["gap", "status", "factors", "counterfactual", "equiv", "calendar"] = "gap", 23 units: list | None = None, 24 show_ci: bool = True, 25 title: str | None = None, 26 figsize: tuple[float, float] = (10, 6), 27 ax=None, 28 **kwargs, 29): 30 """Plot fect results. 31 32 Parameters 33 ---------- 34 result : FectResult 35 Output from ``pyfector.fect()``. 36 kind : str 37 Plot type. 38 units : list, optional 39 Unit IDs for counterfactual plot. 40 show_ci : bool 41 Show confidence intervals (requires ``se=True`` in estimation). 42 """ 43 import matplotlib.pyplot as plt 44 45 if ax is None: 46 fig, ax = plt.subplots(figsize=figsize) 47 else: 48 fig = ax.get_figure() 49 50 if kind == "gap": 51 _plot_gap(result, ax, show_ci, **kwargs) 52 elif kind == "status": 53 _plot_status(result, ax, **kwargs) 54 elif kind == "factors": 55 _plot_factors(result, ax, **kwargs) 56 elif kind == "counterfactual": 57 _plot_counterfactual(result, ax, units, **kwargs) 58 elif kind == "equiv": 59 _plot_equiv(result, ax, **kwargs) 60 elif kind == "calendar": 61 _plot_calendar(result, ax, show_ci, **kwargs) 62 else: 63 raise ValueError(f"Unknown plot kind: {kind}") 64 65 if title: 66 ax.set_title(title) 67 68 fig.tight_layout() 69 return fig
Plot fect results.
Parameters
result : FectResult
Output from pyfector.fect.
kind : str
Plot type.
units : list, optional
Unit IDs for counterfactual plot.
show_ci : bool
Show confidence intervals (requires se=True in estimation).