pyfector 0.1.0__py3-none-any.whl

pyfector/__init__.py

@@ -0,0 +1,41 @@
+"""
+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()
+
+"""
+
+__version__ = "0.1.0"
+
+from .fect import fect, FectResult
+from .backend import set_device, get_device
+from .diagnostics import run_diagnostics, DiagnosticResult
+from .plotting import plot
+
+__all__ = [
+    "fect",
+    "FectResult",
+    "set_device",
+    "get_device",
+    "run_diagnostics",
+    "DiagnosticResult",
+    "plot",
+]

pyfector/backend.py

@@ -0,0 +1,88 @@
+"""
+Array backend abstraction: seamless NumPy / CuPy switching.
+
+Usage:
+    from pyfector.backend import get_backend, set_device
+    set_device("gpu")          # or "cpu" (default)
+    xp = get_backend()         # returns numpy or cupy module
+    A = xp.random.randn(100, 50)
+
+All pyfector internals call ``get_backend()`` to obtain the active array
+module, so every matrix operation automatically runs on CPU or GPU.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from typing import Literal
+
+_DEVICE: Literal["cpu", "gpu"] = "cpu"
+_CUPY_AVAILABLE: bool | None = None
+
+
+def _check_cupy() -> bool:
+    global _CUPY_AVAILABLE
+    if _CUPY_AVAILABLE is None:
+        try:
+            import cupy  # noqa: F401
+            _CUPY_AVAILABLE = True
+        except ImportError:
+            _CUPY_AVAILABLE = False
+    return _CUPY_AVAILABLE
+
+
+def set_device(device: Literal["cpu", "gpu"]) -> None:
+    """Set the compute device globally."""
+    global _DEVICE
+    if device == "gpu" and not _check_cupy():
+        raise ImportError(
+            "CuPy is required for GPU support. Install it with: "
+            "pip install cupy-cuda12x  (adjust for your CUDA version)"
+        )
+    _DEVICE = device
+
+
+def get_device() -> Literal["cpu", "gpu"]:
+    """Return the current device."""
+    return _DEVICE
+
+
+def get_backend():
+    """Return the active array module (numpy or cupy)."""
+    if _DEVICE == "gpu":
+        import cupy
+        return cupy
+    return np
+
+
+def to_numpy(arr) -> np.ndarray:
+    """Ensure *arr* is a host numpy array (move from GPU if needed)."""
+    if _DEVICE == "gpu":
+        import cupy
+        if isinstance(arr, cupy.ndarray):
+            return cupy.asnumpy(arr)
+    return np.asarray(arr)
+
+
+def to_device(arr):
+    """Move a numpy array to the current device."""
+    if _DEVICE == "gpu":
+        import cupy
+        return cupy.asarray(arr)
+    return np.asarray(arr)
+
+
+def make_rng(seed: int | None = None):
+    """Create a seeded RNG for the active backend.
+
+    Returns a ``numpy.random.Generator`` or ``cupy.random.Generator``
+    depending on the current device.
+    """
+    if _DEVICE == "gpu":
+        import cupy
+        if seed is not None:
+            return cupy.random.default_rng(seed)
+        return cupy.random.default_rng()
+    if seed is not None:
+        return np.random.default_rng(seed)
+    return np.random.default_rng()

pyfector/cv.py

@@ -0,0 +1,327 @@
+"""
+Cross-validation for selecting number of factors (IFE) or lambda (MC).
+
+Key performance improvements over R fect:
+- Parallel fold evaluation via joblib
+- Warm-starting: previous fold's fit used as Y0 for next candidate
+- Looser tolerance during CV (max(tol, 1e-3))
+- Seeded RNG for reproducible fold splits
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Literal, NamedTuple
+
+import numpy as np
+from joblib import Parallel, delayed
+
+from .backend import get_backend, to_numpy, to_device, make_rng
+from .estimators import estimate_ife, estimate_mc
+
+
+class CVResult(NamedTuple):
+    best_r: int | None          # selected number of factors (IFE)
+    best_lambda: float | None   # selected lambda (MC)
+    scores: dict                # {candidate: {criterion: score}}
+    all_residuals: dict         # {candidate: residuals array}
+
+
+# ---------------------------------------------------------------------------
+# CV sampling
+# ---------------------------------------------------------------------------
+
+def _make_cv_folds(
+    II: np.ndarray,   # (T, N) control indicator
+    D: np.ndarray,    # (T, N) treatment indicator
+    k: int,
+    cv_prop: float,
+    cv_nobs: int,
+    cv_treat: bool,
+    cv_donut: int,
+    rng,
+) -> list[dict]:
+    """Generate k CV folds by masking control observations.
+
+    Returns list of dicts with keys 'mask' (cells to hide) and 'eval' (cells to score).
+    """
+    xp = get_backend()
+    folds = []
+    T, N = II.shape
+
+    # Eligible cells: observed control cells
+    if cv_treat:
+        # Only mask pre-treatment cells of treated units
+        ever_treated = xp.any(D > 0, axis=0)  # (N,)
+        eligible = II.copy()
+        eligible[:, ~ever_treated] = 0
+    else:
+        eligible = II.copy()
+
+    # Total control obs to mask per fold
+    total_eligible = int(xp.sum(eligible))
+    rm_count = max(1, int(total_eligible * cv_prop))
+
+    # Get indices of eligible cells
+    eligible_np = to_numpy(eligible)
+    elig_rows, elig_cols = np.where(eligible_np > 0)
+    n_eligible = len(elig_rows)
+
+    rng_np = np.random.default_rng(int(rng.integers(2**31)) if hasattr(rng, 'integers') else 42)
+
+    for fold_i in range(k):
+        # Randomly sample cells to mask
+        chosen = rng_np.choice(n_eligible, size=min(rm_count, n_eligible), replace=False)
+        mask_rows = elig_rows[chosen]
+        mask_cols = elig_cols[chosen]
+
+        # Build mask array
+        mask = np.zeros((T, N), dtype=bool)
+        mask[mask_rows, mask_cols] = True
+
+        # Eval indices (same as mask, optionally excluding donut around treatment)
+        eval_mask = mask.copy()
+        if cv_donut > 0:
+            for j in range(N):
+                treat_times = np.where(to_numpy(D[:, j]) > 0)[0]
+                if len(treat_times) > 0:
+                    first_treat = treat_times[0]
+                    donut_start = max(0, first_treat - cv_donut)
+                    donut_end = min(T, first_treat + cv_donut + 1)
+                    eval_mask[donut_start:donut_end, j] = False
+
+        folds.append({
+            "mask": to_device(mask),
+            "eval": to_device(eval_mask),
+        })
+
+    return folds
+
+
+# ---------------------------------------------------------------------------
+# Scoring
+# ---------------------------------------------------------------------------
+
+def _score_residuals(
+    residuals: np.ndarray,
+    eval_mask: np.ndarray,
+    count_weights: np.ndarray | None = None,
+) -> dict[str, float]:
+    """Compute CV scoring criteria from residuals at eval cells."""
+    xp = get_backend()
+
+    e = residuals[eval_mask]
+    n = max(int(xp.sum(eval_mask)), 1)
+
+    mspe = float(xp.sum(e ** 2) / n)
+
+    # GMSPE: geometric mean of squared errors
+    e2 = e ** 2
+    e2_pos = e2[e2 > 0]
+    if len(e2_pos) > 0:
+        gmspe = float(xp.exp(xp.mean(xp.log(e2_pos))))
+    else:
+        gmspe = float("inf")
+
+    # MAD: median absolute deviation
+    mad = float(xp.median(xp.abs(e))) if len(e) > 0 else float("inf")
+
+    return {
+        "mspe": mspe,
+        "gmspe": gmspe,
+        "mad": mad,
+    }
+
+
+# ---------------------------------------------------------------------------
+# CV for IFE (selecting r)
+# ---------------------------------------------------------------------------
+
+def cv_ife(
+    Y: np.ndarray,
+    Y0: np.ndarray,
+    X: np.ndarray | None,
+    I: np.ndarray,
+    II: np.ndarray,
+    D: np.ndarray,
+    W: np.ndarray | None,
+    beta0: np.ndarray | None,
+    force: int,
+    r_range: tuple[int, int],
+    k: int = 10,
+    cv_prop: float = 0.1,
+    cv_nobs: int = 3,
+    cv_treat: bool = True,
+    cv_donut: int = 0,
+    criterion: str = "mspe",
+    tol: float = 1e-5,
+    max_iter: int = 1000,
+    n_jobs: int = 1,
+    seed: int | None = None,
+) -> CVResult:
+    """Cross-validate to select number of factors *r*.
+
+    Parameters
+    ----------
+    r_range : (r_min, r_max)
+        Range of candidate factor numbers.
+    k : int
+        Number of CV folds.
+    n_jobs : int
+        Number of parallel workers for fold evaluation.
+    """
+    xp = get_backend()
+    rng = make_rng(seed)
+
+    # Looser tolerance for CV (like R fect)
+    cv_tol = max(tol, 1e-3)
+
+    # Generate folds
+    folds = _make_cv_folds(II, D, k, cv_prop, cv_nobs, cv_treat, cv_donut, rng)
+
+    r_candidates = list(range(r_range[0], r_range[1] + 1))
+    scores = {}
+
+    for r_cand in r_candidates:
+        fold_residuals = []
+
+        def _run_fold(fold):
+            II_cv = II.copy()
+            II_cv[fold["mask"]] = 0
+            # Re-initialize Y0 for masked cells with simple mean
+            Y0_cv = Y0.copy()
+            if xp.sum(II_cv) > 0:
+                fill_val = xp.sum(Y * II_cv) / (xp.sum(II_cv) + 1e-10)
+                Y0_cv = xp.where(II_cv, Y0_cv, fill_val)
+
+            result = estimate_ife(
+                Y, Y0_cv, X, II_cv, W, beta0,
+                r=r_cand, force=force, tol=cv_tol, max_iter=max_iter,
+            )
+            # Score on eval cells
+            resid = Y - result.fit
+            return _score_residuals(resid, fold["eval"])
+
+        if n_jobs == 1:
+            fold_scores = [_run_fold(f) for f in folds]
+        else:
+            fold_scores = Parallel(n_jobs=n_jobs, prefer="threads")(
+                delayed(_run_fold)(f) for f in folds
+            )
+
+        # Average scores across folds
+        avg_scores = {}
+        for metric in fold_scores[0]:
+            avg_scores[metric] = sum(fs[metric] for fs in fold_scores) / len(fold_scores)
+        scores[r_cand] = avg_scores
+
+    # Select best r (with 1% improvement threshold)
+    best_r = _select_best(scores, r_candidates, criterion)
+
+    return CVResult(best_r=best_r, best_lambda=None, scores=scores, all_residuals={})
+
+
+# ---------------------------------------------------------------------------
+# CV for MC (selecting lambda)
+# ---------------------------------------------------------------------------
+
+def cv_mc(
+    Y: np.ndarray,
+    Y0: np.ndarray,
+    X: np.ndarray | None,
+    I: np.ndarray,
+    II: np.ndarray,
+    D: np.ndarray,
+    W: np.ndarray | None,
+    beta0: np.ndarray | None,
+    force: int,
+    lambda_candidates: np.ndarray | None = None,
+    nlambda: int = 10,
+    k: int = 10,
+    cv_prop: float = 0.1,
+    cv_nobs: int = 3,
+    cv_treat: bool = True,
+    cv_donut: int = 0,
+    criterion: str = "mspe",
+    tol: float = 1e-5,
+    max_iter: int = 1000,
+    n_jobs: int = 1,
+    seed: int | None = None,
+) -> CVResult:
+    """Cross-validate to select lambda (nuclear norm penalty).
+
+    If ``lambda_candidates`` is None, generates a log-spaced grid from the
+    data's largest singular value down to 0.
+    """
+    xp = get_backend()
+    rng = make_rng(seed)
+    cv_tol = max(tol, 1e-3)
+
+    # Generate lambda grid if not provided
+    if lambda_candidates is None:
+        lambda_candidates = _make_lambda_grid(Y, II, nlambda)
+
+    # Generate folds
+    folds = _make_cv_folds(II, D, k, cv_prop, cv_nobs, cv_treat, cv_donut, rng)
+
+    scores = {}
+    for lam in lambda_candidates:
+        lam_key = float(lam)
+
+        def _run_fold(fold):
+            II_cv = II.copy()
+            II_cv[fold["mask"]] = 0
+            Y0_cv = Y0.copy()
+            if xp.sum(II_cv) > 0:
+                fill_val = xp.sum(Y * II_cv) / (xp.sum(II_cv) + 1e-10)
+                Y0_cv = xp.where(II_cv, Y0_cv, fill_val)
+
+            result = estimate_mc(
+                Y, Y0_cv, X, II_cv, W, beta0,
+                lam=lam_key, force=force, tol=cv_tol, max_iter=max_iter,
+            )
+            resid = Y - result.fit
+            return _score_residuals(resid, fold["eval"])
+
+        if n_jobs == 1:
+            fold_scores = [_run_fold(f) for f in folds]
+        else:
+            fold_scores = Parallel(n_jobs=n_jobs, prefer="threads")(
+                delayed(_run_fold)(f) for f in folds
+            )
+
+        avg_scores = {}
+        for metric in fold_scores[0]:
+            avg_scores[metric] = sum(fs[metric] for fs in fold_scores) / len(fold_scores)
+        scores[lam_key] = avg_scores
+
+    best_lambda = _select_best(scores, [float(l) for l in lambda_candidates], criterion)
+
+    return CVResult(best_r=None, best_lambda=best_lambda, scores=scores, all_residuals={})
+
+
+def _make_lambda_grid(Y, II, nlambda: int):
+    """Generate log-spaced lambda grid from data's singular values."""
+    xp = get_backend()
+    T, N = Y.shape
+    Y_obs = Y * II
+    # Get largest singular value
+    _, s, _ = xp.linalg.svd(Y_obs / (T * N), full_matrices=False)
+    lam_max = float(s[0])
+    log_max = math.log10(max(lam_max, 1e-10))
+    step = 3.0 / max(nlambda - 2, 1)
+    grid = [10 ** (log_max - i * step) for i in range(nlambda - 1)]
+    grid.append(0.0)
+    return np.array(grid)
+
+
+def _select_best(scores: dict, candidates: list, criterion: str):
+    """Select best candidate with 1% improvement threshold."""
+    vals = [(c, scores[c][criterion]) for c in candidates]
+    vals.sort(key=lambda x: x[1])
+    best = vals[0]
+    # Apply 1% threshold: prefer simpler model if improvement < 1%
+    for c, v in vals:
+        if v <= best[1] * 1.01:
+            return c
+    return best[0]

pyfector/diagnostics.py

@@ -0,0 +1,213 @@
+"""
+Diagnostic tests for pyfector results.
+
+Implements:
+- Pre-trend F-test (Wald-type joint significance)
+- Equivalence F-test (against non-central F)
+- TOST (two one-sided t-tests) for each pre-treatment period
+- Placebo test
+- Carryover test
+- Leave-one-period-out (LOO) test
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class DiagnosticResult:
+    """Container for diagnostic test results."""
+    # F-test for pre-trends
+    f_stat: float | None = None
+    f_pval: float | None = None
+    f_df1: int | None = None
+    f_df2: int | None = None
+
+    # Equivalence F-test
+    equiv_f_pval: float | None = None
+    equiv_threshold: float | None = None
+
+    # TOST per pre-period
+    tost_pvals: np.ndarray | None = None
+    tost_periods: np.ndarray | None = None
+    tost_threshold: float | None = None
+
+    # Placebo
+    placebo_att: float | None = None
+    placebo_pval: float | None = None
+
+    # Carryover
+    carryover_att: float | None = None
+    carryover_pval: float | None = None
+
+    # Leave-one-out
+    loo_atts: np.ndarray | None = None      # ATT when each pre-period is dropped
+    loo_periods: np.ndarray | None = None
+    loo_max_change: float | None = None     # max |ATT_full - ATT_loo|
+
+    def summary(self) -> str:
+        lines = ["Diagnostic Tests", "=" * 50]
+        if self.f_stat is not None:
+            lines.append(f"Pre-trend F-test:")
+            lines.append(f"  F({self.f_df1},{self.f_df2}) = {self.f_stat:.4f}, p = {self.f_pval:.4f}")
+        if self.equiv_f_pval is not None:
+            lines.append(f"Equivalence F-test:")
+            lines.append(f"  p = {self.equiv_f_pval:.4f} (threshold={self.equiv_threshold})")
+        if self.tost_pvals is not None:
+            lines.append("TOST per pre-period:")
+            for i, t in enumerate(self.tost_periods):
+                p = self.tost_pvals[i]
+                status = "PASS" if p < 0.05 else "fail"
+                lines.append(f"  t={t:+.0f}: p={p:.4f} [{status}]")
+        if self.placebo_pval is not None:
+            lines.append(f"Placebo test:")
+            lines.append(f"  ATT={self.placebo_att:.4f}, p={self.placebo_pval:.4f}")
+        if self.carryover_pval is not None:
+            lines.append(f"Carryover test:")
+            lines.append(f"  ATT={self.carryover_att:.4f}, p={self.carryover_pval:.4f}")
+        if self.loo_max_change is not None:
+            lines.append(f"Leave-one-out:")
+            lines.append(f"  Max ATT change = {self.loo_max_change:.6f}")
+        return "\n".join(lines)
+
+
+def run_diagnostics(
+    result,
+    f_threshold: float = 0.5,
+    tost_threshold: float = 0.36,
+    placebo_period: tuple[int, int] | None = None,
+    carryover_period: tuple[int, int] | None = None,
+    loo: bool = False,
+) -> DiagnosticResult:
+    """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.
+    placebo_period : (start, end), optional
+        Relative time window for placebo test.
+    """
+    from scipy import stats
+
+    diag = DiagnosticResult()
+
+    if result.inference is None:
+        return diag
+
+    inf = result.inference
+    time_on = result.time_on
+    att_on = result.att_on
+
+    if time_on is None or att_on is None:
+        return diag
+
+    # Pre-treatment periods
+    pre_mask = time_on < 0
+    if not np.any(pre_mask):
+        return diag
+
+    pre_idx = np.where(pre_mask)[0]
+    k = len(pre_idx)
+    att_pre = att_on[pre_idx]
+
+    # F-test from bootstrap covariance
+    if inf.att_on_boot is not None and inf.att_on_boot.shape[1] > k:
+        boot_pre = inf.att_on_boot[pre_idx, :]  # (k, nboots)
+        S = np.cov(boot_pre)  # (k, k)
+        if k == 1:
+            S = S.reshape(1, 1)
+
+        n_boot = boot_pre.shape[1]
+        try:
+            cond = np.linalg.cond(S)
+            if cond > 1e12:
+                S_inv = np.linalg.pinv(S)
+            else:
+                S_inv = np.linalg.inv(S)
+            F_stat = float(att_pre @ S_inv @ att_pre)
+            scale = (n_boot - k) / ((n_boot - 1) * k)
+            F_stat_scaled = F_stat * scale
+
+            if np.isfinite(F_stat_scaled) and F_stat_scaled >= 0:
+                diag.f_stat = F_stat_scaled
+                diag.f_df1 = k
+                diag.f_df2 = n_boot - k
+                diag.f_pval = float(1 - stats.f.cdf(F_stat_scaled, k, n_boot - k))
+
+                # Equivalence F-test (non-central F)
+                ncp = n_boot * f_threshold
+                diag.equiv_f_pval = float(stats.ncf.cdf(F_stat_scaled, k, n_boot - k, ncp))
+                diag.equiv_threshold = f_threshold
+        except np.linalg.LinAlgError:
+            pass
+
+    # TOST
+    if inf.att_on_se is not None:
+        se_pre = inf.att_on_se[pre_idx]
+        n_boot = inf.att_on_boot.shape[1] if inf.att_on_boot is not None else 200
+        df = n_boot - 1
+
+        tost_pvals = np.full(k, np.nan)
+        for i in range(k):
+            if se_pre[i] > 0:
+                t_upper = (att_pre[i] - tost_threshold) / se_pre[i]
+                t_lower = (att_pre[i] + tost_threshold) / se_pre[i]
+                p_upper = float(stats.t.cdf(t_upper, df))
+                p_lower = float(1 - stats.t.cdf(t_lower, df))
+                tost_pvals[i] = max(p_upper, p_lower)
+
+        diag.tost_pvals = tost_pvals
+        diag.tost_periods = time_on[pre_idx]
+        diag.tost_threshold = tost_threshold
+
+    # Placebo test
+    if placebo_period is not None:
+        p_start, p_end = placebo_period
+        placebo_mask = (time_on >= p_start) & (time_on <= p_end) & pre_mask
+        if np.any(placebo_mask):
+            p_idx = np.where(placebo_mask)[0]
+            diag.placebo_att = float(np.mean(att_on[p_idx]))
+            if inf.att_on_boot is not None:
+                boot_placebo = np.mean(inf.att_on_boot[p_idx, :], axis=0)
+                diag.placebo_pval = min(
+                    2 * np.mean(boot_placebo >= 0),
+                    2 * np.mean(boot_placebo <= 0),
+                    1.0,
+                )
+
+    # Carryover test (for treatment reversals)
+    if carryover_period is not None and hasattr(result, "att_off") and result.att_off is not None:
+        c_start, c_end = carryover_period
+        time_off = getattr(result, "time_off", None)
+        if time_off is not None:
+            off_mask = (time_off >= c_start) & (time_off <= c_end)
+            if np.any(off_mask):
+                diag.carryover_att = float(np.mean(result.att_off[off_mask]))
+
+    # Leave-one-period-out test
+    if loo and att_on is not None and time_on is not None:
+        post_mask = time_on >= 0
+        post_idx = np.where(post_mask)[0]
+        if len(post_idx) > 1:
+            full_att = result.att_avg
+            loo_atts = []
+            loo_periods = []
+            for drop_i in post_idx:
+                remaining = np.delete(post_idx, np.where(post_idx == drop_i))
+                if len(remaining) > 0:
+                    loo_att = float(np.mean(att_on[remaining]))
+                    loo_atts.append(loo_att)
+                    loo_periods.append(time_on[drop_i])
+            diag.loo_atts = np.array(loo_atts)
+            diag.loo_periods = np.array(loo_periods)
+            diag.loo_max_change = float(np.max(np.abs(diag.loo_atts - full_att)))
+
+    return diag