pyfector.inference

  1"""
  2Non-parametric inference: bootstrap, jackknife, and permutation.
  3
  4This module builds confidence intervals and standard errors around the
  5point estimates returned by :func:`pyfector.fect`.  All three routines
  6resample *units* (not cells), so they naturally respect within-unit
  7dependence:
  8
  9* :func:`bootstrap` — draws ``nboots`` bootstrap samples of unit indices
 10  with replacement, re-estimates on each sample, and returns pivotal
 11  confidence intervals, bootstrap SEs, and two-sided p-values.  Treated,
 12  control, and reversal units are resampled separately so the treated
 13  share stays fixed across replicates (stratified bootstrap); when
 14  ``cluster`` is supplied we fall back to a cluster bootstrap.
 15* :func:`jackknife` — delete-one-unit jackknife with the standard
 16  ``(N-1)/N`` variance correction and t-distribution CIs.
 17* :func:`permutation_test` — shuffles treatment assignment and compares
 18  the observed ATT to the null distribution of permuted ATTs.
 19
 20Implementation notes
 21--------------------
 22* Replications are parallelised via joblib's ``loky`` (process) backend
 23  with ``inner_max_num_threads=1``.  Processes avoid the GIL contention
 24  that throttles the weighted EM path, whose per-replicate work is
 25  dominated by GIL-holding elementwise numpy operations rather than
 26  GIL-releasing BLAS calls.  Pinning each worker to a single BLAS thread
 27  keeps ``n_jobs`` workers from oversubscribing the cores that a
 28  multithreaded BLAS would otherwise claim inside every replicate.
 29* Every replicate receives a deterministic seed derived from the
 30  master ``seed`` so results are reproducible.
 31* The dynamic-ATT grid is pinned to the point-estimate's ``time_on``
 32  so replicate outputs can be stacked into a ``(n_periods, nboots)``
 33  matrix even when a resampled draw happens to miss a relative-time
 34  value.
 35"""
 36
 37from __future__ import annotations
 38
 39from dataclasses import dataclass, field
 40from typing import Literal
 41import warnings
 42
 43import numpy as np
 44from joblib import Parallel, delayed, parallel_config
 45
 46from .backend import get_backend, to_numpy, to_device, make_rng
 47
 48
 49@dataclass
 50class InferenceResult:
 51    """Inference results and distributions.
 52
 53    For bootstrap inference, ``n_boot`` and ``n_boot_failed`` count valid
 54    and failed replications. These counts are zero for jackknife results.
 55    """
 56    att_avg: float                          # point estimate of overall ATT
 57    """Overall ATT point estimate."""
 58    att_avg_se: float                       # standard error
 59    """Standard error of the overall ATT."""
 60    att_avg_ci: tuple[float, float]         # confidence interval
 61    """Confidence interval for the overall ATT."""
 62    att_avg_pval: float                     # p-value
 63    """Two-sided p-value for zero overall ATT."""
 64
 65    att_on: np.ndarray                      # (n_periods,) dynamic ATT
 66    """Dynamic ATT point estimates on time_on."""
 67    att_on_se: np.ndarray                   # (n_periods,) SE
 68    """Standard errors of dynamic ATT estimates."""
 69    att_on_ci_lower: np.ndarray             # (n_periods,)
 70    """Lower confidence limits for dynamic ATT."""
 71    att_on_ci_upper: np.ndarray             # (n_periods,)
 72    """Upper confidence limits for dynamic ATT."""
 73    att_on_pval: np.ndarray                 # (n_periods,)
 74    """Two-sided p-values for zero dynamic ATT."""
 75    time_on: np.ndarray                     # (n_periods,) relative time indices
 76    """Relative event-time grid."""
 77
 78    # Bootstrap distributions (for advanced use)
 79    att_avg_boot: np.ndarray | None = None  # (nboots,)
 80    """Overall ATT distribution across valid resampled fits."""
 81    att_on_boot: np.ndarray | None = None   # (n_periods, nboots)
 82    """Dynamic ATT distribution, shaped (n_periods, n_valid_replications)."""
 83    n_boot: int = 0
 84    """Number of valid bootstrap draws; zero for jackknife."""
 85    n_boot_failed: int = 0
 86    """Number of failed bootstrap draws; zero for jackknife."""
 87
 88
 89def _make_bootstrap_sampler(unit_type, cluster=None):
 90    """Prepare whole-unit or whole-cluster draws once per bootstrap."""
 91    unit_type = np.asarray(unit_type)
 92    if cluster is not None:
 93        cluster = np.asarray(cluster)
 94        if cluster.ndim != 1 or len(cluster) != len(unit_type):
 95            raise ValueError("cluster must have one label per unit")
 96        if any(label is None or label != label for label in cluster):
 97            raise ValueError("cluster labels must not be missing")
 98        labels, inverse, counts = np.unique(cluster, return_inverse=True, return_counts=True)
 99        if len(labels) < 2:
100            raise ValueError("cluster bootstrap requires at least two clusters")
101        order = np.argsort(inverse, kind="stable")
102        groups = np.split(order, np.cumsum(counts)[:-1])
103
104        def sample(rng):
105            selected = rng.integers(len(groups), size=len(groups))
106            return np.concatenate([groups[i] for i in selected])
107    else:
108        groups = [np.flatnonzero(unit_type == kind) for kind in (2, 1, 3)]
109
110        def sample(rng):
111            return np.concatenate([
112                rng.choice(group, size=len(group), replace=True) if len(group) else np.array([], dtype=int)
113                for group in groups
114            ])
115    return sample
116
117
118def _check_bootstrap_replicates(n_valid, n_requested, label="bootstrap"):
119    """Require an estimable SE and report discarded replications."""
120    n_failed = n_requested - n_valid
121    if n_valid == 0:
122        raise RuntimeError(f"All {label} replications failed ({n_requested} requested)")
123    if n_valid < 2:
124        raise RuntimeError(f"Too few valid {label} replications: {n_valid} of {n_requested}")
125    if n_failed:
126        warnings.warn(
127            f"{n_failed} of {n_requested} {label} replications failed; using {n_valid} valid replications.",
128            RuntimeWarning, stacklevel=3,
129        )
130    return n_failed
131
132
133def bootstrap(
134    estimate_fn,
135    Y: np.ndarray,
136    D: np.ndarray,
137    I: np.ndarray,
138    T_on: np.ndarray,
139    unit_type: np.ndarray,
140    nboots: int = 200,
141    alpha: float = 0.05,
142    n_jobs: int = 1,
143    seed: int | None = None,
144    cluster: np.ndarray | None = None,
145    point_estimate=None,
146) -> InferenceResult:
147    """Non-parametric cluster bootstrap.
148
149    Parameters
150    ----------
151    estimate_fn : callable
152        Function that takes (unit_indices,) and returns:
153        - eff: (T, N_sub) treatment effects
154        - D_sub: (T, N_sub) treatment indicators
155        - T_on_sub: (T, N_sub) relative time
156        - I_sub: (T, N_sub) observation indicators
157    unit_type : (N,) array
158        1=control, 2=treated, 3=reversal
159    cluster : (N,) array, optional
160        Cluster assignments. If None, each unit is its own cluster.
161    point_estimate : tuple, optional
162        Precomputed full-sample ``(eff, D, T_on, I)`` result. If supplied,
163        bootstrap reuses it as the original point estimate and only calls
164        ``estimate_fn`` for resampled replications.
165    """
166    rng = np.random.default_rng(seed)
167    if isinstance(nboots, (bool, np.bool_)) or not isinstance(nboots, (int, np.integer)) or nboots < 2:
168        raise ValueError("nboots must be an integer of at least 2")
169    N = Y.shape[1]
170    sample_units = _make_bootstrap_sampler(unit_type, cluster)
171
172    # Point estimate. fect() already has this before entering inference, so
173    # callers may pass it to avoid re-estimating the unchanged full sample.
174    point_result = point_estimate if point_estimate is not None else estimate_fn(np.arange(N))
175    eff_point, D_point, T_on_point, I_point = _unpack_estimate(
176        point_result, I
177    )
178    att_avg_point, att_on_point, time_on = _compute_att(eff_point, D_point, T_on_point, I_point)
179
180    n_periods = len(time_on)
181
182    def _one_boot(boot_seed):
183        boot_rng = np.random.default_rng(boot_seed)
184
185        unit_idx = sample_units(boot_rng)
186
187        try:
188            eff_b, D_b, T_on_b, I_b = _unpack_estimate(
189                estimate_fn(unit_idx), I[:, unit_idx]
190            )
191            att_avg_b, att_on_b, _ = _compute_att(eff_b, D_b, T_on_b, I_b, time_on)
192            return att_avg_b, att_on_b
193        except Exception:
194            return None, None
195
196    # Generate deterministic seeds for each replicate
197    boot_seeds = rng.integers(0, 2**31, size=nboots)
198
199    if n_jobs == 1:
200        results = [_one_boot(s) for s in boot_seeds]
201    else:
202        with parallel_config(backend="loky", inner_max_num_threads=1):
203            results = Parallel(n_jobs=n_jobs)(
204                delayed(_one_boot)(s) for s in boot_seeds
205            )
206
207    # Filter failed boots
208    valid = [(avg, on) for avg, on in results if avg is not None and np.isfinite(avg)]
209    n_failed = _check_bootstrap_replicates(len(valid), nboots)
210
211    att_avg_boot = np.array([v[0] for v in valid])
212    att_on_boot = np.column_stack([v[1] for v in valid])  # (n_periods, nboots_valid)
213
214    # Pivotal confidence intervals: 2*theta - quantile(1-alpha/2) to 2*theta - quantile(alpha/2)
215    att_avg_se = float(np.std(att_avg_boot, ddof=1))
216    att_avg_ci = (
217        2 * att_avg_point - float(np.quantile(att_avg_boot, 1 - alpha / 2)),
218        2 * att_avg_point - float(np.quantile(att_avg_boot, alpha / 2)),
219    )
220    att_avg_pval = min(
221        2 * np.mean(att_avg_boot >= 0),
222        2 * np.mean(att_avg_boot <= 0),
223        1.0,
224    )
225
226    att_on_se = np.nanstd(att_on_boot, axis=1, ddof=1)
227    att_on_ci_lower = 2 * att_on_point - np.nanquantile(att_on_boot, 1 - alpha / 2, axis=1)
228    att_on_ci_upper = 2 * att_on_point - np.nanquantile(att_on_boot, alpha / 2, axis=1)
229    valid_on = np.isfinite(att_on_boot)
230    valid_counts = valid_on.sum(axis=1)
231    with np.errstate(invalid="ignore", divide="ignore"):
232        boot_ge_0 = np.sum((att_on_boot >= 0) & valid_on, axis=1) / valid_counts
233        boot_le_0 = np.sum((att_on_boot <= 0) & valid_on, axis=1) / valid_counts
234    att_on_pval = np.minimum(2 * boot_ge_0, np.minimum(2 * boot_le_0, 1.0))
235    att_on_pval[valid_counts == 0] = np.nan
236
237    return InferenceResult(
238        att_avg=att_avg_point,
239        att_avg_se=att_avg_se,
240        att_avg_ci=att_avg_ci,
241        att_avg_pval=att_avg_pval,
242        att_on=att_on_point,
243        att_on_se=att_on_se,
244        att_on_ci_lower=att_on_ci_lower,
245        att_on_ci_upper=att_on_ci_upper,
246        att_on_pval=att_on_pval,
247        time_on=time_on,
248        att_avg_boot=att_avg_boot,
249        att_on_boot=att_on_boot,
250        n_boot=len(valid),
251        n_boot_failed=n_failed,
252    )
253
254
255def jackknife(
256    estimate_fn,
257    Y: np.ndarray,
258    D: np.ndarray,
259    I: np.ndarray,
260    T_on: np.ndarray,
261    unit_type: np.ndarray,
262    alpha: float = 0.05,
263    n_jobs: int = 1,
264) -> InferenceResult:
265    """Delete-one-unit jackknife inference.
266
267    Drops one unit at a time, re-estimates, computes jackknife SE
268    with the standard (N-1)/N correction.
269    """
270    from scipy import stats
271
272    N = Y.shape[1]
273
274    # Point estimate
275    eff_point, D_point, T_on_point, I_point = _unpack_estimate(
276        estimate_fn(np.arange(N)), I
277    )
278    att_avg_point, att_on_point, time_on = _compute_att(eff_point, D_point, T_on_point, I_point)
279    n_periods = len(time_on)
280
281    def _one_jack(j):
282        idx = np.concatenate([np.arange(j), np.arange(j + 1, N)])
283        try:
284            eff_j, D_j, T_on_j, I_j = _unpack_estimate(
285                estimate_fn(idx), I[:, idx]
286            )
287            att_avg_j, att_on_j, _ = _compute_att(eff_j, D_j, T_on_j, I_j, time_on)
288            return att_avg_j, att_on_j
289        except Exception:
290            return None, None
291
292    if n_jobs == 1:
293        results = [_one_jack(j) for j in range(N)]
294    else:
295        with parallel_config(backend="loky", inner_max_num_threads=1):
296            results = Parallel(n_jobs=n_jobs)(
297                delayed(_one_jack)(j) for j in range(N)
298            )
299
300    valid = [(avg, on) for avg, on in results if avg is not None]
301    n_valid = len(valid)
302    if n_valid != N:
303        raise RuntimeError(
304            f"{N - n_valid} of {N} jackknife deletions failed; "
305            "delete-one-unit inference requires every deletion."
306        )
307    if n_valid < 2:
308        raise RuntimeError("Too few valid jackknife replications")
309
310    att_avg_jack = np.array([v[0] for v in valid])
311    att_on_jack = np.column_stack([v[1] for v in valid])
312
313    # Jackknife variance: (N-1)/N * sum((theta_j - theta_bar)^2)
314    att_avg_se = float(np.sqrt((n_valid - 1) / n_valid * np.sum((att_avg_jack - np.mean(att_avg_jack)) ** 2)))
315    att_on_se = np.sqrt(
316        (n_valid - 1) / n_valid * np.sum((att_on_jack - att_on_jack.mean(axis=1, keepdims=True)) ** 2, axis=1)
317    )
318
319    # t-distribution CI
320    t_crit = stats.t.ppf(1 - alpha / 2, df=n_valid - 1)
321    att_avg_ci = (att_avg_point - t_crit * att_avg_se, att_avg_point + t_crit * att_avg_se)
322    att_on_ci_lower = att_on_point - t_crit * att_on_se
323    att_on_ci_upper = att_on_point + t_crit * att_on_se
324
325    # p-values from t-distribution
326    t_avg = att_avg_point / max(att_avg_se, 1e-10)
327    att_avg_pval = float(2 * stats.t.sf(abs(t_avg), df=n_valid - 1))
328
329    t_on = att_on_point / np.maximum(att_on_se, 1e-10)
330    att_on_pval = 2 * stats.t.sf(np.abs(t_on), df=n_valid - 1)
331
332    return InferenceResult(
333        att_avg=att_avg_point,
334        att_avg_se=att_avg_se,
335        att_avg_ci=att_avg_ci,
336        att_avg_pval=att_avg_pval,
337        att_on=att_on_point,
338        att_on_se=att_on_se,
339        att_on_ci_lower=att_on_ci_lower,
340        att_on_ci_upper=att_on_ci_upper,
341        att_on_pval=att_on_pval,
342        time_on=time_on,
343        att_avg_boot=att_avg_jack,
344        att_on_boot=att_on_jack,
345    )
346
347
348def permutation_test(
349    estimate_fn,
350    att_avg_observed: float,
351    N: int,
352    T: int,
353    n_perms: int = 200,
354    block_size: int = 1,
355    seed: int | None = None,
356    n_jobs: int = 1,
357) -> tuple[float, np.ndarray]:
358    """Permutation test: shuffle treatment assignment, compare ATTs.
359
360    Returns (p_value, permuted_atts).
361    """
362    rng = np.random.default_rng(seed)
363    perm_seeds = rng.integers(0, 2**31, size=n_perms)
364
365    def _one_perm(perm_seed):
366        try:
367            att_perm = estimate_fn(perm_seed)
368            return att_perm
369        except Exception:
370            return None
371
372    if n_jobs == 1:
373        results = [_one_perm(s) for s in perm_seeds]
374    else:
375        with parallel_config(backend="loky", inner_max_num_threads=1):
376            results = Parallel(n_jobs=n_jobs)(
377                delayed(_one_perm)(s) for s in perm_seeds
378            )
379
380    valid = [r for r in results if r is not None]
381    if len(valid) == 0:
382        return 1.0, np.array([])
383
384    permuted_atts = np.array(valid)
385    p_value = float(np.mean(np.abs(permuted_atts) >= abs(att_avg_observed)))
386    return p_value, permuted_atts
387
388
389# ---------------------------------------------------------------------------
390# ATT computation helpers
391# ---------------------------------------------------------------------------
392
393def _unpack_estimate(result, fallback_I: np.ndarray):
394    """Accept old 3-tuple estimate closures and new 4-tuple closures."""
395    if len(result) == 3:
396        eff, D, T_on = result
397        return eff, D, T_on, fallback_I
398    if len(result) == 4:
399        return result
400    raise ValueError("estimate_fn must return (eff, D, T_on) or (eff, D, T_on, I)")
401
402
403def _compute_att(
404    eff: np.ndarray,     # (T, N) effects
405    D: np.ndarray,       # (T, N) treatment indicator
406    T_on: np.ndarray,    # (T, N) relative time
407    I: np.ndarray | None = None,  # (T, N) observation indicator
408    time_on: np.ndarray | None = None,
409) -> tuple[float, np.ndarray, np.ndarray]:
410    """Compute overall and dynamic ATT from the effect matrix.
411
412    Missing raw outcomes are excluded by ``I``.  This keeps bootstrap and
413    jackknife replicates aligned with point estimation: a model-imputed
414    counterfactual can exist at a missing treated cell, but that cell has
415    no observed treated outcome and therefore cannot enter the ATT.
416
417    Uses ``np.bincount`` for the event-time grouping so the cost is linear
418    in the number of observed cells.  When ``time_on`` is supplied the
419    caller's grid is preserved exactly so bootstrap replicates can be
420    stacked into a ``(n_periods, nboots)`` matrix even if some replicates
421    happen to not observe every event time.
422    """
423    if I is not None and np.shape(I) != np.shape(D) and time_on is None:
424        time_on = np.asarray(I)
425        I = None
426
427    if I is None:
428        observed_mask = np.ones_like(D, dtype=bool)
429    else:
430        observed_mask = I > 0
431
432    treated_mask = (D > 0) & observed_mask
433    n_treated = int(np.sum(treated_mask))
434    if n_treated == 0:
435        raise ValueError("No observed treated cells remain for ATT estimation")
436    att_avg = float(np.sum(eff[treated_mask]) / n_treated)
437
438    ever_treated = np.any(D > 0, axis=0)                            # (N,)
439    all_mask = observed_mask & np.broadcast_to(ever_treated[np.newaxis, :], T_on.shape)
440    if not np.all(np.isfinite(eff[all_mask])):
441        raise ValueError("Non-finite effects in observed ever-treated cells")
442    T_on_flat = T_on[all_mask]
443    eff_flat = eff[all_mask]
444    valid = ~np.isnan(T_on_flat)
445    T_on_flat = T_on_flat[valid]
446    eff_flat = eff_flat[valid]
447
448    if T_on_flat.size == 0:
449        if time_on is None:
450            return att_avg, np.array([]), np.array([])
451        return att_avg, np.full(len(time_on), np.nan), time_on
452
453    T_on_int = T_on_flat.astype(np.int64)
454
455    if time_on is None:
456        offset = int(T_on_int.min())
457        minlength = int(T_on_int.max() - offset + 1)
458        idx = T_on_int - offset
459        sums = np.bincount(idx, weights=eff_flat, minlength=minlength)
460        counts = np.bincount(idx, minlength=minlength)
461        keep = counts > 0
462        time_on = (offset + np.arange(minlength))[keep].astype(np.float64)
463        att_on = sums[keep] / counts[keep]
464        return att_avg, att_on, time_on
465
466    # Caller supplied a grid: produce att_on aligned to time_on so that
467    # bootstrap replicates can be stacked with a fixed shape.
468    time_on_int = np.asarray(time_on, dtype=np.int64)
469    grid_offset = int(time_on_int.min())
470    grid_len = int(time_on_int.max() - grid_offset + 1)
471    # Bin flat values that fall inside the grid range.
472    in_range = (T_on_int >= grid_offset) & (T_on_int < grid_offset + grid_len)
473    T_clip = T_on_int[in_range] - grid_offset
474    e_clip = eff_flat[in_range]
475    sums = np.bincount(T_clip, weights=e_clip, minlength=grid_len)
476    counts = np.bincount(T_clip, minlength=grid_len)
477    gather = time_on_int - grid_offset                              # (n_periods,)
478    att_on = np.full(len(time_on), np.nan)
479    counts_at_grid = counts[gather]
480    valid_grid = counts_at_grid > 0
481    att_on[valid_grid] = sums[gather[valid_grid]] / counts_at_grid[valid_grid]
482    return att_avg, att_on, time_on