pyfector.linalg

  1"""
  2Core linear algebra primitives.
  3
  4This module hosts the numerical kernels that the EM engines in
  5:mod:`pyfector.estimators` call on every iteration: truncated SVD, panel
  6factor extraction, nuclear-norm thresholding, additive/interactive
  7fixed-effect assembly, and the partialled-out OLS ``beta`` step.
  8
  9Relationship to R fect
 10----------------------
 11* :func:`panel_factor` matches R fect's ``panel_factor`` (``fe_sub.cpp``):
 12  it eigendecomposes the smaller of the two Gram matrices
 13  ``E @ E.T / (N*T)`` or ``E.T @ E / (N*T)`` and recovers the requested
 14  number of factors.  When ``r`` is small relative to ``min(T, N)``
 15  pyfector dispatches the top-``r`` factor extraction through the
 16  randomised :func:`truncated_svd`, which gives the same subspace up to
 17  rotation at O(T·N·r) cost — a big saving when ``min(T, N)`` is large.
 18  The identification ``factors = sqrt(T) * U_r``,
 19  ``loadings = E.T @ factors / T`` is preserved either way.
 20* :func:`panel_FE` matches R fect's soft/hard singular-value
 21  thresholding operator for the matrix-completion estimator.
 22* :func:`demean` and :func:`fe_add` implement the ``force`` conventions
 23  used throughout R fect: 0 = none, 1 = unit, 2 = time, 3 = two-way.
 24* :func:`panel_beta` / :func:`compute_xxinv` replicate R fect's
 25  ``panel_beta`` / ``(w)XXinv`` (``auxiliary.cpp``) using single einsum
 26  calls instead of Python loops over covariates.
 27
 28Backend abstraction
 29-------------------
 30All functions call :func:`pyfector.backend.get_backend` so they work
 31transparently on CPU (NumPy) and GPU (CuPy).  Host-sync scalar
 32conversions (``float(...)``) are avoided in the inner loops so the GPU
 33path does not stall every iteration.
 34"""
 35
 36from __future__ import annotations
 37
 38from typing import Literal, NamedTuple
 39
 40import numpy as np
 41
 42from .backend import get_backend, to_numpy
 43
 44
 45# ---------------------------------------------------------------------------
 46# SVD helpers
 47# ---------------------------------------------------------------------------
 48
 49class SVDResult(NamedTuple):
 50    """Truncated SVD result: U (T,r), s (r,), Vt (r,N)."""
 51    U: np.ndarray
 52    s: np.ndarray
 53    Vt: np.ndarray
 54
 55
 56def truncated_svd(M, r: int, n_oversamples: int = 10, n_iter: int = 4) -> SVDResult:
 57    """Randomized truncated SVD of *M* (T×N), returning top-*r* components.
 58
 59    Uses the Halko-Martinsson-Tropp algorithm (2011).  Cost is
 60    O(T·N·(r + n_oversamples)) — dramatically cheaper than a full SVD when
 61    r << min(T, N).
 62
 63    Falls back to exact truncated SVD when r is close to min(T, N).
 64    """
 65    xp = get_backend()
 66    T, N = M.shape
 67    k = min(r, T, N)
 68
 69    # For small matrices or when r is large, use exact SVD (more stable)
 70    if k >= min(T, N) * 0.5 or min(T, N) <= 50 or k >= min(T, N) - 1:
 71        return _exact_truncated_svd(M, k)
 72
 73    # Randomized SVD (Halko et al. 2011)
 74    p = min(k + n_oversamples, min(T, N))
 75
 76    # Stage A: form an approximate orthonormal basis Q for range(M)
 77    Omega = xp.random.standard_normal((N, p)).astype(M.dtype)
 78    Y = M @ Omega  # (T, p)
 79
 80    # Power iteration with QR re-orthogonalization (prevents overflow)
 81    for _ in range(n_iter):
 82        Z = M.T @ Y           # (N, p)
 83        Z, _ = xp.linalg.qr(Z)
 84        Y = M @ Z             # (T, p)
 85        Y, _ = xp.linalg.qr(Y)
 86    Q = Y  # already orthonormal from last QR
 87
 88    # Stage B: form B = Q' M, then SVD of the small matrix B
 89    B = Q.T @ M  # (p, N)
 90    Uhat, s, Vt = xp.linalg.svd(B, full_matrices=False)
 91    U = Q @ Uhat  # (T, p)
 92
 93    return SVDResult(U[:, :k], s[:k], Vt[:k, :])
 94
 95
 96def _exact_truncated_svd(M, k: int) -> SVDResult:
 97    """Full SVD truncated to top-k.  Used as fallback."""
 98    xp = get_backend()
 99    U, s, Vt = xp.linalg.svd(M, full_matrices=False)
100    return SVDResult(U[:, :k], s[:k], Vt[:k, :])
101
102
103def full_svd(M):
104    """Full economy SVD.  Use sparingly — prefer truncated_svd."""
105    xp = get_backend()
106    return xp.linalg.svd(M, full_matrices=False)
107
108
109def _gram_svd(M):
110    """Economy SVD through the smaller Gram matrix for rectangular panels."""
111    xp = get_backend()
112    T, N = M.shape
113    dtype = M.dtype
114    eps = xp.finfo(dtype).eps if xp.issubdtype(dtype, xp.floating) else xp.finfo(xp.float64).eps
115
116    if T <= N:
117        vals, U = xp.linalg.eigh(M @ M.T)
118        vals = xp.maximum(vals[::-1], 0.0)
119        U = U[:, ::-1]
120        s = xp.sqrt(vals)
121        nonzero = s > eps * max(T, N) * (s[0] if s.size else 1.0)
122        Vt = xp.zeros((T, N), dtype=dtype)
123        if xp.any(nonzero):
124            Vt[nonzero, :] = (U[:, nonzero].T @ M) / s[nonzero, None]
125        return U, s, Vt
126
127    vals, V = xp.linalg.eigh(M.T @ M)
128    vals = xp.maximum(vals[::-1], 0.0)
129    V = V[:, ::-1]
130    s = xp.sqrt(vals)
131    nonzero = s > eps * max(T, N) * (s[0] if s.size else 1.0)
132    U = xp.zeros((T, N), dtype=dtype)
133    if xp.any(nonzero):
134        U[:, nonzero] = (M @ V[:, nonzero]) / s[nonzero][None, :]
135    return U, s, V.T
136
137
138# ---------------------------------------------------------------------------
139# Panel factor extraction  (equivalent to R fect's panel_factor)
140# ---------------------------------------------------------------------------
141
142class FactorResult(NamedTuple):
143    factors: np.ndarray   # (T, r)
144    loadings: np.ndarray  # (N, r)
145    eigenvalues: np.ndarray  # (r,)
146    FE: np.ndarray        # (T, N) = factors @ loadings.T
147
148
149def panel_factor(E, r: int) -> FactorResult:
150    """Extract ``r`` latent factors and loadings from the ``(T, N)`` error
151    matrix ``E``.
152
153    Matches R fect's ``panel_factor`` identification exactly:
154
155    * factors ``F = sqrt(T) * U_r`` where ``U_r`` is the top-``r`` left
156      singular vectors of ``E @ E.T / (N*T)`` (or equivalently of
157      ``E / sqrt(N*T)``).
158    * loadings ``Lambda = E.T @ F / T``.
159    * ``FE = F @ Lambda.T``.
160
161    When ``r`` is small compared to ``min(T, N)``, pyfector extracts the
162    top-``r`` singular vectors via :func:`truncated_svd` (randomised
163    SVD), which is O(T·N·(r + oversample)) rather than the
164    O(min(T, N)^3) cost of forming the full Gram matrix.  The resulting
165    factors/loadings span the same ``r``-dimensional subspace as the
166    full-SVD path; the only difference is numerical rotation within the
167    subspace, which is immaterial for downstream fit and treatment-effect
168    estimates.  For small problems, or when ``r`` is close to the full
169    rank, the exact path is still used (see the ``use_rsvd`` branch
170    below).
171    """
172    xp = get_backend()
173    T, N = E.shape
174    if r == 0:
175        return FactorResult(
176            xp.empty((T, 0)),
177            xp.empty((N, 0)),
178            xp.empty((0,)),
179            xp.zeros((T, N), dtype=E.dtype),
180        )
181
182    # Decide between randomised truncated SVD and exact full SVD.
183    # The randomised path pays off when the target rank is well below the
184    # full matrix rank.  We use it when r <= min(T, N) / 4 and the
185    # matrix is large enough that the constant factors matter.
186    min_dim = min(T, N)
187    use_rsvd = (r * 4 <= min_dim) and (min_dim > 50)
188
189    if use_rsvd:
190        # Top-r left singular vectors of E match the top-r eigenvectors
191        # of E E'; scaled by sqrt(T) they equal pyfector/fect factors.
192        svd = truncated_svd(E, r=r)
193        U = svd.U                                  # (T, r)
194        s = svd.s                                  # (r,)
195        factors = U * xp.sqrt(xp.asarray(T, dtype=E.dtype))        # (T, r)
196        loadings = E.T @ factors / T                                 # (N, r)
197        # Eigenvalues of E E' / (N T) are s^2 / (N T).
198        eigenvalues = (s * s) / (N * T)
199    elif T < N:
200        EE = E @ E.T / (N * T)
201        U, s, _ = xp.linalg.svd(EE, full_matrices=False)
202        factors = U[:, :r] * xp.sqrt(xp.asarray(T, dtype=E.dtype))  # (T, r)
203        loadings = E.T @ factors / T                                  # (N, r)
204        eigenvalues = s[:r]
205    else:
206        EE = E.T @ E / (N * T)
207        U, s, _ = xp.linalg.svd(EE, full_matrices=False)
208        loadings = U[:, :r] * xp.sqrt(xp.asarray(N, dtype=E.dtype))  # (N, r)
209        factors = E @ loadings / N                                    # (T, r)
210        eigenvalues = s[:r]
211
212    FE = factors @ loadings.T
213    return FactorResult(factors, loadings, eigenvalues, FE)
214
215
216# ---------------------------------------------------------------------------
217# Nuclear norm thresholding  (equivalent to R fect's panel_FE)
218# ---------------------------------------------------------------------------
219
220def panel_FE(E, lam: float, hard: bool = False) -> np.ndarray:
221    """Soft (or hard) singular value thresholding of *E* (T×N).
222
223    For matrix completion: shrink singular values by *lam* (soft) or zero out
224    those below *lam* (hard).
225
226    Returns the thresholded T×N matrix.
227    """
228    xp = get_backend()
229    T, N = E.shape
230    scale = T * N
231
232    M = E / scale
233    min_dim = min(T, N)
234    max_dim = max(T, N)
235    if min_dim > 0 and max_dim >= 4 * min_dim:
236        U, s, Vt = _gram_svd(M)
237    else:
238        U, s, Vt = full_svd(M)
239
240    if hard:
241        mask = s > lam
242        s_thresh = xp.where(mask, s, xp.zeros_like(s))
243    else:
244        s_thresh = xp.maximum(s - lam, 0.0)
245
246    # Reconstruct only with nonzero singular values
247    keep = s_thresh > 0
248    if not xp.any(keep):
249        return xp.zeros_like(E)
250
251    n_keep = int(xp.sum(keep))
252    FE = (U[:, :n_keep] * s_thresh[:n_keep][None, :]) @ Vt[:n_keep, :] * scale
253    return FE
254
255
256# ---------------------------------------------------------------------------
257# Demeaning (equivalent to Y_demean / Y_wdemean)
258# ---------------------------------------------------------------------------
259
260class DemeanResult(NamedTuple):
261    YY: np.ndarray       # demeaned matrix (T, N)
262    mu: float            # grand mean
263    alpha: np.ndarray | None   # unit FE (N,)
264    xi: np.ndarray | None      # time FE (T,)
265
266
267def demean(Y, force: int, W=None) -> DemeanResult:
268    """Remove additive fixed effects from *Y* (T×N).
269
270    Parameters
271    ----------
272    force : int
273        0=none, 1=unit, 2=time, 3=two-way.
274    W : array (T, N), optional
275        Observation weights.  If given, computes weighted means.
276    """
277    xp = get_backend()
278    T, N = Y.shape
279    alpha = None
280    xi = None
281
282    if W is not None:
283        # Weighted demeaning
284        Wsum = xp.sum(W)
285        mu = float(xp.sum(Y * W) / Wsum)
286        YY = Y.copy()
287
288        if force in (1, 3):
289            col_w = xp.sum(W, axis=0)  # (N,)
290            col_w = xp.where(col_w > 0, col_w, xp.ones_like(col_w))
291            alpha = xp.sum(Y * W, axis=0) / col_w  # (N,)
292            YY = YY - alpha[None, :]
293
294        if force in (2, 3):
295            row_w = xp.sum(W, axis=1)  # (T,)
296            row_w = xp.where(row_w > 0, row_w, xp.ones_like(row_w))
297            xi = xp.sum(Y * W, axis=1) / row_w  # (T,)
298            YY = YY - xi[:, None]
299
300        if force == 3:
301            YY = YY + mu
302        elif force == 0:
303            YY = YY - mu
304    else:
305        mu = float(xp.mean(Y))
306        YY = Y.copy()
307
308        if force in (1, 3):
309            alpha = xp.mean(Y, axis=0)  # (N,)
310            YY = YY - alpha[None, :]
311
312        if force in (2, 3):
313            xi = xp.mean(Y, axis=1)  # (T,)
314            YY = YY - xi[:, None]
315
316        if force == 3:
317            YY = YY + mu
318        elif force == 0:
319            YY = YY - mu
320
321    return DemeanResult(YY, mu, alpha, xi)
322
323
324def fe_add(mu: float, alpha, xi, T: int, N: int, force: int):
325    """Reconstruct additive FE matrix (T×N) from components.
326
327    R fect convention:
328      force=0: FE = mu
329      force=1: FE = mu + (alpha_Y - mu) = alpha_Y  (per-unit mean)
330      force=2: FE = mu + (xi_Y - mu) = xi_Y  (per-time mean)
331      force=3: FE = mu + (alpha_Y - mu) + (xi_Y - mu) = alpha_Y + xi_Y - mu
332    """
333    xp = get_backend()
334    FE = xp.full((T, N), mu, dtype=xp.float64)
335    if force in (1, 3) and alpha is not None:
336        FE = FE + (alpha - mu)[None, :]
337    if force in (2, 3) and xi is not None:
338        FE = FE + (xi - mu)[:, None]
339    return FE
340
341
342# ---------------------------------------------------------------------------
343# IFE: combined additive + interactive FE  (equivalent to R's ife())
344# ---------------------------------------------------------------------------
345
346class IFEResult(NamedTuple):
347    FE: np.ndarray          # total FE (T, N)
348    mu: float
349    alpha: np.ndarray | None
350    xi: np.ndarray | None
351    factors: np.ndarray | None
352    loadings: np.ndarray | None
353    eigenvalues: np.ndarray | None
354    FE_interactive: np.ndarray | None
355
356
357def ife(E, force: int, mc: bool = False, r: int = 0,
358        hard: bool = False, lam: float = 0.0) -> IFEResult:
359    """Additive + interactive fixed effects extraction.
360
361    Equivalent to fect's C++ ``ife()`` function.
362    """
363    xp = get_backend()
364    T, N = E.shape
365
366    # Step 1: demean
367    dm = demean(E, force)
368    EE = dm.YY
369
370    # Step 2: additive FE
371    FE_ad = fe_add(dm.mu, dm.alpha, dm.xi, T, N, force)
372
373    # Step 3: interactive FE
374    factors = loadings = eigenvalues = FE_inter = None
375    if r > 0 or mc:
376        if mc:
377            FE_inter = panel_FE(EE, lam, hard=hard)
378        else:
379            fr = panel_factor(EE, r)
380            factors = fr.factors
381            loadings = fr.loadings
382            eigenvalues = fr.eigenvalues
383            FE_inter = fr.FE
384
385    FE_total = FE_ad
386    if FE_inter is not None:
387        FE_total = FE_total + FE_inter
388
389    return IFEResult(FE_total, dm.mu, dm.alpha, dm.xi,
390                     factors, loadings, eigenvalues, FE_inter)
391
392
393# ---------------------------------------------------------------------------
394# Panel beta estimation (equivalent to panel_beta / wpanel_beta)
395# ---------------------------------------------------------------------------
396
397def panel_beta(X, Y, FE, xxinv=None, W=None):
398    """OLS ``beta`` with the fixed-effect fit partialled out.
399
400    Given a ``(T, N, p)`` covariate tensor, an outcome ``Y``, and the
401    current fixed-effect fit ``FE`` (also ``(T, N)``), solve
402
403    .. math::
404
405        \\beta = (X' W X)^{-1} X' W (Y - FE)
406
407    where ``W`` defaults to the identity observation weight.  This is the
408    closed-form step inside the EM loop whenever covariates are present.
409
410    Parameters
411    ----------
412    X : ``(T, N, p)`` ndarray
413        Covariate tensor.
414    Y : ``(T, N)`` ndarray
415    FE : ``(T, N)`` ndarray
416    xxinv : ``(p, p)`` ndarray, optional
417        Precomputed ``(X' W X)^{-1}``.  If supplied, reused as-is.
418    W : ``(T, N)`` ndarray, optional
419        Observation weights.  When ``None``, unweighted OLS.
420
421    Notes
422    -----
423    Implemented as single einsum calls so the work is pushed into BLAS
424    instead of a per-covariate Python loop.  Equivalent to R fect's
425    ``panel_beta`` / ``wpanel_beta`` (``auxiliary.cpp``).
426    """
427    xp = get_backend()
428    resid = Y - FE  # (T, N)
429
430    if W is not None:
431        # xy_k = sum_{t,n} W_{tn} X_{tnk} resid_{tn}
432        xy = xp.einsum("tnk,tn,tn->k", X, W, resid)
433        if xxinv is None:
434            xx = xp.einsum("tni,tnj,tn->ij", X, X, W)
435            xxinv = xp.linalg.inv(xx)
436    else:
437        xy = xp.einsum("tnk,tn->k", X, resid)
438        if xxinv is None:
439            xx = xp.einsum("tni,tnj->ij", X, X)
440            xxinv = xp.linalg.inv(xx)
441
442    return xxinv @ xy
443
444
445def compute_xxinv(X, W=None):
446    """Precompute ``(X' X)^{-1}`` or ``(X' W X)^{-1}`` via einsum.
447
448    Same contract as in R fect's ``XXinv`` / ``wXXinv``; vectorised so
449    the per-``p`` Python loop of earlier versions is replaced by a
450    single BLAS call.
451    """
452    xp = get_backend()
453    if W is not None:
454        xx = xp.einsum("tni,tnj,tn->ij", X, X, W)
455    else:
456        xx = xp.einsum("tni,tnj->ij", X, X)
457    return xp.linalg.inv(xx)
458
459
460def covar_fit(X, beta):
461    """Compute X @ beta contribution: sum_k X[:,:,k] * beta[k]."""
462    xp = get_backend()
463    # Vectorised: (T, N, p) @ (p,) -> (T, N) via einsum
464    return xp.einsum("ijk,k->ij", X, beta)