pyfector.panel

  1"""
  2Panel data ingestion and reshaping.
  3
  4This module converts long-format panel data (polars, polars LazyFrame, or
  5pandas DataFrames) into the ``(T, N)`` matrix layout used by the estimation
  6engine.  Everything downstream of ``prepare_panel`` operates on dense
  7``(T, N)`` arrays keyed by (unit, time) position rather than by the original
  8unit/time identifiers.
  9
 10Responsibilities
 11----------------
 12* Assign integer unit and time positions via sorted unique values of the
 13  ``index`` columns; original identifiers are preserved on ``PanelData``.
 14* Build the outcome matrix ``Y``, treatment indicator ``D``, observation
 15  indicator ``I`` (1 if present, 0 if missing), and the control mask
 16  ``II = I * (1 - D)``.  Missing ``Y`` entries are stored as a 0 placeholder;
 17  missingness is carried by ``I`` and the MC EM loop overwrites missing cells
 18  with model imputations before using them.
 19* Build the optional covariate tensor ``X`` of shape ``(T, N, p)`` and the
 20  weight matrix ``W`` of shape ``(T, N)``.
 21* Compute relative event time ``T_on`` (treatment onset) and, for reversal
 22  units, ``T_off`` (treatment removal).
 23* Classify each unit as control (1), treated (2), or treatment-reversal (3).
 24* Apply unit-level filters: ``min_T0`` minimum pre-treatment periods and
 25  ``max_missing`` fraction of missing observations.
 26* Compute initial imputation ``Y0`` and initial OLS beta on the control
 27  subsample via :func:`initial_fit`, matching R fect's ``initialFit``.
 28
 29Conventions
 30-----------
 31* Matrices are row-major ``(T, N)`` with ``float64`` elements.
 32* ``force`` encodes the additive fixed-effects choice: ``0 = none``,
 33  ``1 = unit``, ``2 = time``, ``3 = two-way``.
 34* Missing raw outcomes are not zero-valued observations.  The 0 in ``Y`` is
 35  only a dense-array placeholder, while ``I`` defines which cells enter the
 36  fitting objective and treatment-effect averages.
 37* Units with no observed outcome are dropped even when ``max_missing=1.0``;
 38  they provide no untreated fitting information and no observed treated
 39  outcome for ATT construction.
 40* ``T_on`` is ``NaN`` for never-treated units and for always-treated units
 41  (no identifiable event).  For treated units, the first observed switch
 42  from ``D == 0`` to ``D == 1`` defines the reference period and
 43  ``T_on[t] = t - t0`` for every observed ``t``.
 44"""
 45
 46from __future__ import annotations
 47
 48from dataclasses import dataclass, field
 49
 50import numpy as np
 51
 52try:
 53    import polars as pl
 54    HAS_POLARS = True
 55except ImportError:
 56    HAS_POLARS = False
 57
 58try:
 59    import pandas as pd
 60    HAS_PANDAS = True
 61except ImportError:
 62    HAS_PANDAS = False
 63
 64from .backend import get_backend, to_device
 65
 66
 67@dataclass
 68class PanelData:
 69    """Structured panel data ready for estimation."""
 70    # Core matrices (all T × N)
 71    Y: np.ndarray               # outcome
 72    """Outcome matrix, shaped (T, N); I identifies observed cells."""
 73    D: np.ndarray               # treatment indicator (0/1)
 74    """Treatment indicator matrix, shaped (T, N)."""
 75    I: np.ndarray               # observation indicator (1=observed, 0=missing)
 76    """Observed-outcome indicator, shaped (T, N)."""
 77    II: np.ndarray              # control indicator (I with D==1 cells set to 0)
 78    """Untreated fitting mask, shaped (T, N)."""
 79
 80    # Covariates (T × N × p), may be None
 81    X: np.ndarray | None
 82    """Covariate tensor, shaped (T, N, p), when supplied."""
 83
 84    # Weights (T × N), may be None
 85    W: np.ndarray | None
 86    """Observation weights, shaped (T, N), when supplied."""
 87
 88    # Treatment timing (T × N)
 89    T_on: np.ndarray            # relative time to treatment onset
 90    """Relative time from the first observed treatment onset; NaN without an identified onset."""
 91    T_off: np.ndarray | None    # relative time to treatment reversal
 92    """Relative time from treatment reversal when populated."""
 93
 94    # Metadata
 95    unit_ids: np.ndarray        # original unit identifiers
 96    """Retained unit identifiers in matrix-column order."""
 97    time_ids: np.ndarray        # original time identifiers
 98    """Period identifiers in matrix-row order."""
 99    unit_type: np.ndarray       # (N,) 1=control, 2=treated, 3=reversal
100    """Unit classification: 1 control, 2 treated, 3 treatment reversal."""
101    covariate_names: list[str]
102    """Names of supplied covariate columns."""
103    N: int
104    """Number of retained units."""
105    T: int
106    """Number of periods."""
107
108    # Status matrix (T × N): 1=treated, 2=control, 3=missing, 4=removed
109    obs_status: np.ndarray
110    """Cell status: 1 observed treated, 2 observed untreated, 3 missing outcome."""
111    cluster: np.ndarray | None = None
112    """Cluster labels aligned with unit_ids, when supplied."""
113
114
115def prepare_panel(
116    data,
117    Y: str,
118    D: str,
119    index: tuple[str, str],
120    X: list[str] | None = None,
121    W: str | None = None,
122    group: str | None = None,
123    min_T0: int = 1,
124    min_T0_strict: bool = False,
125    max_missing: float = 1.0,
126    cluster: str | None = None,
127) -> PanelData:
128    """Convert long-format data to panel matrices.
129
130    Parameters
131    ----------
132    data : polars.DataFrame, polars.LazyFrame, or pandas.DataFrame
133        Long-format panel data.
134    Y : str
135        Column name for outcome variable.
136    D : str
137        Column name for binary treatment indicator (0/1).
138    index : (str, str)
139        Column names for (unit_id, time_period).
140    X : list of str, optional
141        Column names for time-varying covariates.
142    W : str, optional
143        Column name for observation weights.
144    group : str, optional
145        Column name for cohort/group variable.
146    min_T0 : int
147        Minimum untreated/pre-treatment observed periods. By default this is
148        enforced only for treated and treatment-reversal units.
149    min_T0_strict : bool
150        If True, enforce ``min_T0`` on all units, including controls. If False,
151        controls are kept regardless of ``min_T0`` as long as they have at
152        least one observed outcome and satisfy ``max_missing``.
153    max_missing : float
154        Maximum fraction of missing observations per unit (0-1). A value of
155        1.0 permits sparse units, but units with zero observed outcomes are
156        always removed.
157    cluster : str, optional
158        Cluster column, non-missing and constant within each unit.
159
160    Returns
161    -------
162    PanelData
163        Observed panel matrices and aligned unit metadata.
164    """
165    unit_col, time_col = index
166    X = X or []
167
168    # Convert to Polars if needed
169    df = _to_polars(data)
170
171    # Sort by unit, time
172    df = df.sort([unit_col, time_col])
173
174    # Get unique units and times
175    unit_ids = df[unit_col].unique().sort().to_numpy()
176    time_ids = df[time_col].unique().sort().to_numpy()
177    N = len(unit_ids)
178    TT = len(time_ids)
179
180    cluster_ids = None
181    if cluster is not None:
182        if not isinstance(cluster, str) or cluster not in df.columns:
183            raise ValueError("cluster must name a column in data")
184        labels = df[cluster]
185        if labels.null_count() or (labels.dtype.is_float() and not labels.is_finite().all()):
186            raise ValueError("cluster labels must not be missing or non-finite")
187        assignments = df.select(list(dict.fromkeys([unit_col, cluster]))).unique()
188        if assignments.height != N:
189            raise ValueError("cluster must be constant within each unit")
190        cluster_ids = assignments.sort(unit_col)[cluster].to_numpy()
191
192    # Initialize matrices
193    Y_mat = np.full((TT, N), np.nan, dtype=np.float64)
194    D_mat = np.zeros((TT, N), dtype=np.float64)
195    I_mat = np.zeros((TT, N), dtype=np.float64)
196    W_mat = np.ones((TT, N), dtype=np.float64) if W else None
197
198    X_mat = np.zeros((TT, N, len(X)), dtype=np.float64) if X else None
199
200    # --- Vectorised matrix fill ---
201    # Map each long-form row to (t_idx, u_idx) via searchsorted on the
202    # already-sorted unique id arrays.  This avoids the per-row Python
203    # loop used in earlier versions and is ~100x faster on large panels.
204    units_np = df[unit_col].to_numpy()
205    times_np = df[time_col].to_numpy()
206    y_np = df[Y].to_numpy().astype(np.float64)
207    d_np = df[D].to_numpy().astype(np.float64)
208
209    t_idx = np.searchsorted(time_ids, times_np)
210    u_idx = np.searchsorted(unit_ids, units_np)
211
212    Y_mat[t_idx, u_idx] = y_np
213    D_mat[t_idx, u_idx] = d_np
214    I_mat[t_idx, u_idx] = (~np.isnan(y_np)).astype(np.float64)
215
216    if W_mat is not None:
217        w_np = df[W].to_numpy().astype(np.float64)
218        W_mat[t_idx, u_idx] = w_np
219
220    if X_mat is not None:
221        for k, xname in enumerate(X):
222            X_mat[t_idx, u_idx, k] = df[xname].to_numpy().astype(np.float64)
223
224    # Initialize missing outcome cells with a storage placeholder. The
225    # observation mask I carries missingness; MC EM iterations overwrite
226    # I == 0 cells with current model imputations before using them.
227    Y_mat = np.nan_to_num(Y_mat, nan=0.0)
228
229    # Control indicator: observed AND not treated
230    II_mat = I_mat * (1 - D_mat)
231
232    # --- Vectorised unit classification ---
233    # A unit is treated if any observed period has D == 1.
234    # A treated unit has a reversal if any observed D sequence decreases
235    # from 1 back to 0 at some later observed period.
236    obs_mask = I_mat > 0                                            # (T, N)
237    D_obs = np.where(obs_mask, D_mat, 0.0)
238    ever_treated = np.any(D_obs > 0, axis=0)                        # (N,)
239    # Reversals: diff down from 1->0 on observed cells only
240    # Walk observed D per column; a reversal is any 1 -> 0 transition.
241    D_obs_shift = np.vstack([np.zeros((1, N)), D_obs[:-1]])
242    drops = (D_obs < D_obs_shift) & obs_mask
243    has_reversal = np.any(drops, axis=0) & ever_treated
244
245    unit_type = np.ones(N, dtype=np.int32)
246    unit_type[ever_treated & ~has_reversal] = 2
247    unit_type[has_reversal] = 3
248
249    # --- Vectorised treatment timing (T_on / T_off) ---
250    T_on_mat = _compute_T_on(D_mat, I_mat)
251    T_off_mat = _compute_T_off(D_mat, I_mat, unit_type)
252
253    # --- Vectorised unit filters ---
254    keep = np.ones(N, dtype=bool)
255
256    # Units with no observed outcome cannot contribute to the untreated
257    # fitting objective or to observed treated effects. Drop them before
258    # applying estimand-specific min_T0 rules.
259    observed_counts = I_mat.sum(axis=0)                              # (N,)
260    keep &= observed_counts > 0
261
262    # Min untreated/pre-treatment periods. The default preserves historical
263    # pyfector behavior: only treated / reversal units are filtered. Strict
264    # mode also requires controls to have enough observed untreated cells,
265    # matching R fect's more conservative sparse-panel handling.
266    pre_periods = (II_mat > 0).sum(axis=0)                          # (N,)
267    if min_T0_strict:
268        keep &= pre_periods >= min_T0
269    else:
270        treated_mask = (unit_type == 2) | (unit_type == 3)
271        keep &= ~(treated_mask & (pre_periods < min_T0))
272
273    # Max missing (fraction observed must exceed 1 - max_missing)
274    obs_rate = I_mat.mean(axis=0)                                   # (N,)
275    keep &= obs_rate >= (1.0 - max_missing)
276
277    if not np.any(keep):
278        raise ValueError(
279            "All units were removed by missing-outcome filters. "
280            "Check missing Y values, min_T0, min_T0_strict, and max_missing."
281        )
282
283    if not np.all(keep):
284        kept_idx = np.where(keep)[0]
285        Y_mat = Y_mat[:, kept_idx]
286        D_mat = D_mat[:, kept_idx]
287        I_mat = I_mat[:, kept_idx]
288        II_mat = II_mat[:, kept_idx]
289        T_on_mat = T_on_mat[:, kept_idx]
290        T_off_mat = T_off_mat[:, kept_idx]
291        unit_ids = unit_ids[kept_idx]
292        unit_type = unit_type[kept_idx]
293        if cluster_ids is not None:
294            cluster_ids = cluster_ids[kept_idx]
295        if X_mat is not None:
296            X_mat = X_mat[:, kept_idx, :]
297        if W_mat is not None:
298            W_mat = W_mat[:, kept_idx]
299        N = len(kept_idx)
300
301    # Build status matrix
302    # 1 = treated observed, 2 = control observed, 3 = missing
303    obs_status = np.where(
304        I_mat == 0, 3,
305        np.where(D_mat > 0, 1, 2),
306    ).astype(np.int32)
307    # Note: removed units (status 4) are dropped above, so `obs_status`
308    # after the slice only covers kept units.
309
310    return PanelData(
311        Y=Y_mat, D=D_mat, I=I_mat, II=II_mat,
312        X=X_mat, W=W_mat,
313        T_on=T_on_mat, T_off=T_off_mat,
314        unit_ids=unit_ids, time_ids=time_ids,
315        unit_type=unit_type,
316        covariate_names=X,
317        N=N, T=TT,
318        obs_status=obs_status,
319        cluster=cluster_ids,
320    )
321
322
323def _to_polars(data):
324    """Convert input data to Polars DataFrame."""
325    if HAS_POLARS:
326        if isinstance(data, pl.DataFrame):
327            return data
328        if isinstance(data, pl.LazyFrame):
329            return data.collect()
330    if HAS_PANDAS:
331        if isinstance(data, pd.DataFrame):
332            if HAS_POLARS:
333                return pl.from_pandas(data)
334            # Fallback: use pandas directly via our adapter
335            return _PandasAdapter(data)
336    raise TypeError(
337        f"Expected polars.DataFrame, polars.LazyFrame, or pandas.DataFrame, "
338        f"got {type(data).__name__}"
339    )
340
341
342class _PandasAdapter:
343    """Minimal adapter so panel.py works with pandas if polars unavailable."""
344
345    def __init__(self, df):
346        self._df = df
347
348    def sort(self, cols):
349        return _PandasAdapter(self._df.sort_values(cols).reset_index(drop=True))
350
351    def __getitem__(self, key):
352        return _PandasColumnAdapter(self._df[key])
353
354    def __len__(self):
355        return len(self._df)
356
357
358class _PandasColumnAdapter:
359    def __init__(self, series):
360        self._s = series
361
362    def unique(self):
363        return _PandasColumnAdapter(self._s.drop_duplicates())
364
365    def sort(self):
366        return _PandasColumnAdapter(self._s.sort_values())
367
368    def to_numpy(self):
369        return self._s.to_numpy()
370
371
372def _compute_T_on(D: np.ndarray, I: np.ndarray) -> np.ndarray:
373    """Relative event time to treatment onset, vectorised over units.
374
375    For every unit (column) this computes ``t - t0`` for every observed
376    period, where ``t0`` is the first observed period in which
377    ``D`` switches from 0 to 1 (an observed 0 immediately before an
378    observed 1).  Units whose observed ``D`` sequence is all-zero or
379    all-one, and units with no observed 0->1 transition but whose first
380    observed period already has ``D == 1``, are handled as follows:
381
382    * all-zero (never treated): return NaN.
383    * all-one (always treated, no identifiable onset): return NaN.
384    * already-treated-at-first-observation (no 0->1 switch but some
385      observed period has ``D == 1``): ``t0 = first observed period``.
386
387    Unobserved periods are left as NaN.
388
389    Returns a ``(T, N)`` float matrix, NaN where relative time is
390    undefined.
391
392    This replaces the per-unit Python loop of earlier versions; behaviour
393    is identical to the previous ``_get_term(..., kind="on")`` on all the
394    DGPs used by the test suite and by R fect's examples.
395    """
396    T, N = D.shape
397    obs = I > 0
398    D_obs = np.where(obs, D, 0.0)
399
400    # 0->1 transition on observed cells.  Previous-observed-D is taken
401    # with ``prepend=0``, matching the C++ loop-start condition.
402    D_prev = np.vstack([np.zeros((1, N)), D_obs[:-1]])
403    obs_prev = np.vstack([np.zeros((1, N), dtype=bool), obs[:-1]])
404    switch_on = (D_obs > 0) & (D_prev == 0) & obs & obs_prev      # (T, N)
405
406    # First switch index per unit; columns with no switch get a sentinel
407    # value we'll replace below.
408    has_switch = switch_on.any(axis=0)                             # (N,)
409    t0 = np.where(has_switch, switch_on.argmax(axis=0), -1)        # (N,)
410
411    # Fallback: no switch but first observed period already has D == 1.
412    no_switch = ~has_switch
413    if no_switch.any():
414        # first observed period per column, NaN-safe
415        first_obs_idx = np.where(obs.any(axis=0), obs.argmax(axis=0), -1)
416        # treated at first obs?
417        at_first = np.zeros(N, dtype=bool)
418        valid = first_obs_idx >= 0
419        at_first[valid] = D[first_obs_idx[valid], np.arange(N)[valid]] > 0
420        use_first = no_switch & at_first
421        t0 = np.where(use_first, first_obs_idx, t0)
422
423    # Also exclude never-treated and always-treated units (no obs with D
424    # differing from the constant value across its observed window).
425    # Never-treated: all observed D == 0 -> no onset.
426    # Always-treated: all observed D == 1 and no prior 0 to switch from.
427    # Already handled: never-treated falls through to t0 == -1;
428    # always-treated without a fallback hit also t0 == -1.
429    valid_col = t0 >= 0                                             # (N,)
430
431    term = np.full((T, N), np.nan)
432    if valid_col.any():
433        t_idx = np.arange(T)[:, None]                               # (T, 1)
434        rel = (t_idx - t0[None, :]).astype(np.float64)              # (T, N)
435        # Only observed cells of valid columns get a value.
436        mask = obs & valid_col[None, :]
437        term = np.where(mask, rel, np.nan)
438    return term
439
440
441def _compute_T_off(
442    D: np.ndarray, I: np.ndarray, unit_type: np.ndarray
443) -> np.ndarray:
444    """Relative time to first treatment reversal, vectorised over units.
445
446    Only units classified as reversal (``unit_type == 3``) get a
447    non-NaN column.  ``t0`` is the first observed 1->0 transition; the
448    resulting relative time is ``t - t0`` at every observed period.
449    """
450    T, N = D.shape
451    obs = I > 0
452    D_obs = np.where(obs, D, 0.0)
453
454    D_prev = np.vstack([np.zeros((1, N)), D_obs[:-1]])
455    obs_prev = np.vstack([np.zeros((1, N), dtype=bool), obs[:-1]])
456    switch_off = (D_obs == 0) & (D_prev > 0) & obs & obs_prev
457
458    has_switch = switch_off.any(axis=0)
459    t0 = np.where(has_switch, switch_off.argmax(axis=0), -1)
460
461    valid_col = (unit_type == 3) & (t0 >= 0)
462    term = np.full((T, N), np.nan)
463    if valid_col.any():
464        t_idx = np.arange(T)[:, None]
465        rel = (t_idx - t0[None, :]).astype(np.float64)
466        mask = obs & valid_col[None, :]
467        term = np.where(mask, rel, np.nan)
468    return term
469
470
471def initial_fit(
472    Y: np.ndarray,
473    X: np.ndarray | None,
474    II: np.ndarray,
475    force: int,
476) -> tuple[np.ndarray, np.ndarray | None]:
477    """Compute initial fitted values (Y0) and beta0 using OLS on control obs.
478
479    Equivalent to R fect's ``initialFit()`` using fixest::feols.
480    Uses simple demeaning on control observations for initial FE estimates.
481    """
482    xp = get_backend()
483    T, N = Y.shape
484
485    # Compute means only over control observations
486    Y_ctrl = Y * II
487    ctrl_count = II.copy()
488
489    # Grand mean of control obs
490    total = xp.sum(Y_ctrl)
491    n_ctrl = xp.sum(ctrl_count)
492    mu = total / (n_ctrl + 1e-10)
493
494    # Unit means
495    alpha = None
496    if force in (1, 3):
497        col_sum = xp.sum(Y_ctrl, axis=0)
498        col_n = xp.sum(ctrl_count, axis=0) + 1e-10
499        alpha = col_sum / col_n
500
501    # Time means
502    xi = None
503    if force in (2, 3):
504        row_sum = xp.sum(Y_ctrl, axis=1)
505        row_n = xp.sum(ctrl_count, axis=1) + 1e-10
506        xi = row_sum / row_n
507
508    # Build Y0
509    Y0 = xp.full_like(Y, float(mu))
510    if alpha is not None:
511        Y0 = Y0 + (alpha - mu)[None, :]
512    if xi is not None:
513        Y0 = Y0 + (xi - mu)[:, None]
514
515    # Beta from control obs via OLS: (Y - FE) ~ X on the II subset.
516    # Vectorised via einsum: replaces the p*p Python loop of earlier
517    # versions with a single BLAS call per matrix.  Because II is 0/1,
518    # multiplying either operand by II is equivalent to restricting the
519    # sum to control observations.
520    beta0 = None
521    if X is not None and X.shape[2] > 0:
522        p = X.shape[2]
523        X_w = X * II[:, :, None]                                    # (T, N, p)
524        xx = xp.einsum("tni,tnj->ij", X_w, X)                       # (p, p)
525        xy = xp.einsum("tnk,tn->k", X_w, Y - Y0)                    # (p,)
526        try:
527            beta0 = xp.linalg.solve(xx, xy)
528        except np.linalg.LinAlgError:
529            beta0 = xp.zeros(p)
530        # Update Y0 with covariate contribution X @ beta0
531        Y0 = Y0 + xp.einsum("tnk,k->tn", X, beta0)
532
533    return Y0, beta0