diff-stream 0.1.0.dev0__py3-none-any.whl

diff_stream/linear.py

@@ -0,0 +1,339 @@
+"""Covariate/extra-FE SA and BJS with local weighted unit elimination.
+
+Unit intercepts never enter a global table. Other terms use the bounded family
+model-matrix compiler; numerical state depends on columns and effect cells.
+"""
+
+import numpy as np
+
+from ._core import covariate_kernels, twostage_kernels
+from ._core.nuisance import solve_checked
+from ._core.within import center_groups, independent_columns, moments, row_sums, weighted_squares
+
+
+def fixed_effect_terms(rhs, panel):
+    """(covariates, additional fixed-effect terms as lists of column names)."""
+    pieces = rhs.split("|")
+    if len(pieces) != 2:
+        raise ValueError("use one fixed-effect section; IV formulas are unsupported")
+    covariates, effects = (p.strip() for p in pieces)
+    if not effects:
+        raise ValueError("the fixed-effect section cannot be empty")
+    terms = []
+    for term in effects.split("+"):
+        names = [s.strip() for s in term.strip().split("^")]
+        if any(not s.isidentifier() for s in names):
+            raise ValueError(
+                "additional fixed effects must be column names or name^name interactions"
+            )
+        if len(names) == 1 and names[0] in panel:
+            continue  # Unit/time intercepts are always included.
+        if panel[0] in names:
+            raise ValueError("unit-interacted fixed effects are not supported")
+        terms.append(names)
+    return covariates or "1", terms
+
+
+def regression_formula(rhs, panel):
+    pieces = rhs.split("|")
+    if len(pieces) != 2:
+        raise ValueError("use one fixed-effect section; IV formulas are unsupported")
+    covariates, effects = (p.strip() for p in pieces)
+    if not effects:
+        raise ValueError("the fixed-effect section cannot be empty")
+    terms = []
+    for term in effects.split("+"):
+        names = [s.strip() for s in term.strip().split("^")]
+        if any(not s.isidentifier() for s in names):
+            raise ValueError(
+                "additional fixed effects must be column names or name^name interactions"
+            )
+        if len(names) == 1 and names[0] in panel:
+            continue  # Unit/time intercepts are always included.
+        if panel[0] in names:
+            raise ValueError("unit-interacted fixed effects are not supported")
+        terms.append(":".join(f"C({name})" for name in names))
+    return " + ".join([covariates, *terms])
+
+
+class RowDesign:
+    def __init__(self, store):
+        self.store = store
+        self.time = np.eye(store.periods)[:, 1:]
+        names = store.design.names if store.design is not None else []
+        self.names = [f"time::{t}" for t in store.times[1:]] + list(names)
+        self.keep = np.arange(len(self.names))
+        self.dropped = []
+
+    def raw_batches(self, *, actual_cohort=False):
+        """The store's batches (g, y, x, w) with covariates only, for the compiled passes."""
+        for g, y, x, w in self.store.batches(
+            include_x=True, include_weights=True, row_weights=True, actual_cohort=actual_cohort
+        ):
+            yield g, y, (x if self.store.design is not None else x[:, :, :0]), w
+
+    def batches(self, *, actual_cohort=False):
+        for g, y, x, w in self.store.batches(
+            include_x=True, include_weights=True, row_weights=True, actual_cohort=actual_cohort
+        ):
+            x = x if self.store.design is not None else x[:, :, :0]
+            X = np.concatenate(
+                [np.broadcast_to(self.time, (*y.shape, self.time.shape[1])), x], axis=2
+            )
+            yield g, y, np.nan_to_num(X, nan=0.0)[:, :, self.keep], w
+
+    def reduce(self, clean=None, source=None):
+        """Select identified columns; optionally also the clean-row moments.
+
+        ``clean(g)`` gives the rows (units by periods) of a second, clean-sample
+        regression. Its within-unit moments are accumulated in the same pass
+        over all columns and returned restricted to the identified ones:
+        centering acts column by column, so restricting afterwards is exact.
+        """
+        P = len(self.names)
+        if twostage_kernels.ENABLED:
+            kernels = twostage_kernels.TwoStageKernels(self.store.times, 0)
+            Q, _, norms, Qc, qc = kernels.reduce(source or self.raw_batches, P, clean)
+        else:
+            Q = np.zeros((P, P))
+            norms = np.zeros(P)
+            Qc, qc = (np.zeros((P, P)), np.zeros(P)) if clean is not None else (None, None)
+            for g, y, X, w in self.batches():
+                y = np.nan_to_num(y, nan=0.0)
+                q, _ = moments(X, y, w)
+                Q += q
+                norms += weighted_squares(X, w)
+                if clean is not None:
+                    a, b = moments(X, y, w * clean(g))
+                    Qc += a
+                    qc += b
+        self.keep = independent_columns(Q, original_norm=norms)
+        self.dropped = [n for j, n in enumerate(self.names) if j not in self.keep]
+        if not len(self.keep):
+            raise ValueError("the within-unit regression has no identified nuisance columns")
+        if clean is not None:
+            return Qc[np.ix_(self.keep, self.keep)], qc[self.keep]
+
+    def residuals(self, g, y, X, w, theta, anticipation, *, label="BJS"):
+        observed = np.isfinite(y)
+        use = observed & ((g[:, None] == 0) | (self.store.times < g[:, None] - anticipation))
+        clean_weight = w * use
+        raw = np.where(observed, y - X @ theta, 0)
+        mass = clean_weight.sum(axis=1)
+        if np.any(mass <= 0):
+            raise ValueError(f"{label} requires at least one untreated observation per unit")
+        xm = row_sums(clean_weight, X) / mass[:, None]
+        alpha = (clean_weight * raw).sum(axis=1) / mass
+        residual = np.where(observed, raw - alpha[:, None], 0)
+        first_score = row_sums(clean_weight * residual, X - xm[:, None])
+        return observed, use, xm, residual, first_score
+
+
+def fit_bjs(store, cells, avgeffectsby, leaveout, workdir, anticipation=0):
+    from .unbalanced import ObservedBJSInference
+
+    model = RowDesign(store)
+    try:
+        Q, q = model.reduce(
+            clean=lambda g: (g[:, None] == 0) | (store.times < g[:, None] - anticipation)
+        )
+    except ValueError as error:
+        if "positive-weight estimation observation" not in str(error):
+            raise
+        raise ValueError(
+            "BJS requires at least one untreated observation per unit; "
+            "exclude units without one explicitly"
+        ) from None
+    P, M = len(model.keep), len(cells)
+    theta = solve_checked(
+        Q,
+        q,
+        "BJS untreated design is unidentified (a disconnected observed design, "
+        "or a collinear covariate or fixed effect)",
+    )
+    counts = np.array([store.cell_weights[g][t - store.times[0]] for g, t in cells])
+    beta, D = np.zeros(M), np.zeros((M, P))
+    if covariate_kernels.ENABLED:
+        kernels = covariate_kernels.BJSKernels(store, cells, model.keep, theta, anticipation)
+        beta, D = kernels.cell_sums(model.raw_batches, counts)
+    else:
+        for g, y, X, w in model.batches():
+            obs, _, xm, resid, _ = model.residuals(g, y, X, w, theta, anticipation)
+            for j, (cg, t) in enumerate(cells):
+                ti = t - store.times[0]
+                mask = (g == cg) & obs[:, ti]
+                sw = w[mask, ti] / counts[j]
+                beta[j] += sw @ resid[mask, ti]
+                D[j] += sw @ (X[mask, ti] - xm[mask])
+    H = solve_checked(Q, D.T, "BJS untreated covariate/fixed-effect design").T
+    context = ObservedBJSInference(
+        store,
+        cells,
+        counts,
+        None,
+        theta,
+        H,
+        avgeffectsby,
+        leaveout,
+        workdir,
+        anticipation,
+        model=model,
+    )
+    try:
+        V = context.covariance(np.eye(M), allow_unidentified=avgeffectsby != "cohort_time")
+    except BaseException:
+        context.cleanup()
+        raise
+    return beta, counts, V, context, model
+
+
+def fit_sunab(store, cells, anticipation=0):
+    model = RowDesign(store)
+    model.reduce()
+    K, M = len(model.keep), len(cells)
+    P = K + M
+    if covariate_kernels.ENABLED:
+        kernels = covariate_kernels.SunAbrahamKernels(store, cells, model.keep)
+        Q, q = kernels.moments(model.raw_batches)
+        coefficients = solve_checked(
+            Q, q, "Sun–Abraham covariate/fixed-effect interactions are unidentified"
+        )
+        H = solve_checked(Q, np.eye(P)[:, K:], "Sun–Abraham joint design")
+        V = np.zeros((M, M))
+        for raw in model.raw_batches():
+            score = kernels.scores(raw, coefficients, H)
+            V += score.T @ score
+        counts = np.array([store.cell_weights[g][t - store.times[0]] for g, t in cells])
+        return coefficients[K:], counts, (V + V.T) / 2, model
+    Q, q = np.zeros((P, P)), np.zeros(P)
+
+    def batches():
+        for g, y, X, w in model.batches():
+            for cg in np.unique(g):
+                mask = g == cg
+                js = [j for j, (cohort, _) in enumerate(cells) if cohort == cg]
+                interaction = (
+                    np.column_stack([store.times == cells[j][1] for j in js])
+                    if js
+                    else np.empty((store.periods, 0))
+                )
+                active = np.r_[np.arange(K), K + np.array(js, dtype=int)]
+                design = np.concatenate(
+                    [X[mask], np.broadcast_to(interaction, (int(mask.sum()), *interaction.shape))],
+                    axis=2,
+                )
+                yield mask, active, design, y[mask], w[mask]
+
+    for _, active, X, y, w in batches():
+        a, b = moments(X, np.nan_to_num(y, nan=0.0), w)
+        Q[np.ix_(active, active)] += a
+        q[active] += b
+    coefficients = solve_checked(
+        Q, q, "Sun–Abraham covariate/fixed-effect interactions are unidentified"
+    )
+    H = solve_checked(Q, np.eye(P)[:, K:], "Sun–Abraham joint design")
+    V = np.zeros((M, M))
+    for _, active, X, y, w in batches():
+        xc, yc, _, _ = center_groups(X, np.nan_to_num(y, nan=0.0), w)
+        residual = yc - xc @ coefficients[active]
+        score = row_sums(w * residual, xc) @ H[active]
+        V += score.T @ score
+    counts = np.array([store.cell_weights[g][t - store.times[0]] for g, t in cells])
+    return coefficients[K:], counts, (V + V.T) / 2, model
+
+
+def fit_placebos(store, pretrends, anticipation, *, fml, depvar, panel, vcov):
+    """Separate clean lead regression with the same richer nuisance design."""
+    import polars as pl
+
+    from .bjs_placebo import BJSPlaceboResult
+
+    model = RowDesign(store)
+    model.reduce()
+    L, K = pretrends, len(model.keep)
+    horizons = np.arange(-anticipation - L, -anticipation, dtype=np.int64)
+    support, rows = np.zeros(L, dtype=int), []
+    for g, counts in sorted(store.actual_cell_counts.items()):
+        if not g:
+            continue
+        for ti, t in enumerate(store.times):
+            j = int(t) - g + anticipation + L
+            if 0 <= j < L and counts[ti]:
+                support[j] += counts[ti]
+                rows.append(
+                    dict(
+                        aggregate=f"event::{horizons[j]}",
+                        cohort=int(g),
+                        time=int(t),
+                        event_time=int(horizons[j]),
+                        observed_units=int(counts[ti]),
+                        weight_total=float(
+                            store.actual_cell_weights[g][ti]
+                            * store.weight_scale
+                            / store.weight_normalizer
+                        ),
+                    )
+                )
+    if np.any(support == 0):
+        raise ValueError(
+            "BJS requested clean placebo horizons are not observed; reduce pretrend_leads"
+        )
+
+    def batches():
+        for g, y, X, w in model.batches(actual_cohort=True):
+            lead = ((store.times[None, :, None] - g[:, None, None]) == horizons) & (
+                g[:, None, None] != 0
+            )
+            use = np.isfinite(y) & ((g[:, None] == 0) | (store.times < g[:, None] - anticipation))
+            yield np.concatenate([X, lead], axis=2), np.nan_to_num(y, nan=0.0), w * use, use
+
+    P = K + L
+    Q, q = np.zeros((P, P)), np.zeros(P)
+    n_obs = 0
+    for X, y, w, use in batches():
+        a, b = moments(X, y, w)
+        Q += a
+        q += b
+        n_obs += int(use.sum())
+    keep = independent_columns(Q)
+    if not np.isin(np.arange(K, P), keep).all():
+        raise ValueError("BJS clean placebo design is unidentified; reduce pretrend_leads")
+    Q = Q[np.ix_(keep, keep)]
+    context = "BJS clean placebo covariate/fixed-effect design"
+    coefficient = solve_checked(Q, q[keep], context)
+    H = solve_checked(Q, np.eye(len(keep))[:, -L:], context)
+    V = np.zeros((L, L))
+    for X, y, w, _ in batches():
+        xc, yc, _, _ = center_groups(X[:, :, keep], y, w)
+        score = row_sums(w * (yc - xc @ coefficient), xc) @ H
+        V += score.T @ score
+    V *= store.n_units / (store.n_units - 1) if vcov == "CRV1" else 1.0
+    all_names = [*(model.names[j] for j in model.keep), *(f"event::{h}" for h in horizons)]
+    return BJSPlaceboResult(
+        fml=fml,
+        depvar=depvar,
+        coefnames=[f"event::{h}" for h in horizons],
+        beta=coefficient[-L:],
+        vcov=(V + V.T) / 2,
+        vcov_type=vcov,
+        df_t=store.n_units - 1,
+        all_vcovs={vcov: (V, store.n_units - 1)},
+        n_obs=n_obs,
+        n_clusters={panel[0]: store.n_units},
+        n_units=store.n_units,
+        weights=None,
+        _support=pl.DataFrame(rows),
+        diagnostics=dict(
+            aggregation="event_time",
+            anticipation=anticipation,
+            reference_event_time=None,
+            fitted_horizons=horizons.tolist(),
+            pretrends=L,
+            base_period="all other clean observations; no fixed event-time reference",
+            inference="unit clustered; clean-sample covariate/fixed-effect WLS regression",
+            regression_columns=[all_names[j] for j in keep],
+            absorbed_or_redundant_columns=model.dropped
+            + [n for j, n in enumerate(all_names) if j not in keep],
+            warnings=[],
+        ),
+    )

diff_stream/periods.py

@@ -0,0 +1,72 @@
+"""Period labels, and cohorts from a treatment column.
+
+Panels with evenly spaced periods (every year, every second year, ...) are
+numbered consecutively inside the store, starting at the first period, so the
+estimators count periods: event time 1 is the next period, whatever its
+spacing, and anticipation and pre-trend leads are numbers of periods. Results
+show the original labels of cohorts and periods.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class PeriodLabels:
+    """Original period = origin + (internal period - origin) * step."""
+
+    origin: int
+    step: int = 1
+
+    def internal(self, t):
+        """Internal periods of on-grid original periods."""
+        return self.origin + (np.asarray(t) - self.origin) // self.step
+
+    def internal_cohorts(self, g, last_internal):
+        """Internal cohorts; 0 stays 0. Cohorts within the sample must be periods of it."""
+        g = np.asarray(g).copy()
+        nonzero = g != 0
+        offset = g[nonzero] - self.origin
+        inside = (offset > 0) & (g[nonzero] <= self.label(last_internal))
+        off_grid = inside & (offset % self.step != 0)
+        if off_grid.any():
+            raise ValueError(
+                f"cohort {int(g[nonzero][off_grid][0])} is not one of the panel's periods "
+                f"({self.origin}, {self.origin + self.step}, ...): give the first period in "
+                "which the unit is observed treated, or use treatment="
+            )
+        # Cohorts after the sample stay after it (they are never treated within it).
+        g[nonzero] = self.origin - (-offset // self.step)
+        return g
+
+    def label(self, t):
+        """Original labels of internal periods (scalars or arrays)."""
+        if self.step == 1:
+            return t
+        return self.origin + (np.asarray(t) - self.origin) * self.step
+
+    def label_int(self, t):
+        return int(self.label(int(t)))
+
+    def label_cohort(self, g):
+        """Original cohort label; 0 (never treated) stays 0."""
+        return 0 if int(g) == 0 else self.label_int(g)
+
+
+def cohorts_from_treatment(treatment, t, ids, starts, unit):
+    """Each row's cohort: its unit's first period with treatment 1, else 0.
+
+    Rows are sorted by unit and period. Treatment must be 0/1 and absorbing.
+    """
+    if np.any((treatment != 0) & (treatment != 1)):
+        raise ValueError("treatment must be 0 or 1")
+    same_unit = ids[1:] == ids[:-1]
+    if np.any(same_unit & (treatment[1:] < treatment[:-1])):
+        raise ValueError(
+            "treatment must be absorbing: once 1, a unit stays treated in its later periods"
+        )
+    never = np.iinfo(np.int64).max
+    first = np.minimum.reduceat(np.where(treatment == 1, t, never), starts)
+    first[first == never] = 0
+    return first[unit]

diff_stream/pretrend.py

@@ -0,0 +1,97 @@
+"""Clean pre-treatment selection and transparent joint Wald tests."""
+
+import numpy as np
+import polars as pl
+
+from ._present import internal
+from .aggregation import integer
+
+
+def pretrend_test(result, *, level="auto", horizons=None):
+    if result.estimator == "bjs":
+        return result.placebos().pretrend_test(level=level, horizons=horizons)
+    if level not in ("auto", "cell", "event_time"):
+        raise ValueError("level must be 'auto', 'cell' or 'event_time'")
+    A = internal(result).get("anticipation", 0)
+    if result.cells:
+        level = "cell" if level == "auto" else level
+        events = np.array([t - g for g, t in result.cells])
+    else:
+        if internal(result).get("aggregation") != "event_time" or level == "cell":
+            raise ValueError(
+                "pretrend_test requires cohort-time estimates or an event-time aggregate"
+            )
+        level = "event_time"
+        events = np.array([int(name.split("::")[-1]) for name in result.coefnames])
+    clean = events < -A
+    requested = None
+    if horizons is not None:
+        try:
+            requested = list(horizons)
+        except TypeError as exc:
+            raise ValueError(
+                "horizons must contain clean pre-treatment integer event times"
+            ) from exc
+        if not requested or not all(integer(h) and h < -A for h in requested):
+            raise ValueError(
+                "horizons must contain clean pre-treatment integer event times < -anticipation"
+            )
+        requested = sorted(set(int(h) for h in requested))
+        missing = set(requested) - set(events[clean].tolist())
+        if missing:
+            raise ValueError(f"requested clean horizons are not estimated: {sorted(missing)}")
+        clean &= np.isin(events, requested)
+    if not clean.any():
+        raise ValueError(
+            "no estimated clean pre-treatment coefficients; normalized references and anticipation periods are excluded"
+        )
+    if result.cells and level == "event_time":
+        result = result.aggregate(horizons=sorted(set(events[clean].tolist())))
+        events = np.array([int(name.split("::")[-1]) for name in result.coefnames])
+        clean = np.ones(len(events), dtype=bool)
+    ix = np.flatnonzero(clean)
+    names = [result.coefnames[j] for j in ix]
+    V = result.vcov[np.ix_(ix, ix)]
+    sd = np.sqrt(np.maximum(np.diag(V), 0))
+    if not np.isfinite(V).all() or np.any(sd <= 0):
+        raise ValueError(
+            "pre-trend covariance has zero or non-finite variance; select estimable horizons"
+        )
+    correlation = V / np.outer(sd, sd)
+    if np.linalg.eigvalsh((correlation + correlation.T) / 2)[0] <= 1e-10:
+        raise ValueError(
+            "pre-trend covariance is singular or ill-conditioned; select fewer cells/horizons or use level='event_time'"
+        )
+    # Rescaling restrictions leaves the null unchanged and lets the inherited
+    # Wald solver work on the correlation scale. Otherwise its pseudoinverse
+    # can discard a small-variance restriction solely because of units/scales.
+    R = np.zeros((len(ix), len(result.beta)))
+    R[np.arange(len(ix)), ix] = 1 / sd
+    test = result.wald_test(R=R, distribution="chi2")
+    if result.cells:
+        cells = [dict(cohort=int(result.cells[j][0]), time=int(result.cells[j][1])) for j in ix]
+    else:
+        table = (
+            result.support() if result.estimator == "bjs_placebo" else result.contributions()
+        ).filter(pl.col("aggregate").is_in(names))
+        cells = table.select("cohort", "time").unique().sort(["cohort", "time"]).to_dicts()
+    return dict(
+        **test,
+        level=level,
+        coefficients=names,
+        horizons=sorted(set(int(h) for h in events[ix])),
+        cells=cells,
+        n_restrictions=len(ix),
+        anticipation=A,
+        reference_event_time=internal(result).get("reference_event_time", -A - 1),
+        null="all selected clean pre-treatment effects equal zero",
+        base_period=internal(result).get("base_period"),
+        **(
+            dict(
+                fitted_horizons=result.diagnostics["fitted_horizons"],
+                sample="separate clean-sample lead regression",
+            )
+            if result.estimator == "bjs_placebo"
+            else {}
+        ),
+    )

diff_stream/reference.py

@@ -0,0 +1,125 @@
+"""Explicit clean last-treated reference windows on an already validated panel."""
+
+import warnings
+from dataclasses import dataclass
+
+import polars as pl
+
+from .store import PanelStore
+
+
+class ReferenceWindowWarning(UserWarning):
+    """Estimation uses only dates before the last reference cohort is affected."""
+
+
+@dataclass
+class FrozenDesign:
+    """Reuse validated model columns when restaging a restricted panel."""
+
+    names: list
+
+    def discover(self, frame):
+        pass
+
+    def compile(self):
+        pass
+
+    def evaluate(self, frame):
+        return frame.select([f"__x{j}" for j in range(len(self.names))])
+
+
+def reference_store(store, policy, anticipation):
+    """Auto selects last-treated only when no clean never-treated group exists.
+
+    Complete-panel selection precedes reference selection. The reference cohort
+    and existing never-treated units are internally pooled as clean controls;
+    the original reference date and all restrictions remain explicit metadata.
+    """
+    use_last = policy == "lasttreated" or (policy == "auto" and 0 not in store.gc)
+    if not use_last:
+        return store, dict(
+            reference_group_requested=policy,
+            reference_group="nevertreated",
+            reference_cohort=0,
+            reference_cutoff=None,
+            reference_dropped_rows=0,
+            reference_dropped_units=0,
+        )
+    treated = [g for g in store.gc if g]
+    if len(treated) < 2:
+        raise ValueError(
+            "last-treated reference requires at least two treated cohorts after panel selection"
+        )
+    reference = max(treated)
+    cutoff = reference - anticipation
+    if cutoff - int(store.times[0]) < 2:
+        raise ValueError(
+            "fewer than two clean dates remain before the last-treated reference is affected"
+        )
+    data = (
+        pl.scan_parquet(store.paths, low_memory=True)
+        .filter(pl.col("t") < cutoff)
+        .with_columns(pl.when(pl.col("g") == reference).then(0).otherwise(pl.col("g")).alias("g"))
+    )
+    restricted = restage(
+        store,
+        data,
+        "reference",
+        "balanced" if store.panel_mode in ("balanced", "complete") else "unbalanced",
+        anticipation,
+    )
+    dropped_rows = store.n_rows_used - restricted.n_rows_used
+    dropped_units = store.n_units - restricted.n_units
+    metadata = dict(
+        reference_group_requested=policy,
+        reference_group="lasttreated",
+        reference_cohort=int(reference),
+        reference_cohort_units=store.gc[reference],
+        reference_never_units=store.gc.get(0, 0),
+        input_time_min=int(store.times[0]),
+        input_time_max=int(store.times[-1]),
+        reference_cutoff=int(cutoff),
+        reference_last_clean_time=int(restricted.times[-1]),
+        reference_dropped_rows=dropped_rows,
+        reference_dropped_units=dropped_units,
+        reference_includes_never_treated=0 in store.gc,
+        reference_selection="after original balanced/complete/unbalanced panel policy; before reference is affected",
+    )
+    for path in store.paths:
+        path.unlink()
+    warnings.warn(
+        f"Last-treated cohort {reference} is a reference only: retain time < {cutoff} (anticipation={anticipation}); "
+        f"excluded {dropped_rows} observed rows. Effects after the reference becomes affected are not identified by this design.",
+        ReferenceWindowWarning,
+        stacklevel=3,
+    )
+    return restricted, metadata
+
+
+def restage(store, data, name, panel_mode, anticipation):
+    """Stage a subset of a store's rows (``data``, with the store's columns) as
+    a new store, keeping the input counts and the original period labels."""
+    root = store.root / name
+    root.mkdir()
+    data = data.with_columns(pl.col("u").alias("id")).drop("u")
+    if store.weighted:
+        data = data.with_columns(pl.col("w") * store.weight_scale)
+    design = FrozenDesign(list(store.design.names)) if store.design is not None else None
+    restricted = PanelStore(
+        data,
+        root,
+        store.batch_units,
+        store.budget,
+        store.leaf_bytes * 4 / 1e6,
+        design,
+        panel_mode,
+        anticipation,
+        allow_time_varying_weights=store.allow_time_varying_weights,
+        retain_actual_cohort=store.retain_actual_cohort,
+        actual_cohort_column="g_actual" if store.retain_actual_cohort else None,
+    )
+    # The restaged rows are already numbered consecutively; keep the original labels.
+    restricted.period_labels = store.period_labels
+    restricted.n_rows = store.n_rows
+    restricted.n_units_input = store.n_units_input
+    return restricted