pyfector.cv

  1"""
  2Cross-validation for hyperparameter selection.
  3
  4This module provides two entry points used by :func:`pyfector.fect`:
  5
  6* :func:`cv_ife` picks the number of latent factors ``r`` for the IFE
  7  estimator by grid search over ``r_range = (r_min, r_max)``.
  8* :func:`cv_mc` picks the nuclear-norm penalty ``lam`` for the
  9  matrix-completion estimator over an automatically generated
 10  log-spaced grid (or a user-supplied one).
 11
 12CV scheme
 13---------
 14* Hold-out cells are sampled *only* from control observations
 15  (``II``).  When ``cv_treat=True`` (the default) we further restrict
 16  the hold-out cells to pre-treatment periods of ever-treated units, so
 17  CV mimics the prediction task the estimator faces on treated cells.
 18* For each candidate we fit on the non-masked cells and score on the
 19  masked cells using one of ``mspe`` (mean squared prediction error),
 20  ``gmspe`` (geometric mean of squared errors), or ``mad`` (median
 21  absolute deviation).  The paper-faithful default is
 22  ``cv_rule="min"``, which chooses the candidate with the lowest
 23  validation score.  ``cv_rule="onepct"`` is an explicit slack option:
 24  it chooses the simplest candidate within 1% of the best score (lower
 25  ``r`` for IFE and higher ``lam`` for MC).
 26* Fold parallelism uses a thread-backed joblib pool so the heavy
 27  numpy/BLAS work runs concurrently without pickling.
 28* Randomness (fold construction) is seeded via :func:`pyfector.backend.make_rng`
 29  so CV is reproducible given ``seed``.
 30
 31The CV inner loop calls into :func:`pyfector.estimators.estimate_ife`
 32and :func:`pyfector.estimators.estimate_mc` with a relaxed convergence
 33tolerance ``max(tol, 1e-3)`` — the same trick R fect uses to keep CV
 34cheap without affecting the final fit, which always uses the tight
 35tolerance.
 36"""
 37
 38from __future__ import annotations
 39
 40import math
 41import warnings
 42from typing import Literal, NamedTuple
 43
 44import numpy as np
 45from joblib import Parallel, delayed, parallel_config
 46
 47from .backend import get_backend, to_numpy, to_device, make_rng
 48from .estimators import estimate_ife, estimate_mc
 49
 50
 51class CVResult(NamedTuple):
 52    """Selected hyperparameters and validation scores."""
 53    best_r: int | None          # selected number of factors (IFE)
 54    """Selected IFE rank, or None for MC."""
 55    best_lambda: float | None   # selected lambda (MC)
 56    """Selected MC penalty, or None for IFE."""
 57    scores: dict                # {candidate: {criterion: score}}
 58    """Validation scores by candidate and criterion."""
 59    all_residuals: dict         # {candidate: residuals array}
 60    """Validation residuals for each candidate when retained."""
 61
 62
 63# ---------------------------------------------------------------------------
 64# CV sampling
 65# ---------------------------------------------------------------------------
 66
 67def _make_cv_folds(
 68    II: np.ndarray,   # (T, N) control indicator
 69    D: np.ndarray,    # (T, N) treatment indicator
 70    k: int,
 71    cv_prop: float,
 72    cv_nobs: int,
 73    cv_treat: bool,
 74    cv_donut: int,
 75    rng,
 76) -> list[dict]:
 77    """Generate k CV folds by masking control observations.
 78
 79    Returns list of dicts with keys 'mask' (cells to hide) and 'eval' (cells to score).
 80    """
 81    xp = get_backend()
 82    folds = []
 83    T, N = II.shape
 84
 85    # Eligible cells: observed control cells
 86    if cv_treat:
 87        # Only mask pre-treatment cells of treated units
 88        ever_treated = xp.any(D > 0, axis=0)  # (N,)
 89        eligible = II.copy()
 90        eligible[:, ~ever_treated] = 0
 91    else:
 92        eligible = II.copy()
 93
 94    # Total control obs to mask per fold
 95    total_eligible = int(xp.sum(eligible))
 96    rm_count = max(1, int(total_eligible * cv_prop))
 97
 98    # Get indices of eligible cells
 99    II_np = to_numpy(II)
100    D_np = to_numpy(D)
101    eligible_np = to_numpy(eligible)
102    elig_rows, elig_cols = np.where(eligible_np > 0)
103    n_eligible = len(elig_rows)
104
105    rng_np = np.random.default_rng(int(rng.integers(2**31)) if hasattr(rng, 'integers') else 42)
106
107    if n_eligible == 0:
108        raise ValueError("No eligible control observations for cross-validation")
109
110    for fold_i in range(k):
111        mask = np.zeros((T, N), dtype=bool)
112
113        if cv_nobs <= 1 or rm_count <= 2:
114            chosen = rng_np.choice(n_eligible, size=min(rm_count, n_eligible), replace=False)
115            mask[elig_rows[chosen], elig_cols[chosen]] = True
116        else:
117            # Match R fect's block-style CV more closely: hide short
118            # within-unit runs rather than only isolated cells.
119            target = min(rm_count, n_eligible)
120            n_masked = 0
121            while n_masked < target:
122                base = int(rng_np.integers(n_eligible))
123                row = int(elig_rows[base])
124                col = int(elig_cols[base])
125                stop = min(T, row + cv_nobs)
126                rows = np.arange(row, stop)
127                rows = rows[eligible_np[rows, col] > 0]
128                if len(rows) == 0:
129                    if not mask[row, col]:
130                        mask[row, col] = True
131                        n_masked += 1
132                else:
133                    new_rows = rows[~mask[rows, col]]
134                    if len(new_rows) > 0:
135                        mask[new_rows, col] = True
136                        n_masked += int(len(new_rows))
137
138            selected = np.flatnonzero(mask.ravel())
139            if len(selected) > rm_count:
140                drop = rng_np.choice(selected, size=len(selected) - rm_count, replace=False)
141                mask.ravel()[drop] = False
142
143        # Keep every row and usable unit identifiable after masking.
144        II_after = II_np.copy()
145        II_after[mask] = 0
146        bad_rows = np.where(II_after.sum(axis=1) < 1)[0]
147        if len(bad_rows) > 0:
148            mask[bad_rows, :] = False
149            II_after = II_np.copy()
150            II_after[mask] = 0
151        orig_col_counts = II_np.sum(axis=0)
152        bad_cols = np.where((orig_col_counts > 0) & (II_after.sum(axis=0) < 1))[0]
153        if len(bad_cols) > 0:
154            mask[:, bad_cols] = False
155
156        # Eval indices (same as mask, optionally excluding donut around treatment)
157        eval_mask = mask.copy()
158        if cv_donut > 0:
159            for j in range(N):
160                treat_times = np.where(D_np[:, j] > 0)[0]
161                if len(treat_times) > 0:
162                    first_treat = treat_times[0]
163                    donut_start = max(0, first_treat - cv_donut)
164                    donut_end = min(T, first_treat + cv_donut + 1)
165                    eval_mask[donut_start:donut_end, j] = False
166
167        folds.append({
168            "mask": to_device(mask),
169            "eval": to_device(eval_mask),
170        })
171
172    return folds
173
174
175# ---------------------------------------------------------------------------
176# Scoring
177# ---------------------------------------------------------------------------
178
179def _score_residuals(
180    residuals: np.ndarray,
181    eval_mask: np.ndarray,
182    count_weights: np.ndarray | None = None,
183) -> dict[str, float]:
184    """Compute CV scoring criteria from residuals at eval cells."""
185    e = to_numpy(residuals[eval_mask])
186    n = max(int(e.size), 1)
187
188    mspe = float(np.sum(e ** 2) / n)
189
190    # GMSPE: geometric mean of squared errors
191    e2 = e ** 2
192    e2_pos = e2[e2 > 0]
193    if len(e2_pos) > 0:
194        gmspe = float(np.exp(np.mean(np.log(e2_pos))))
195    else:
196        gmspe = float("inf")
197
198    # MAD: median absolute deviation
199    mad = float(np.median(np.abs(e))) if len(e) > 0 else float("inf")
200
201    return {
202        "mspe": mspe,
203        "gmspe": gmspe,
204        "mad": mad,
205    }
206
207
208# ---------------------------------------------------------------------------
209# CV for IFE (selecting r)
210# ---------------------------------------------------------------------------
211
212def cv_ife(
213    Y: np.ndarray,
214    Y0: np.ndarray,
215    X: np.ndarray | None,
216    I: np.ndarray,
217    II: np.ndarray,
218    D: np.ndarray,
219    W: np.ndarray | None,
220    beta0: np.ndarray | None,
221    force: int,
222    r_range: tuple[int, int],
223    k: int = 10,
224    cv_prop: float = 0.1,
225    cv_nobs: int = 3,
226    cv_treat: bool = True,
227    cv_donut: int = 0,
228    criterion: str = "mspe",
229    cv_rule: Literal["min", "onepct"] = "min",
230    tol: float = 1e-5,
231    max_iter: int = 1000,
232    n_jobs: int = 1,
233    seed: int | None = None,
234) -> CVResult:
235    """Cross-validate to select number of factors *r*.
236
237    Parameters
238    ----------
239    r_range : (r_min, r_max)
240        Range of candidate factor numbers.
241    k : int
242        Number of CV folds.
243    cv_rule : {"min", "onepct"}
244        ``"min"`` selects the lowest-score candidate, matching the
245        paper's cross-validation description. ``"onepct"`` selects the
246        lowest rank within 1% of the global best score.
247    n_jobs : int
248        Number of parallel workers for fold evaluation.
249    """
250    xp = get_backend()
251    rng = make_rng(seed)
252
253    # Looser tolerance for CV (like R fect)
254    cv_tol = max(tol, 1e-3)
255
256    # Generate folds
257    folds = _make_cv_folds(II, D, k, cv_prop, cv_nobs, cv_treat, cv_donut, rng)
258
259    r_candidates = list(range(r_range[0], r_range[1] + 1))
260    scores = {}
261
262    def _score_candidate(r_cand: int, parallel: Parallel | None = None) -> dict[str, float]:
263        def _run_fold(fold):
264            II_cv = II.copy()
265            II_cv[fold["mask"]] = 0
266            # Re-initialize Y0 for masked cells with simple mean
267            Y0_cv = Y0.copy()
268            if xp.sum(II_cv) > 0:
269                fill_val = xp.sum(Y * II_cv) / (xp.sum(II_cv) + 1e-10)
270                Y0_cv = xp.where(II_cv, Y0_cv, fill_val)
271
272            result = estimate_ife(
273                Y, Y0_cv, X, II_cv, W, beta0,
274                r=r_cand, force=force, tol=cv_tol, max_iter=max_iter,
275            )
276            # Score on eval cells
277            resid = Y - result.fit
278            return _score_residuals(resid, fold["eval"])
279
280        if parallel is None:
281            fold_scores = [_run_fold(f) for f in folds]
282        else:
283            fold_scores = parallel(delayed(_run_fold)(f) for f in folds)
284
285        # Average scores across folds
286        return {
287            metric: sum(fs[metric] for fs in fold_scores) / len(fold_scores)
288            for metric in fold_scores[0]
289        }
290
291    if n_jobs == 1:
292        for r_cand in r_candidates:
293            scores[r_cand] = _score_candidate(r_cand)
294    else:
295        with parallel_config(backend="loky", inner_max_num_threads=1), \
296                Parallel(n_jobs=n_jobs) as parallel:
297            for r_cand in r_candidates:
298                scores[r_cand] = _score_candidate(r_cand, parallel)
299
300    # r_candidates is ordered from simpler to more flexible.
301    best_r = _select_best(scores, r_candidates, criterion, cv_rule=cv_rule)
302
303    return CVResult(best_r=best_r, best_lambda=None, scores=scores, all_residuals={})
304
305
306def cv_mc(
307    Y: np.ndarray,
308    Y0: np.ndarray,
309    X: np.ndarray | None,
310    I: np.ndarray,
311    II: np.ndarray,
312    D: np.ndarray,
313    W: np.ndarray | None,
314    beta0: np.ndarray | None,
315    force: int,
316    lambda_candidates: np.ndarray | None = None,
317    nlambda: int = 10,
318    k: int = 10,
319    cv_prop: float = 0.1,
320    cv_nobs: int = 3,
321    cv_treat: bool = True,
322    cv_donut: int = 0,
323    criterion: str = "mspe",
324    cv_rule: Literal["min", "onepct"] = "min",
325    tol: float = 1e-5,
326    max_iter: int = 1000,
327    n_jobs: int = 1,
328    seed: int | None = None,
329) -> CVResult:
330    """Cross-validate to select lambda (nuclear norm penalty).
331
332    If ``lambda_candidates`` is None, generates a log-spaced grid from the
333    data's largest singular value down to 0. With ``cv_rule="onepct"``, the
334    largest penalty within 1% of the global best score is selected.
335    """
336    xp = get_backend()
337    rng = make_rng(seed)
338    cv_tol = max(tol, 1e-3)
339
340    # Generate lambda grid if not provided
341    if lambda_candidates is None:
342        lambda_candidates = _make_lambda_grid(Y, II, nlambda)
343    else:
344        lambda_candidates = np.asarray(lambda_candidates, dtype=float)
345        if lambda_candidates.ndim != 1 or lambda_candidates.size == 0:
346            raise ValueError("lambda_candidates must be a non-empty one-dimensional array")
347        if not np.all(np.isfinite(lambda_candidates)):
348            raise ValueError("lambda_candidates must contain only finite values")
349        if np.any(lambda_candidates < 0):
350            raise ValueError("lambda_candidates must be non-negative")
351
352    # Generate folds
353    folds = _make_cv_folds(II, D, k, cv_prop, cv_nobs, cv_treat, cv_donut, rng)
354
355    scores = {}
356    def _score_candidate(lam: float, parallel: Parallel | None = None) -> dict[str, float]:
357        lam_key = float(lam)
358
359        def _run_fold(fold):
360            II_cv = II.copy()
361            II_cv[fold["mask"]] = 0
362            Y0_cv = Y0.copy()
363            if xp.sum(II_cv) > 0:
364                fill_val = xp.sum(Y * II_cv) / (xp.sum(II_cv) + 1e-10)
365                Y0_cv = xp.where(II_cv, Y0_cv, fill_val)
366
367            result = estimate_mc(
368                Y, Y0_cv, X, II_cv, W, beta0,
369                lam=lam_key, force=force, tol=cv_tol, max_iter=max_iter,
370            )
371            resid = Y - result.fit
372            return _score_residuals(resid, fold["eval"])
373
374        if parallel is None:
375            fold_scores = [_run_fold(f) for f in folds]
376        else:
377            fold_scores = parallel(delayed(_run_fold)(f) for f in folds)
378
379        return {
380            metric: sum(fs[metric] for fs in fold_scores) / len(fold_scores)
381            for metric in fold_scores[0]
382        }
383
384    if n_jobs == 1:
385        for lam in lambda_candidates:
386            scores[float(lam)] = _score_candidate(float(lam))
387    else:
388        with parallel_config(backend="loky", inner_max_num_threads=1), \
389                Parallel(n_jobs=n_jobs) as parallel:
390            for lam in lambda_candidates:
391                scores[float(lam)] = _score_candidate(float(lam), parallel)
392
393    # Select from heavier to lighter regularization, i.e. from simpler
394    # to more flexible. Sort here so user-supplied grids behave the same
395    # way as the generated grid.
396    lambda_order = sorted([float(l) for l in lambda_candidates], reverse=True)
397    best_lambda = _select_best(
398        scores, lambda_order, criterion,
399        cv_rule=cv_rule,
400    )
401    _warn_if_lambda_on_boundary(best_lambda, lambda_candidates)
402
403    return CVResult(best_r=None, best_lambda=best_lambda, scores=scores, all_residuals={})
404
405
406def _make_lambda_grid(Y, II, nlambda: int):
407    """Generate log-spaced lambda grid from data's singular values."""
408    xp = get_backend()
409    T, N = Y.shape
410    Y_obs = Y * II
411    # Get largest singular value
412    _, s, _ = xp.linalg.svd(Y_obs / (T * N), full_matrices=False)
413    lam_max = float(s[0])
414    log_max = math.log10(max(lam_max, 1e-10))
415    step = 3.0 / max(nlambda - 2, 1)
416    grid = [10 ** (log_max - i * step) for i in range(nlambda - 1)]
417    grid.append(0.0)
418    return np.array(grid)
419
420
421def _select_best(
422    scores: dict,
423    candidates: list,
424    criterion: str,
425    cv_rule: Literal["min", "onepct"] = "min",
426):
427    """Select a CV candidate by min score or an explicit 1% slack rule.
428
429    ``"min"`` is the paper-faithful rule: select the candidate with the
430    lowest validation score.  For ``cv_rule="onepct"``, ``candidates``
431    must be ordered from simpler to more flexible; the selected candidate
432    is the first candidate within 1% of the global best score.
433    """
434    if not candidates:
435        raise ValueError("candidates must not be empty")
436    if cv_rule not in {"min", "onepct"}:
437        raise ValueError("cv_rule must be 'min' or 'onepct'")
438
439    vals = [(c, float(scores[c][criterion])) for c in candidates]
440    if cv_rule == "min":
441        return min(vals, key=lambda x: x[1])[0]
442
443    best_score = min(v for _, v in vals)
444    threshold = best_score * 1.01
445    for c, v in vals:
446        if v <= threshold:
447            return c
448
449    return min(vals, key=lambda x: x[1])[0]
450
451
452def _warn_if_lambda_on_boundary(best_lambda: float | None, lambda_candidates) -> None:
453    """Warn when MC cross-validation selects a grid boundary."""
454    if best_lambda is None:
455        return
456
457    candidates = np.asarray([float(lam) for lam in lambda_candidates], dtype=float)
458    if candidates.size <= 1:
459        return
460
461    best = float(best_lambda)
462    lower = float(np.min(candidates))
463    upper = float(np.max(candidates))
464
465    if np.isclose(best, lower, rtol=1e-12, atol=1e-12):
466        warnings.warn(
467            (
468                "MC cross-validation selected the lowest lambda in the candidate "
469                f"grid ({lower:g}). This boundary choice may indicate that little "
470                "or no regularization is preferred; compare cv_rule='onepct' or "
471                "set lam explicitly after validating the boundary selection."
472            ),
473            UserWarning,
474            stacklevel=3,
475        )
476    elif np.isclose(best, upper, rtol=1e-12, atol=1e-12):
477        warnings.warn(
478            (
479                "MC cross-validation selected the highest lambda in the candidate "
480                f"grid ({upper:g}). This boundary choice may indicate that the "
481                "grid should be extended upward."
482            ),
483            UserWarning,
484            stacklevel=3,
485        )