pyfector.estimators
EM iteration engines for the FE, IFE, MC, and CFE estimators.
Each public estimate_* function takes panel matrices in the shape
used by pyfector.panel and iterates an EM/SQUAREM loop to
convergence, returning fitted values, residuals, and estimated
parameters packed into an EstimationResult.
Structure
estimate_ife()is the dispatch for the FE / IFE family. It picks one of four inner loops depending on whether covariates are present and whether the target rankris zero or positive:_em_fe_ad— additive FE only, no covariates_em_fe_inter— additive + interactive FE, no covariates_em_fe_ad_covar— additive FE + covariates_em_fe_inter_covar— full model: FE + interactive + covariates
estimate_mc()implements the matrix-completion estimator via SQUAREM-accelerated soft singular-value thresholding.estimate_cfe()implements the complex FE estimator withZ * gamma_tandkappa_i * Qinteraction terms.
Correspondence with R fect
The inner loops reproduce the algorithms in R fect's
src/ife.cpp, src/ife_sub.cpp and src/mc.cpp up to the exact
order of operations. In particular:
_em_fe_inter_covar()follows R fect'sinter_fe_ubEM step (two-block coordinate update of covariate coefficientbetaand interactive fixed effects) with an added SQUAREM acceleration step (Varadhan & Roland 2008) that is mathematically equivalent at convergence.estimate_mc()follows R fect'sinter_fe_mcwith SQUAREM added. The soft-thresholded singular values are recovered throughpyfector.linalg.panel_FE().estimate_cfe()mirrors R fect'scomplex_fe_uband cycles over thebeta,gamma,kappa, additive-FE, and interactive-FE blocks.
The CV-selection logic for r (IFE) and lambda (MC) lives in
pyfector.cv and is called before these routines by
pyfector.fect.
1""" 2EM iteration engines for the FE, IFE, MC, and CFE estimators. 3 4Each public ``estimate_*`` function takes panel matrices in the shape 5used by :mod:`pyfector.panel` and iterates an EM/SQUAREM loop to 6convergence, returning fitted values, residuals, and estimated 7parameters packed into an :class:`EstimationResult`. 8 9Structure 10--------- 11* :func:`estimate_ife` is the dispatch for the FE / IFE family. It 12 picks one of four inner loops depending on whether covariates are 13 present and whether the target rank ``r`` is zero or positive: 14 15 - ``_em_fe_ad`` — additive FE only, no covariates 16 - ``_em_fe_inter`` — additive + interactive FE, no covariates 17 - ``_em_fe_ad_covar`` — additive FE + covariates 18 - ``_em_fe_inter_covar`` — full model: FE + interactive + covariates 19 20* :func:`estimate_mc` implements the matrix-completion estimator via 21 SQUAREM-accelerated soft singular-value thresholding. 22* :func:`estimate_cfe` implements the complex FE estimator with 23 ``Z * gamma_t`` and ``kappa_i * Q`` interaction terms. 24 25Correspondence with R fect 26-------------------------- 27The inner loops reproduce the algorithms in R fect's 28``src/ife.cpp``, ``src/ife_sub.cpp`` and ``src/mc.cpp`` up to the exact 29order of operations. In particular: 30 31* :func:`_em_fe_inter_covar` follows R fect's ``inter_fe_ub`` EM step 32 (two-block coordinate update of covariate coefficient ``beta`` and 33 interactive fixed effects) with an added SQUAREM acceleration step 34 (Varadhan & Roland 2008) that is mathematically equivalent at 35 convergence. 36* :func:`estimate_mc` follows R fect's ``inter_fe_mc`` with SQUAREM 37 added. The soft-thresholded singular values are recovered through 38 :func:`pyfector.linalg.panel_FE`. 39* :func:`estimate_cfe` mirrors R fect's ``complex_fe_ub`` and cycles 40 over the ``beta``, ``gamma``, ``kappa``, additive-FE, and 41 interactive-FE blocks. 42 43The CV-selection logic for ``r`` (IFE) and ``lambda`` (MC) lives in 44:mod:`pyfector.cv` and is called before these routines by 45:func:`pyfector.fect`. 46""" 47 48from __future__ import annotations 49 50from typing import NamedTuple 51 52import numpy as np 53 54from .backend import get_backend, to_numpy 55from .linalg import ( 56 ife, 57 panel_beta, 58 panel_factor, 59 panel_FE, 60 compute_xxinv, 61 covar_fit, 62 demean, 63 fe_add, 64) 65 66 67# --------------------------------------------------------------------------- 68# E-step helpers 69# --------------------------------------------------------------------------- 70 71def _e_adj(Y, fit, I, out=None): 72 """Fill missing entries (I==0) with current fit values. 73 74 When *out* is supplied the result is written into it, avoiding a 75 fresh allocation on every EM iteration. 76 """ 77 xp = get_backend() 78 if out is None: 79 return xp.where(I, Y, fit) 80 xp.copyto(out, fit) 81 xp.copyto(out, Y, where=I.astype(bool)) 82 return out 83 84 85def _we_adj(Y, fit, W, I, out=None): 86 """Weighted E-step: observed entries shrunk toward fit by weight. 87 88 When *out* is supplied the result is written into it. 89 """ 90 xp = get_backend() 91 result = xp.where(I, W * Y + (1.0 - W) * fit, fit) 92 if out is not None: 93 xp.copyto(out, result) 94 return out 95 return result 96 97 98def _fe_adj(FE, I): 99 """Zero out entries where I==0 (keep only observed residuals).""" 100 xp = get_backend() 101 return xp.where(I, FE, 0.0) 102 103 104def _active_fe_rank(mask: np.ndarray, force: int) -> int: 105 """Rank of the additive FE design on the observed estimation cells.""" 106 mask = np.asarray(mask, dtype=bool) 107 obs = int(mask.sum()) 108 if obs == 0: 109 return 0 110 if force == 0: 111 return 1 112 113 active_rows = mask.any(axis=1) 114 active_cols = mask.any(axis=0) 115 n_rows = int(active_rows.sum()) 116 n_cols = int(active_cols.sum()) 117 118 if force == 1: 119 return n_cols 120 if force == 2: 121 return n_rows 122 123 # Two-way FE rank is active rows + active columns minus the number of 124 # connected components in the observed row-column bipartite graph. 125 from scipy import sparse 126 from scipy.sparse.csgraph import connected_components 127 128 row_idx, col_idx = np.nonzero(mask) 129 row_map = np.full(mask.shape[0], -1, dtype=np.int64) 130 col_map = np.full(mask.shape[1], -1, dtype=np.int64) 131 row_map[np.where(active_rows)[0]] = np.arange(n_rows) 132 col_map[np.where(active_cols)[0]] = np.arange(n_cols) 133 134 graph_rows = row_map[row_idx] 135 graph_cols = n_rows + col_map[col_idx] 136 rows = np.concatenate([graph_rows, graph_cols]) 137 cols = np.concatenate([graph_cols, graph_rows]) 138 data = np.ones(rows.size, dtype=np.int8) 139 graph = sparse.csr_matrix( 140 (data, (rows, cols)), 141 shape=(n_rows + n_cols, n_rows + n_cols), 142 ) 143 components = int(connected_components(graph, directed=False, return_labels=False)) 144 return n_rows + n_cols - components 145 146 147def _fe_design_sparse(mask: np.ndarray, force: int): 148 """Sparse additive-FE design for observed cells, used for df accounting.""" 149 from scipy import sparse 150 151 mask = np.asarray(mask, dtype=bool) 152 row_idx, col_idx = np.nonzero(mask) 153 n_obs = row_idx.size 154 obs_idx = np.arange(n_obs, dtype=np.int64) 155 156 if force == 0: 157 return sparse.csr_matrix(np.ones((n_obs, 1), dtype=np.float64)) 158 159 if force == 1: 160 active_cols = np.where(mask.any(axis=0))[0] 161 col_map = np.full(mask.shape[1], -1, dtype=np.int64) 162 col_map[active_cols] = np.arange(active_cols.size) 163 return sparse.csr_matrix( 164 (np.ones(n_obs), (obs_idx, col_map[col_idx])), 165 shape=(n_obs, active_cols.size), 166 ) 167 168 if force == 2: 169 active_rows = np.where(mask.any(axis=1))[0] 170 row_map = np.full(mask.shape[0], -1, dtype=np.int64) 171 row_map[active_rows] = np.arange(active_rows.size) 172 return sparse.csr_matrix( 173 (np.ones(n_obs), (obs_idx, row_map[row_idx])), 174 shape=(n_obs, active_rows.size), 175 ) 176 177 active_rows = np.where(mask.any(axis=1))[0] 178 active_cols = np.where(mask.any(axis=0))[0] 179 row_map = np.full(mask.shape[0], -1, dtype=np.int64) 180 col_map = np.full(mask.shape[1], -1, dtype=np.int64) 181 row_map[active_rows] = np.arange(active_rows.size) 182 col_map[active_cols] = np.arange(active_cols.size) 183 rows = np.concatenate([obs_idx, obs_idx]) 184 cols = np.concatenate([ 185 row_map[row_idx], 186 active_rows.size + col_map[col_idx], 187 ]) 188 data = np.ones(rows.size, dtype=np.float64) 189 return sparse.csr_matrix( 190 (data, (rows, cols)), 191 shape=(n_obs, active_rows.size + active_cols.size), 192 ) 193 194 195def _residual_df_r0(I, force: int, X=None) -> int: 196 """Residual df for the additive FE baseline on observed estimation cells. 197 198 The r=0 baseline is the variance source for ``sigma2_fect``. Its 199 denominator should reflect the rank of the design identified by the 200 observed untreated cells, not just a nominal balanced-panel parameter 201 count. 202 """ 203 mask = np.asarray(to_numpy(I), dtype=bool) 204 obs = int(mask.sum()) 205 if obs == 0: 206 return 1 207 208 model_rank = _active_fe_rank(mask, force) 209 210 if X is not None and X.shape[2] > 0: 211 from scipy.sparse.linalg import lsmr 212 213 X_obs = np.asarray(to_numpy(X), dtype=np.float64)[mask, :] 214 finite_cols = np.all(np.isfinite(X_obs), axis=0) 215 if np.any(finite_cols): 216 X_obs = X_obs[:, finite_cols] 217 A = _fe_design_sparse(mask, force) 218 residualized = np.empty_like(X_obs, dtype=np.float64) 219 for j in range(X_obs.shape[1]): 220 sol = lsmr(A, X_obs[:, j], atol=1e-10, btol=1e-10) 221 residualized[:, j] = X_obs[:, j] - A @ sol[0] 222 if residualized.size: 223 model_rank += int(np.linalg.matrix_rank(residualized)) 224 225 return max(obs - model_rank, 1) 226 227 228def _frob_ratio(new, old, W=None, buf=None): 229 """Relative Frobenius norm change: ``||new - old||_F / ||old||_F``. 230 231 Computed via a single square root (``sqrt(diff2 / base2)``) to avoid 232 the double sqrt of earlier versions; the subtraction ``new - old`` 233 is unavoidable but is done only once. On GPU we compute the ratio 234 device-side and sync only the final scalar, which saves a host 235 round-trip per EM iteration. 236 237 When *buf* is supplied (same shape as *new*), the ``new - old`` 238 difference is computed into it to avoid allocating a temporary. 239 """ 240 xp = get_backend() 241 if buf is not None: 242 xp.subtract(new, old, out=buf) 243 delta = buf 244 else: 245 delta = new - old 246 if W is not None: 247 diff2 = xp.sum(W * delta * delta) 248 base2 = xp.sum(W * old * old) + 1e-20 249 else: 250 diff2 = xp.sum(delta * delta) 251 base2 = xp.sum(old * old) + 1e-20 252 return float(xp.sqrt(diff2 / base2)) 253 254 255# --------------------------------------------------------------------------- 256# Result containers 257# --------------------------------------------------------------------------- 258 259class EstimationResult(NamedTuple): 260 beta: np.ndarray | None # (p,) covariate coefficients 261 mu: float # grand mean 262 alpha: np.ndarray | None # (N,) unit FE 263 xi: np.ndarray | None # (T,) time FE 264 factors: np.ndarray | None # (T, r) latent factors 265 loadings: np.ndarray | None # (N, r) factor loadings 266 eigenvalues: np.ndarray | None # (r,) eigenvalues 267 fit: np.ndarray # (T, N) fitted values (counterfactual) 268 residuals: np.ndarray # (T, N) residuals (observed only) 269 sigma2: float # error variance estimate 270 IC: float # information criterion (BIC-type) 271 PC: float # PC criterion 272 niter: int # iterations to convergence 273 converged: bool # whether converged 274 275 276# --------------------------------------------------------------------------- 277# FE / IFE estimator (inter_fe_ub equivalent) 278# --------------------------------------------------------------------------- 279 280def estimate_ife( 281 Y: np.ndarray, # (T, N) outcome 282 Y0: np.ndarray, # (T, N) initial imputation 283 X: np.ndarray | None, # (T, N, p) covariates or None 284 I: np.ndarray, # (T, N) observation indicator 285 W: np.ndarray | None, # (T, N) weights or None 286 beta0: np.ndarray | None, # (p,) initial beta 287 r: int, # number of factors 288 force: int, # FE type: 0/1/2/3 289 tol: float = 1e-5, 290 max_iter: int = 5000, 291) -> EstimationResult: 292 """EM algorithm for interactive fixed effects on unbalanced panels. 293 294 Equivalent to R fect's ``inter_fe_ub()``. 295 """ 296 xp = get_backend() 297 T, N = Y.shape 298 use_weight = W is not None and W.shape == Y.shape 299 has_covar = X is not None and X.shape[2] > 0 300 301 # Normalize weights 302 if use_weight: 303 W = W / (xp.max(W) + 1e-10) 304 WI = W * I 305 else: 306 WI = None 307 308 # Check covariate variation and drop invariant ones. 309 # Vectorised over covariates: |X[:,:,k]|.sum() > tol. 310 valid_x = [] 311 if has_covar: 312 col_sums = xp.abs(X).sum(axis=(0, 1)) # (p,) 313 valid_x = list(np.flatnonzero(to_numpy(col_sums) > 1e-5)) 314 if valid_x: 315 X = X[:, :, valid_x] 316 else: 317 X = None 318 has_covar = False 319 320 p = X.shape[2] if has_covar else 0 321 322 # Precompute (X'X)^-1 323 xxinv = None 324 if has_covar: 325 xxinv = compute_xxinv(X, W=WI if use_weight else None) 326 327 # Initialize 328 beta = beta0[:len(valid_x)] if beta0 is not None and has_covar else ( 329 xp.zeros(p) if has_covar else None 330 ) 331 fit = Y0.copy() 332 333 # Dispatch to appropriate EM loop 334 if not has_covar and r == 0 and force == 0: 335 # Trivial case: just grand mean 336 if use_weight: 337 mu = float(xp.sum(Y * WI) / (xp.sum(WI) + 1e-10)) 338 else: 339 mu = float(xp.sum(Y * I) / (xp.sum(I) + 1e-10)) 340 fit = xp.full_like(Y, mu) 341 residuals = _fe_adj(Y - fit, I) 342 denom = _residual_df_r0(I, force) 343 sigma2 = float(xp.sum(residuals ** 2) / denom) 344 return EstimationResult( 345 beta=None, mu=mu, alpha=None, xi=None, 346 factors=None, loadings=None, eigenvalues=None, 347 fit=fit, residuals=residuals, sigma2=sigma2, 348 IC=0.0, PC=0.0, niter=0, converged=True, 349 ) 350 351 if not has_covar and r > 0: 352 return _em_fe_inter(Y, Y0, I, W, WI, use_weight, force, r, tol, max_iter) 353 elif not has_covar and r == 0: 354 return _em_fe_ad(Y, Y0, I, W, WI, use_weight, force, tol, max_iter) 355 elif has_covar and r == 0: 356 return _em_fe_ad_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, tol, max_iter) 357 else: 358 return _em_fe_inter_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, r, tol, max_iter) 359 360 361# --------------------------------------------------------------------------- 362# EM: additive FE only, no covariates 363# --------------------------------------------------------------------------- 364 365def _em_fe_ad(Y, Y0, I, W, WI, use_weight, force, tol, max_iter): 366 xp = get_backend() 367 T, N = Y.shape 368 fit = Y0.copy() 369 converged = False 370 371 # Pre-allocate reusable buffers for the hot loop. 372 fit_old = xp.empty_like(Y) 373 YY = xp.empty_like(Y) 374 _buf = xp.empty_like(Y) # scratch for _frob_ratio 375 376 for niter in range(1, max_iter + 1): 377 xp.copyto(fit_old, fit) 378 379 # E-step 380 if use_weight: 381 _we_adj(Y, fit, W, I, out=YY) 382 else: 383 _e_adj(Y, fit, I, out=YY) 384 385 # M-step: demean and reconstruct 386 ife_res = ife(YY, force, mc=False, r=0) 387 fit = ife_res.FE 388 389 # Convergence 390 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 391 if dif < tol: 392 converged = True 393 break 394 395 residuals = _fe_adj(Y - fit, I) 396 denom = _residual_df_r0(I, force) 397 sigma2 = float(xp.sum(residuals ** 2) / denom) 398 399 return EstimationResult( 400 beta=None, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 401 factors=None, loadings=None, eigenvalues=None, 402 fit=fit, residuals=residuals, sigma2=sigma2, 403 IC=0.0, PC=0.0, niter=niter, converged=converged, 404 ) 405 406 407# --------------------------------------------------------------------------- 408# EM: additive + interactive FE, no covariates (with burn-in) 409# --------------------------------------------------------------------------- 410 411def _em_fe_inter(Y, Y0, I, W, WI, use_weight, force, r, tol, max_iter): 412 xp = get_backend() 413 T, N = Y.shape 414 fit = Y0.copy() 415 d = min(T, N) 416 converged = False 417 burn_in_done = False 418 niter_total = 0 419 420 # Pre-allocate reusable buffers. 421 fit_old = xp.empty_like(Y) 422 YY = xp.empty_like(Y) 423 _buf = xp.empty_like(Y) 424 425 for phase in range(2): 426 # Phase 0: burn-in (start with high rank, reduce gradually) 427 # Phase 1: final estimation with target r 428 if phase == 0 and not use_weight: 429 # Skip burn-in for unweighted estimation 430 continue 431 432 fit_phase = fit.copy() if phase == 0 else Y0.copy() 433 fit_current = fit_phase.copy() 434 435 for niter in range(1, max_iter + 1): 436 xp.copyto(fit_old, fit_current) 437 438 # E-step 439 if use_weight: 440 _we_adj(Y, fit_current, W, I, out=YY) 441 else: 442 _e_adj(Y, fit_current, I, out=YY) 443 444 # Current rank (burn-in decreases from d) 445 if phase == 0: 446 r_current = max(r, d - niter) 447 else: 448 r_current = r 449 450 # M-step 451 ife_res = ife(YY, force, mc=False, r=r_current) 452 fit_current = ife_res.FE 453 454 # Convergence 455 dif = _frob_ratio(fit_current, fit_old, W=WI if use_weight else None, buf=_buf) 456 niter_total += 1 457 458 if dif < tol: 459 if phase == 0: 460 burn_in_done = True 461 break 462 else: 463 converged = True 464 break 465 466 if phase == 0 and burn_in_done: 467 fit = fit_current 468 continue 469 elif phase == 1: 470 fit = fit_current 471 472 residuals = _fe_adj(Y - fit, I) 473 obs = float(xp.sum(I)) 474 np_param = r * (N + T) - r * r + 1 475 if force in (1, 3): 476 np_param += N - 1 477 if force in (2, 3): 478 np_param += T - 1 479 denom = max(obs - np_param, 1) 480 sigma2 = float(xp.sum(residuals ** 2) / denom) 481 482 IC = _compute_ic(sigma2, np_param, obs) 483 PC = _compute_pc(residuals, np_param, T, N, obs) 484 485 return EstimationResult( 486 beta=None, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 487 factors=ife_res.factors, loadings=ife_res.loadings, 488 eigenvalues=ife_res.eigenvalues, 489 fit=fit, residuals=residuals, sigma2=sigma2, 490 IC=IC, PC=PC, niter=niter_total, converged=converged, 491 ) 492 493 494# --------------------------------------------------------------------------- 495# EM: additive FE + covariates, no interactive FE 496# --------------------------------------------------------------------------- 497 498def _em_fe_ad_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, tol, max_iter): 499 xp = get_backend() 500 T, N = Y.shape 501 p = X.shape[2] 502 fit = Y0.copy() 503 cf = covar_fit(X, beta) 504 FE = fit - cf 505 converged = False 506 507 def _one_em_step(fit_in, FE_in, beta_in): 508 """One complete EM iteration. Returns (fit, FE, beta, ife_res).""" 509 if use_weight: 510 YY = _we_adj(Y, fit_in, W, I) 511 else: 512 YY = _e_adj(Y, fit_in, I) 513 b = panel_beta(X, YY, FE_in, xxinv, W=WI if use_weight else None) 514 c = covar_fit(X, b) 515 if use_weight: 516 U = _we_adj(YY - c, FE_in, W, I) 517 else: 518 U = _e_adj(YY - c, FE_in, I) 519 ires = ife(U, force, mc=False, r=0) 520 return c + ires.FE, ires.FE, b, ires 521 522 # SQUAREM acceleration (Varadhan & Roland 2008). 523 # Every 3rd step we extrapolate; otherwise plain EM. 524 niter = 0 525 ife_res = None 526 _buf = xp.empty_like(Y) # scratch for _frob_ratio 527 while niter < max_iter: 528 fit_old = fit.copy() 529 530 # Step 1: EM 531 fit1, FE1, beta1, ife_res1 = _one_em_step(fit, FE, beta) 532 niter += 1 533 534 dif = _frob_ratio(fit1, fit_old, W=WI if use_weight else None, buf=_buf) 535 if dif < tol: 536 fit, FE, beta, ife_res = fit1, FE1, beta1, ife_res1 537 converged = True 538 break 539 540 # Step 2: second EM 541 fit2, FE2, beta2, ife_res2 = _one_em_step(fit1, FE1, beta1) 542 niter += 1 543 544 # SQUAREM extrapolation 545 r_vec = fit1 - fit 546 v_vec = (fit2 - fit1) - r_vec 547 r_norm2 = float(xp.sum(r_vec ** 2)) 548 v_norm2 = float(xp.sum(v_vec ** 2)) 549 550 if v_norm2 > 1e-30: 551 alpha_sq = -max(1.0, r_norm2 / v_norm2) ** 0.5 552 # Candidate: fit0 - 2*alpha*r + alpha^2*v 553 fit_cand = fit - 2.0 * alpha_sq * r_vec + alpha_sq * alpha_sq * v_vec 554 # One EM "stabilization" step from candidate 555 fit3, FE3, beta3, ife_res3 = _one_em_step(fit_cand, FE2, beta2) 556 niter += 1 557 # Accept if objective improved (fit3 closer to data on observed cells) 558 res3 = float(xp.sum(I * (Y - fit3) ** 2)) 559 res2 = float(xp.sum(I * (Y - fit2) ** 2)) 560 if res3 <= res2: 561 fit, FE, beta, ife_res = fit3, FE3, beta3, ife_res3 562 else: 563 fit, FE, beta, ife_res = fit2, FE2, beta2, ife_res2 564 else: 565 fit, FE, beta, ife_res = fit2, FE2, beta2, ife_res2 566 567 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 568 if dif < tol: 569 converged = True 570 break 571 572 if ife_res is None: 573 # Edge case: max_iter=0 574 _, _, _, ife_res = _one_em_step(fit, FE, beta) 575 576 residuals = _fe_adj(Y - fit, I) 577 denom = _residual_df_r0(I, force, X) 578 sigma2 = float(xp.sum(residuals ** 2) / denom) 579 580 return EstimationResult( 581 beta=beta, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 582 factors=None, loadings=None, eigenvalues=None, 583 fit=fit, residuals=residuals, sigma2=sigma2, 584 IC=0.0, PC=0.0, niter=niter, converged=converged, 585 ) 586 587 588# --------------------------------------------------------------------------- 589# EM: full model — additive FE + interactive FE + covariates (with burn-in) 590# --------------------------------------------------------------------------- 591 592def _em_fe_inter_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, r, tol, max_iter): 593 xp = get_backend() 594 T, N = Y.shape 595 p = X.shape[2] 596 d = min(T, N) 597 cf = covar_fit(X, beta) 598 FE = Y0 - cf 599 fit = Y0.copy() 600 converged = False 601 burn_in_done = False 602 niter_total = 0 603 604 def _one_step(fit_in, FE_in, beta_in, r_cur): 605 if use_weight: 606 YY = _we_adj(Y, fit_in, W, I) 607 else: 608 YY = _e_adj(Y, fit_in, I) 609 b = panel_beta(X, YY, FE_in, xxinv, W=WI if use_weight else None) 610 c = covar_fit(X, b) 611 if use_weight: 612 U = _we_adj(YY - c, FE_in, W, I) 613 else: 614 U = _e_adj(YY - c, FE_in, I) 615 ires = ife(U, force, mc=False, r=r_cur) 616 return c + ires.FE, ires.FE, b, ires 617 618 _buf = xp.empty_like(Y) # scratch for _frob_ratio 619 620 for phase in range(2): 621 if phase == 0 and not use_weight: 622 continue 623 624 fit_current = fit.copy() if phase == 0 else Y0.copy() 625 cf_current = cf.copy() 626 FE_current = FE.copy() 627 beta_current = beta.copy() 628 ife_res = None 629 630 niter = 0 631 while niter < max_iter: 632 fit_old = fit_current.copy() 633 r_current = max(r, d - niter) if phase == 0 else r 634 635 # Step 1 636 fit1, FE1, beta1, ires1 = _one_step(fit_current, FE_current, beta_current, r_current) 637 niter += 1 638 niter_total += 1 639 640 dif = _frob_ratio(fit1, fit_old, W=WI if use_weight else None, buf=_buf) 641 if dif < tol: 642 fit_current, FE_current, beta_current, ife_res = fit1, FE1, beta1, ires1 643 if phase == 0: 644 burn_in_done = True 645 else: 646 converged = True 647 break 648 649 # Step 2 650 fit2, FE2, beta2, ires2 = _one_step(fit1, FE1, beta1, r_current) 651 niter += 1 652 niter_total += 1 653 654 # SQUAREM extrapolation 655 r_vec = fit1 - fit_current 656 v_vec = (fit2 - fit1) - r_vec 657 r_n2 = float(xp.sum(r_vec ** 2)) 658 v_n2 = float(xp.sum(v_vec ** 2)) 659 660 if v_n2 > 1e-30: 661 alpha_sq = -max(1.0, r_n2 / v_n2) ** 0.5 662 fit_cand = fit_current - 2.0 * alpha_sq * r_vec + alpha_sq * alpha_sq * v_vec 663 fit3, FE3, beta3, ires3 = _one_step(fit_cand, FE2, beta2, r_current) 664 niter += 1 665 niter_total += 1 666 res3 = float(xp.sum(I * (Y - fit3) ** 2)) 667 res2 = float(xp.sum(I * (Y - fit2) ** 2)) 668 if res3 <= res2: 669 fit_current, FE_current, beta_current, ife_res = fit3, FE3, beta3, ires3 670 else: 671 fit_current, FE_current, beta_current, ife_res = fit2, FE2, beta2, ires2 672 else: 673 fit_current, FE_current, beta_current, ife_res = fit2, FE2, beta2, ires2 674 675 dif = _frob_ratio(fit_current, fit_old, W=WI if use_weight else None, buf=_buf) 676 if dif < tol: 677 if phase == 0: 678 burn_in_done = True 679 else: 680 converged = True 681 break 682 683 beta = beta_current 684 cf = cf_current if ife_res is None else covar_fit(X, beta_current) 685 FE = FE_current 686 fit = fit_current 687 688 if ife_res is None: 689 _, _, _, ife_res = _one_step(fit, FE, beta, r) 690 691 residuals = _fe_adj(Y - fit, I) 692 obs = float(xp.sum(I)) 693 np_param = r * (N + T) - r * r + p + 1 694 if force in (1, 3): 695 np_param += N - 1 696 if force in (2, 3): 697 np_param += T - 1 698 denom = max(obs - np_param, 1) 699 sigma2 = float(xp.sum(residuals ** 2) / denom) 700 701 IC = _compute_ic(sigma2, np_param, obs) 702 PC = _compute_pc(residuals, np_param, T, N, obs) 703 704 return EstimationResult( 705 beta=beta, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 706 factors=ife_res.factors, loadings=ife_res.loadings, 707 eigenvalues=ife_res.eigenvalues, 708 fit=fit, residuals=residuals, sigma2=sigma2, 709 IC=IC, PC=PC, niter=niter_total, converged=converged, 710 ) 711 712 713# --------------------------------------------------------------------------- 714# MC estimator (inter_fe_mc equivalent) 715# --------------------------------------------------------------------------- 716 717def estimate_mc( 718 Y: np.ndarray, 719 Y0: np.ndarray, 720 X: np.ndarray | None, 721 I: np.ndarray, 722 W: np.ndarray | None, 723 beta0: np.ndarray | None, 724 lam: float, 725 force: int, 726 tol: float = 1e-5, 727 max_iter: int = 5000, 728) -> EstimationResult: 729 """Matrix completion estimator via nuclear norm penalization. 730 731 Equivalent to R fect's ``inter_fe_mc()``. Uses soft singular value 732 thresholding in each EM iteration. 733 """ 734 xp = get_backend() 735 T, N = Y.shape 736 use_weight = W is not None and W.shape == Y.shape 737 has_covar = X is not None and X.shape[2] > 0 738 739 if use_weight: 740 W = W / (xp.max(W) + 1e-10) 741 WI = W * I 742 else: 743 WI = None 744 745 # Check covariate variation 746 valid_x = [] 747 if has_covar: 748 for k in range(X.shape[2]): 749 if xp.sum(xp.abs(X[:, :, k])) > 1e-5: 750 valid_x.append(k) 751 if valid_x: 752 X = X[:, :, valid_x] 753 else: 754 X = None 755 has_covar = False 756 757 p = X.shape[2] if has_covar else 0 758 759 xxinv = None 760 if has_covar: 761 xxinv = compute_xxinv(X, W=WI if use_weight else None) 762 763 beta = beta0[:len(valid_x)] if beta0 is not None and has_covar else ( 764 xp.zeros(p) if has_covar else None 765 ) 766 fit = Y0.copy() 767 cf = covar_fit(X, beta) if has_covar else xp.zeros_like(Y) 768 FE = fit - cf 769 converged = False 770 771 def _one_mc_step(fit_in, FE_in, beta_in, cf_in): 772 if use_weight: 773 YY = _we_adj(Y, fit_in, W, I) 774 else: 775 YY = _e_adj(Y, fit_in, I) 776 b, c = beta_in, cf_in 777 if has_covar: 778 b = panel_beta(X, YY, FE_in, xxinv, W=WI if use_weight else None) 779 c = covar_fit(X, b) 780 if use_weight: 781 U = _we_adj(YY - c, FE_in, W, I) 782 else: 783 U = _e_adj(YY - c, FE_in, I) 784 ires = ife(U, force, mc=True, r=1, hard=False, lam=lam) 785 return c + ires.FE, ires.FE, b, c, ires 786 787 # SQUAREM-accelerated EM 788 niter = 0 789 ife_res = None 790 _buf = xp.empty_like(Y) # scratch for _frob_ratio 791 while niter < max_iter: 792 fit_old = fit.copy() 793 794 fit1, FE1, beta1, cf1, ires1 = _one_mc_step(fit, FE, beta, cf) 795 niter += 1 796 dif = _frob_ratio(fit1, fit_old, W=WI if use_weight else None, buf=_buf) 797 if dif < tol: 798 fit, FE, beta, cf, ife_res = fit1, FE1, beta1, cf1, ires1 799 converged = True 800 break 801 802 fit2, FE2, beta2, cf2, ires2 = _one_mc_step(fit1, FE1, beta1, cf1) 803 niter += 1 804 805 r_vec = fit1 - fit 806 v_vec = (fit2 - fit1) - r_vec 807 r_n2 = float(xp.sum(r_vec ** 2)) 808 v_n2 = float(xp.sum(v_vec ** 2)) 809 810 if v_n2 > 1e-30: 811 alpha_sq = -max(1.0, r_n2 / v_n2) ** 0.5 812 fit_cand = fit - 2.0 * alpha_sq * r_vec + alpha_sq * alpha_sq * v_vec 813 fit3, FE3, beta3, cf3, ires3 = _one_mc_step(fit_cand, FE2, beta2, cf2) 814 niter += 1 815 res3 = float(xp.sum(I * (Y - fit3) ** 2)) 816 res2 = float(xp.sum(I * (Y - fit2) ** 2)) 817 if res3 <= res2: 818 fit, FE, beta, cf, ife_res = fit3, FE3, beta3, cf3, ires3 819 else: 820 fit, FE, beta, cf, ife_res = fit2, FE2, beta2, cf2, ires2 821 else: 822 fit, FE, beta, cf, ife_res = fit2, FE2, beta2, cf2, ires2 823 824 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 825 if dif < tol: 826 converged = True 827 break 828 829 if ife_res is None: 830 _, _, _, _, ife_res = _one_mc_step(fit, FE, beta, cf) 831 832 residuals = _fe_adj(Y - fit, I) 833 obs = float(xp.sum(I)) 834 sigma2 = float(xp.sum(residuals ** 2) / max(obs - 1, 1)) 835 836 return EstimationResult( 837 beta=beta, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 838 factors=ife_res.factors, loadings=ife_res.loadings, 839 eigenvalues=ife_res.eigenvalues, 840 fit=fit, residuals=residuals, sigma2=sigma2, 841 IC=0.0, PC=0.0, niter=niter, converged=converged, 842 ) 843 844 845# --------------------------------------------------------------------------- 846# CFE estimator (complex_fe_ub equivalent) 847# --------------------------------------------------------------------------- 848 849def estimate_cfe( 850 Y: np.ndarray, 851 Y0: np.ndarray, 852 X: np.ndarray | None, 853 I: np.ndarray, 854 W: np.ndarray | None, 855 beta0: np.ndarray | None, 856 r: int, 857 force: int, 858 Z: np.ndarray | None = None, # (N, p_z) unit-invariant covariates 859 Q: np.ndarray | None = None, # (p_q, T) time-invariant covariates 860 gamma_groups: np.ndarray | None = None, # (T,) group assignments for gamma 861 kappa_groups: np.ndarray | None = None, # (N,) group assignments for kappa 862 tol: float = 1e-5, 863 max_iter: int = 5000, 864) -> EstimationResult: 865 """Complex fixed effects estimator with Z*gamma_t and kappa_i*Q. 866 867 Equivalent to R fect's ``complex_fe_ub()`` (simplified: no extra FE dims). 868 Uses block coordinate descent / EM. 869 """ 870 xp = get_backend() 871 T, N = Y.shape 872 use_weight = W is not None and W.shape == Y.shape 873 has_covar = X is not None and X.shape[2] > 0 874 has_Z = Z is not None and Z.shape[1] > 0 875 has_Q = Q is not None and Q.shape[0] > 0 876 877 if use_weight: 878 W = W / (xp.max(W) + 1e-10) 879 WI = W * I 880 else: 881 WI = None 882 883 p = X.shape[2] if has_covar else 0 884 xxinv = compute_xxinv(X, W=WI if use_weight else None) if has_covar else None 885 beta = beta0 if beta0 is not None and has_covar else ( 886 xp.zeros(p) if has_covar else None 887 ) 888 889 # Precompute Z'Z inverse and Q Q' inverse 890 zzinv = xp.linalg.inv(Z.T @ Z) if has_Z else None 891 qqinv = xp.linalg.inv(Q @ Q.T) if has_Q else None 892 893 # Default group assignments (each period/unit is its own group) 894 if gamma_groups is None and has_Z: 895 gamma_groups = xp.arange(T) 896 if kappa_groups is None and has_Q: 897 kappa_groups = xp.arange(N) 898 899 # Initialize fits 900 fit = Y0.copy() 901 cf = covar_fit(X, beta) if has_covar else xp.zeros_like(Y) 902 gamma_fit = xp.zeros_like(Y) 903 kappa_fit = xp.zeros_like(Y) 904 FE_ad = xp.zeros_like(Y) 905 FE_inter = xp.zeros_like(Y) 906 907 converged = False 908 consec_converged = 0 909 _buf = xp.empty_like(Y) # scratch for _frob_ratio 910 911 for niter in range(1, max_iter + 1): 912 fit_old = fit.copy() 913 914 # E-step 915 if use_weight: 916 YY = _we_adj(Y, fit, W, I) 917 else: 918 YY = _e_adj(Y, fit, I) 919 920 # Block 1: beta 921 if has_covar: 922 residual_for_beta = YY - gamma_fit - kappa_fit - FE_ad - FE_inter 923 beta = panel_beta(X, residual_for_beta + cf, FE_ad + FE_inter + gamma_fit + kappa_fit, 924 xxinv, W=WI if use_weight else None) 925 cf = covar_fit(X, beta) 926 927 # Block 2: gamma (time-varying coefficients on Z) 928 if has_Z: 929 E_gamma = YY - cf - kappa_fit - FE_ad - FE_inter 930 gamma_fit = _estimate_gamma(E_gamma, Z, zzinv, gamma_groups, T, N, xp) 931 932 # Block 3: kappa (unit-varying coefficients on Q) 933 if has_Q: 934 E_kappa = YY - cf - gamma_fit - FE_ad - FE_inter 935 kappa_fit = _estimate_kappa(E_kappa, Q, qqinv, kappa_groups, T, N, xp) 936 937 # Block 4: additive FE 938 E_ad = YY - cf - gamma_fit - kappa_fit - FE_inter 939 dm = demean(E_ad, force) 940 FE_ad = fe_add(dm.mu, dm.alpha, dm.xi, T, N, force) 941 942 # Block 5: interactive FE 943 E_inter = YY - cf - gamma_fit - kappa_fit - FE_ad 944 if r > 0: 945 fr = panel_factor(E_inter - xp.mean(E_inter), r) 946 FE_inter = fr.FE 947 else: 948 FE_inter = xp.zeros_like(Y) 949 950 fit = cf + gamma_fit + kappa_fit + FE_ad + FE_inter 951 952 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 953 if dif < tol: 954 consec_converged += 1 955 if consec_converged >= 3: 956 converged = True 957 break 958 else: 959 consec_converged = 0 960 961 residuals = _fe_adj(Y - fit, I) 962 obs = float(xp.sum(I)) 963 sigma2 = float(xp.sum(residuals ** 2) / max(obs - 1, 1)) 964 965 return EstimationResult( 966 beta=beta, mu=float(dm.mu), alpha=dm.alpha, xi=dm.xi, 967 factors=fr.factors if r > 0 else None, 968 loadings=fr.loadings if r > 0 else None, 969 eigenvalues=fr.eigenvalues if r > 0 else None, 970 fit=fit, residuals=residuals, sigma2=sigma2, 971 IC=0.0, PC=0.0, niter=niter, converged=converged, 972 ) 973 974 975def _estimate_gamma(E, Z, zzinv, gamma_groups, T, N, xp): 976 """Estimate time-varying coefficients gamma: Y ≈ Z @ gamma_t.""" 977 unique_groups = xp.unique(gamma_groups) 978 gamma_coefs = xp.zeros((Z.shape[1], T), dtype=E.dtype) 979 for g in unique_groups: 980 g_int = int(g) 981 mask = gamma_groups == g 982 E_sum = xp.sum(E[mask, :], axis=0) # (N,) 983 count = int(xp.sum(mask)) 984 if count > 0: 985 y_bar = E_sum / count 986 gamma_g = zzinv @ Z.T @ y_bar # (p_z,) 987 gamma_coefs[:, mask] = gamma_g[:, None] 988 return (Z @ gamma_coefs).T # (T, N) 989 990 991def _estimate_kappa(E, Q, qqinv, kappa_groups, T, N, xp): 992 """Estimate unit-varying coefficients kappa: Y ≈ kappa_i @ Q.""" 993 unique_groups = xp.unique(kappa_groups) 994 kappa_coefs = xp.zeros((N, Q.shape[0]), dtype=E.dtype) 995 for g in unique_groups: 996 g_int = int(g) 997 mask = kappa_groups == g 998 E_sum = xp.sum(E[:, mask], axis=1) # (T,) 999 count = int(xp.sum(mask)) 1000 if count > 0: 1001 y_bar = E_sum / count 1002 kappa_g = qqinv @ Q @ y_bar # (p_q,) 1003 kappa_coefs[mask, :] = kappa_g[None, :] 1004 return (kappa_coefs @ Q).T if Q.shape[0] > 0 else xp.zeros((T, N), dtype=E.dtype) 1005 1006 1007# --------------------------------------------------------------------------- 1008# IC / PC helpers 1009# --------------------------------------------------------------------------- 1010 1011def _compute_ic(sigma2: float, np_param: int, obs: float) -> float: 1012 """BIC-type information criterion.""" 1013 import math 1014 if sigma2 <= 0: 1015 return float("inf") 1016 return math.log(sigma2) + np_param * math.log(obs) / obs 1017 1018 1019def _compute_pc(residuals, np_param: int, T: int, N: int, obs: float) -> float: 1020 """PC criterion (Li 2018).""" 1021 xp = get_backend() 1022 mse = float(xp.sum(residuals ** 2) / max(obs, 1)) 1023 # Penalty term 1024 penalty = np_param * mse * (T + N) / (T * N) if T * N > 0 else 0 1025 return mse + penalty
class
EstimationResult(typing.NamedTuple):
260class EstimationResult(NamedTuple): 261 beta: np.ndarray | None # (p,) covariate coefficients 262 mu: float # grand mean 263 alpha: np.ndarray | None # (N,) unit FE 264 xi: np.ndarray | None # (T,) time FE 265 factors: np.ndarray | None # (T, r) latent factors 266 loadings: np.ndarray | None # (N, r) factor loadings 267 eigenvalues: np.ndarray | None # (r,) eigenvalues 268 fit: np.ndarray # (T, N) fitted values (counterfactual) 269 residuals: np.ndarray # (T, N) residuals (observed only) 270 sigma2: float # error variance estimate 271 IC: float # information criterion (BIC-type) 272 PC: float # PC criterion 273 niter: int # iterations to convergence 274 converged: bool # whether converged
EstimationResult(beta, mu, alpha, xi, factors, loadings, eigenvalues, fit, residuals, sigma2, IC, PC, niter, converged)
EstimationResult( beta: ForwardRef('np.ndarray | None'), mu: float, alpha: ForwardRef('np.ndarray | None'), xi: ForwardRef('np.ndarray | None'), factors: ForwardRef('np.ndarray | None'), loadings: ForwardRef('np.ndarray | None'), eigenvalues: ForwardRef('np.ndarray | None'), fit: ForwardRef('np.ndarray'), residuals: ForwardRef('np.ndarray'), sigma2: float, IC: float, PC: float, niter: int, converged: bool)
Create new instance of EstimationResult(beta, mu, alpha, xi, factors, loadings, eigenvalues, fit, residuals, sigma2, IC, PC, niter, converged)
def
estimate_ife( Y: numpy.ndarray, Y0: numpy.ndarray, X: numpy.ndarray | None, I: numpy.ndarray, W: numpy.ndarray | None, beta0: numpy.ndarray | None, r: int, force: int, tol: float = 1e-05, max_iter: int = 5000) -> EstimationResult:
281def estimate_ife( 282 Y: np.ndarray, # (T, N) outcome 283 Y0: np.ndarray, # (T, N) initial imputation 284 X: np.ndarray | None, # (T, N, p) covariates or None 285 I: np.ndarray, # (T, N) observation indicator 286 W: np.ndarray | None, # (T, N) weights or None 287 beta0: np.ndarray | None, # (p,) initial beta 288 r: int, # number of factors 289 force: int, # FE type: 0/1/2/3 290 tol: float = 1e-5, 291 max_iter: int = 5000, 292) -> EstimationResult: 293 """EM algorithm for interactive fixed effects on unbalanced panels. 294 295 Equivalent to R fect's ``inter_fe_ub()``. 296 """ 297 xp = get_backend() 298 T, N = Y.shape 299 use_weight = W is not None and W.shape == Y.shape 300 has_covar = X is not None and X.shape[2] > 0 301 302 # Normalize weights 303 if use_weight: 304 W = W / (xp.max(W) + 1e-10) 305 WI = W * I 306 else: 307 WI = None 308 309 # Check covariate variation and drop invariant ones. 310 # Vectorised over covariates: |X[:,:,k]|.sum() > tol. 311 valid_x = [] 312 if has_covar: 313 col_sums = xp.abs(X).sum(axis=(0, 1)) # (p,) 314 valid_x = list(np.flatnonzero(to_numpy(col_sums) > 1e-5)) 315 if valid_x: 316 X = X[:, :, valid_x] 317 else: 318 X = None 319 has_covar = False 320 321 p = X.shape[2] if has_covar else 0 322 323 # Precompute (X'X)^-1 324 xxinv = None 325 if has_covar: 326 xxinv = compute_xxinv(X, W=WI if use_weight else None) 327 328 # Initialize 329 beta = beta0[:len(valid_x)] if beta0 is not None and has_covar else ( 330 xp.zeros(p) if has_covar else None 331 ) 332 fit = Y0.copy() 333 334 # Dispatch to appropriate EM loop 335 if not has_covar and r == 0 and force == 0: 336 # Trivial case: just grand mean 337 if use_weight: 338 mu = float(xp.sum(Y * WI) / (xp.sum(WI) + 1e-10)) 339 else: 340 mu = float(xp.sum(Y * I) / (xp.sum(I) + 1e-10)) 341 fit = xp.full_like(Y, mu) 342 residuals = _fe_adj(Y - fit, I) 343 denom = _residual_df_r0(I, force) 344 sigma2 = float(xp.sum(residuals ** 2) / denom) 345 return EstimationResult( 346 beta=None, mu=mu, alpha=None, xi=None, 347 factors=None, loadings=None, eigenvalues=None, 348 fit=fit, residuals=residuals, sigma2=sigma2, 349 IC=0.0, PC=0.0, niter=0, converged=True, 350 ) 351 352 if not has_covar and r > 0: 353 return _em_fe_inter(Y, Y0, I, W, WI, use_weight, force, r, tol, max_iter) 354 elif not has_covar and r == 0: 355 return _em_fe_ad(Y, Y0, I, W, WI, use_weight, force, tol, max_iter) 356 elif has_covar and r == 0: 357 return _em_fe_ad_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, tol, max_iter) 358 else: 359 return _em_fe_inter_covar(Y, Y0, X, I, W, WI, use_weight, beta, xxinv, force, r, tol, max_iter)
EM algorithm for interactive fixed effects on unbalanced panels.
Equivalent to R fect's inter_fe_ub().
def
estimate_mc( Y: numpy.ndarray, Y0: numpy.ndarray, X: numpy.ndarray | None, I: numpy.ndarray, W: numpy.ndarray | None, beta0: numpy.ndarray | None, lam: float, force: int, tol: float = 1e-05, max_iter: int = 5000) -> EstimationResult:
718def estimate_mc( 719 Y: np.ndarray, 720 Y0: np.ndarray, 721 X: np.ndarray | None, 722 I: np.ndarray, 723 W: np.ndarray | None, 724 beta0: np.ndarray | None, 725 lam: float, 726 force: int, 727 tol: float = 1e-5, 728 max_iter: int = 5000, 729) -> EstimationResult: 730 """Matrix completion estimator via nuclear norm penalization. 731 732 Equivalent to R fect's ``inter_fe_mc()``. Uses soft singular value 733 thresholding in each EM iteration. 734 """ 735 xp = get_backend() 736 T, N = Y.shape 737 use_weight = W is not None and W.shape == Y.shape 738 has_covar = X is not None and X.shape[2] > 0 739 740 if use_weight: 741 W = W / (xp.max(W) + 1e-10) 742 WI = W * I 743 else: 744 WI = None 745 746 # Check covariate variation 747 valid_x = [] 748 if has_covar: 749 for k in range(X.shape[2]): 750 if xp.sum(xp.abs(X[:, :, k])) > 1e-5: 751 valid_x.append(k) 752 if valid_x: 753 X = X[:, :, valid_x] 754 else: 755 X = None 756 has_covar = False 757 758 p = X.shape[2] if has_covar else 0 759 760 xxinv = None 761 if has_covar: 762 xxinv = compute_xxinv(X, W=WI if use_weight else None) 763 764 beta = beta0[:len(valid_x)] if beta0 is not None and has_covar else ( 765 xp.zeros(p) if has_covar else None 766 ) 767 fit = Y0.copy() 768 cf = covar_fit(X, beta) if has_covar else xp.zeros_like(Y) 769 FE = fit - cf 770 converged = False 771 772 def _one_mc_step(fit_in, FE_in, beta_in, cf_in): 773 if use_weight: 774 YY = _we_adj(Y, fit_in, W, I) 775 else: 776 YY = _e_adj(Y, fit_in, I) 777 b, c = beta_in, cf_in 778 if has_covar: 779 b = panel_beta(X, YY, FE_in, xxinv, W=WI if use_weight else None) 780 c = covar_fit(X, b) 781 if use_weight: 782 U = _we_adj(YY - c, FE_in, W, I) 783 else: 784 U = _e_adj(YY - c, FE_in, I) 785 ires = ife(U, force, mc=True, r=1, hard=False, lam=lam) 786 return c + ires.FE, ires.FE, b, c, ires 787 788 # SQUAREM-accelerated EM 789 niter = 0 790 ife_res = None 791 _buf = xp.empty_like(Y) # scratch for _frob_ratio 792 while niter < max_iter: 793 fit_old = fit.copy() 794 795 fit1, FE1, beta1, cf1, ires1 = _one_mc_step(fit, FE, beta, cf) 796 niter += 1 797 dif = _frob_ratio(fit1, fit_old, W=WI if use_weight else None, buf=_buf) 798 if dif < tol: 799 fit, FE, beta, cf, ife_res = fit1, FE1, beta1, cf1, ires1 800 converged = True 801 break 802 803 fit2, FE2, beta2, cf2, ires2 = _one_mc_step(fit1, FE1, beta1, cf1) 804 niter += 1 805 806 r_vec = fit1 - fit 807 v_vec = (fit2 - fit1) - r_vec 808 r_n2 = float(xp.sum(r_vec ** 2)) 809 v_n2 = float(xp.sum(v_vec ** 2)) 810 811 if v_n2 > 1e-30: 812 alpha_sq = -max(1.0, r_n2 / v_n2) ** 0.5 813 fit_cand = fit - 2.0 * alpha_sq * r_vec + alpha_sq * alpha_sq * v_vec 814 fit3, FE3, beta3, cf3, ires3 = _one_mc_step(fit_cand, FE2, beta2, cf2) 815 niter += 1 816 res3 = float(xp.sum(I * (Y - fit3) ** 2)) 817 res2 = float(xp.sum(I * (Y - fit2) ** 2)) 818 if res3 <= res2: 819 fit, FE, beta, cf, ife_res = fit3, FE3, beta3, cf3, ires3 820 else: 821 fit, FE, beta, cf, ife_res = fit2, FE2, beta2, cf2, ires2 822 else: 823 fit, FE, beta, cf, ife_res = fit2, FE2, beta2, cf2, ires2 824 825 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 826 if dif < tol: 827 converged = True 828 break 829 830 if ife_res is None: 831 _, _, _, _, ife_res = _one_mc_step(fit, FE, beta, cf) 832 833 residuals = _fe_adj(Y - fit, I) 834 obs = float(xp.sum(I)) 835 sigma2 = float(xp.sum(residuals ** 2) / max(obs - 1, 1)) 836 837 return EstimationResult( 838 beta=beta, mu=ife_res.mu, alpha=ife_res.alpha, xi=ife_res.xi, 839 factors=ife_res.factors, loadings=ife_res.loadings, 840 eigenvalues=ife_res.eigenvalues, 841 fit=fit, residuals=residuals, sigma2=sigma2, 842 IC=0.0, PC=0.0, niter=niter, converged=converged, 843 )
Matrix completion estimator via nuclear norm penalization.
Equivalent to R fect's inter_fe_mc(). Uses soft singular value
thresholding in each EM iteration.
def
estimate_cfe( Y: numpy.ndarray, Y0: numpy.ndarray, X: numpy.ndarray | None, I: numpy.ndarray, W: numpy.ndarray | None, beta0: numpy.ndarray | None, r: int, force: int, Z: numpy.ndarray | None = None, Q: numpy.ndarray | None = None, gamma_groups: numpy.ndarray | None = None, kappa_groups: numpy.ndarray | None = None, tol: float = 1e-05, max_iter: int = 5000) -> EstimationResult:
850def estimate_cfe( 851 Y: np.ndarray, 852 Y0: np.ndarray, 853 X: np.ndarray | None, 854 I: np.ndarray, 855 W: np.ndarray | None, 856 beta0: np.ndarray | None, 857 r: int, 858 force: int, 859 Z: np.ndarray | None = None, # (N, p_z) unit-invariant covariates 860 Q: np.ndarray | None = None, # (p_q, T) time-invariant covariates 861 gamma_groups: np.ndarray | None = None, # (T,) group assignments for gamma 862 kappa_groups: np.ndarray | None = None, # (N,) group assignments for kappa 863 tol: float = 1e-5, 864 max_iter: int = 5000, 865) -> EstimationResult: 866 """Complex fixed effects estimator with Z*gamma_t and kappa_i*Q. 867 868 Equivalent to R fect's ``complex_fe_ub()`` (simplified: no extra FE dims). 869 Uses block coordinate descent / EM. 870 """ 871 xp = get_backend() 872 T, N = Y.shape 873 use_weight = W is not None and W.shape == Y.shape 874 has_covar = X is not None and X.shape[2] > 0 875 has_Z = Z is not None and Z.shape[1] > 0 876 has_Q = Q is not None and Q.shape[0] > 0 877 878 if use_weight: 879 W = W / (xp.max(W) + 1e-10) 880 WI = W * I 881 else: 882 WI = None 883 884 p = X.shape[2] if has_covar else 0 885 xxinv = compute_xxinv(X, W=WI if use_weight else None) if has_covar else None 886 beta = beta0 if beta0 is not None and has_covar else ( 887 xp.zeros(p) if has_covar else None 888 ) 889 890 # Precompute Z'Z inverse and Q Q' inverse 891 zzinv = xp.linalg.inv(Z.T @ Z) if has_Z else None 892 qqinv = xp.linalg.inv(Q @ Q.T) if has_Q else None 893 894 # Default group assignments (each period/unit is its own group) 895 if gamma_groups is None and has_Z: 896 gamma_groups = xp.arange(T) 897 if kappa_groups is None and has_Q: 898 kappa_groups = xp.arange(N) 899 900 # Initialize fits 901 fit = Y0.copy() 902 cf = covar_fit(X, beta) if has_covar else xp.zeros_like(Y) 903 gamma_fit = xp.zeros_like(Y) 904 kappa_fit = xp.zeros_like(Y) 905 FE_ad = xp.zeros_like(Y) 906 FE_inter = xp.zeros_like(Y) 907 908 converged = False 909 consec_converged = 0 910 _buf = xp.empty_like(Y) # scratch for _frob_ratio 911 912 for niter in range(1, max_iter + 1): 913 fit_old = fit.copy() 914 915 # E-step 916 if use_weight: 917 YY = _we_adj(Y, fit, W, I) 918 else: 919 YY = _e_adj(Y, fit, I) 920 921 # Block 1: beta 922 if has_covar: 923 residual_for_beta = YY - gamma_fit - kappa_fit - FE_ad - FE_inter 924 beta = panel_beta(X, residual_for_beta + cf, FE_ad + FE_inter + gamma_fit + kappa_fit, 925 xxinv, W=WI if use_weight else None) 926 cf = covar_fit(X, beta) 927 928 # Block 2: gamma (time-varying coefficients on Z) 929 if has_Z: 930 E_gamma = YY - cf - kappa_fit - FE_ad - FE_inter 931 gamma_fit = _estimate_gamma(E_gamma, Z, zzinv, gamma_groups, T, N, xp) 932 933 # Block 3: kappa (unit-varying coefficients on Q) 934 if has_Q: 935 E_kappa = YY - cf - gamma_fit - FE_ad - FE_inter 936 kappa_fit = _estimate_kappa(E_kappa, Q, qqinv, kappa_groups, T, N, xp) 937 938 # Block 4: additive FE 939 E_ad = YY - cf - gamma_fit - kappa_fit - FE_inter 940 dm = demean(E_ad, force) 941 FE_ad = fe_add(dm.mu, dm.alpha, dm.xi, T, N, force) 942 943 # Block 5: interactive FE 944 E_inter = YY - cf - gamma_fit - kappa_fit - FE_ad 945 if r > 0: 946 fr = panel_factor(E_inter - xp.mean(E_inter), r) 947 FE_inter = fr.FE 948 else: 949 FE_inter = xp.zeros_like(Y) 950 951 fit = cf + gamma_fit + kappa_fit + FE_ad + FE_inter 952 953 dif = _frob_ratio(fit, fit_old, W=WI if use_weight else None, buf=_buf) 954 if dif < tol: 955 consec_converged += 1 956 if consec_converged >= 3: 957 converged = True 958 break 959 else: 960 consec_converged = 0 961 962 residuals = _fe_adj(Y - fit, I) 963 obs = float(xp.sum(I)) 964 sigma2 = float(xp.sum(residuals ** 2) / max(obs - 1, 1)) 965 966 return EstimationResult( 967 beta=beta, mu=float(dm.mu), alpha=dm.alpha, xi=dm.xi, 968 factors=fr.factors if r > 0 else None, 969 loadings=fr.loadings if r > 0 else None, 970 eigenvalues=fr.eigenvalues if r > 0 else None, 971 fit=fit, residuals=residuals, sigma2=sigma2, 972 IC=0.0, PC=0.0, niter=niter, converged=converged, 973 )
Complex fixed effects estimator with Zgamma_t and kappa_iQ.
Equivalent to R fect's complex_fe_ub() (simplified: no extra FE dims).
Uses block coordinate descent / EM.