bossanova 0.1.0.dev21__tar.gz → 0.1.0.dev23__tar.gz

{bossanova-0.1.0.dev21 → bossanova-0.1.0.dev23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bossanova
-Version: 0.1.0.dev21
+Version: 0.1.0.dev23
 Summary: Bridging statistical cultures with some jazz
 Author: Eshin Jolly
 License-Expression: MIT

{bossanova-0.1.0.dev21 → bossanova-0.1.0.dev23}/bossanova/internal/infer/bootstrap.py

@@ -104,6 +104,20 @@ def compute_bootstrap_params(
     # Generate all bootstrap indices at once
     boot_indices = generate_bootstrap_indices(rng, n, n_boot)
 
+    # LM fast path: batched OLS via chunked vmap
+    is_lm = (
+        spec.family == "gaussian"
+        and spec.link == "identity"
+        and not spec.has_random_effects
+    )
+
+    if is_lm:
+        from bossanova.internal.infer.resample.lm_operators import (
+            compute_batched_ols_bootstrap,
+        )
+
+        return compute_batched_ols_bootstrap(boot_indices, bundle.X, bundle.y)
+
     if n_jobs == 1:
         # Sequential execution
         boot_coefs = np.zeros((n_boot, p))

{bossanova-0.1.0.dev21 → bossanova-0.1.0.dev23}/bossanova/internal/infer/permutation.py

@@ -184,6 +184,42 @@ def explore_with_fit(
     )
 
 
+def _compute_permutation_lm_fast(
+    X: np.ndarray,
+    y: np.ndarray,
+    n_perm: int,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Compute all permutation coefficients in a single matmul.
+
+    For LM models, coefficients are a linear map of y when X is fixed:
+    ``coef = (X'X)^{-1} X' y``. Pre-computing the factorization once
+    and applying it to all permuted y vectors simultaneously avoids
+    the per-permutation solver overhead.
+
+    Args:
+        X: Design matrix, shape ``(n, p)``.
+        y: Original response vector, shape ``(n,)``.
+        n_perm: Number of permutations.
+        rng: NumPy random generator for reproducibility.
+
+    Returns:
+        Array of null coefficients, shape ``(n_perm, p)``.
+    """
+    from bossanova.internal.infer.resample.lm_operators import (
+        build_lm_coefficient_operator,
+    )
+
+    n = len(y)
+    perm_indices = np.stack([rng.permutation(n) for _ in range(n_perm)])
+    Y_perm = y[perm_indices]  # (n_perm, n)
+
+    operator = build_lm_coefficient_operator(X)
+    # operator accepts [n, m] → [p, m], so transpose in and out
+    null_coef = operator(Y_perm.T).T  # (n_perm, p)
+    return np.asarray(null_coef)
+
+
 # =============================================================================
 # Public API
 # =============================================================================
@@ -233,22 +269,33 @@ def compute_params_permutation(
 
     n_coef = len(coef)
 
-    # Store null distribution of coefficients and t-statistics
-    null_coef = np.zeros((n_perm, n_coef))
-    null_t = np.zeros((n_perm, n_coef))
-
     # Original y for permutation
     y_orig = bundle.y.copy()
     n_obs = len(y_orig)
 
-    for i in range(n_perm):
-        perm_idx = rng.permutation(n_obs)
-        y_perm = y_orig[perm_idx]
+    # LM fast path: single matmul via hat-matrix operator
+    is_lm = (
+        spec.family == "gaussian"
+        and spec.link == "identity"
+        and not spec.has_random_effects
+    )
 
-        fit_perm = fit_with_permuted_y(spec, bundle, y_perm)
+    if is_lm:
+        null_coef = _compute_permutation_lm_fast(bundle.X, y_orig, n_perm, rng)
+        null_t = (null_coef - null) / se_obs
+    else:
+        # Generic path: refit model on each permutation
+        null_coef = np.zeros((n_perm, n_coef))
+        null_t = np.zeros((n_perm, n_coef))
+
+        for i in range(n_perm):
+            perm_idx = rng.permutation(n_obs)
+            y_perm = y_orig[perm_idx]
+
+            fit_perm = fit_with_permuted_y(spec, bundle, y_perm)
 
-        null_coef[i, :] = fit_perm.coef
-        null_t[i, :] = (fit_perm.coef - null) / se_obs
+            null_coef[i, :] = fit_perm.coef
+            null_t[i, :] = (fit_perm.coef - null) / se_obs
 
     # Permutation p-values based on alternative
     if alternative == "two-sided":

{bossanova-0.1.0.dev21 → bossanova-0.1.0.dev23}/bossanova/internal/infer/resample/__init__.py

@@ -26,6 +26,7 @@ from bossanova.internal.infer.resample.core import (
 from bossanova.internal.infer.resample.glmer import glmer_bootstrap
 from bossanova.internal.infer.resample.lm_operators import (
     build_lm_coefficient_operator,
+    compute_batched_ols_bootstrap,
 )
 from bossanova.internal.infer.resample.lmer import lmer_bootstrap
 from bossanova.internal.infer.resample.results import BootstrapResult
@@ -35,6 +36,7 @@ __all__ = [
     "BCa_MIN_NBOOT",
     "BootstrapResult",
     "build_lm_coefficient_operator",
+    "compute_batched_ols_bootstrap",
     "compute_bootstrap_ci_basic",
     "compute_bootstrap_ci_bca",
     "compute_bootstrap_ci_percentile",

bossanova-0.1.0.dev23/bossanova/internal/infer/resample/lm_operators.py

@@ -0,0 +1,280 @@
+"""Linear model coefficient operator (hat-matrix trick).
+
+Provides ``build_lm_coefficient_operator`` which pre-computes the expensive
+factorization of (X'X) once and returns a lightweight function that maps
+any response vector Y to OLS coefficients. Used by hypothesis property
+tests (BS.4) and the generic bootstrap engine.
+"""
+
+from typing import Any, Callable, Literal
+
+import numpy as np
+import scipy.linalg as la
+from bossanova.internal.maths.backend import get_backend
+
+# Type alias for arrays (works with both JAX and NumPy)
+Array = Any
+
+__all__ = [
+    "Array",
+    "build_lm_coefficient_operator",
+    "compute_batched_ols_bootstrap",
+]
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def _ensure_2d(xp: Any, Y: Array) -> tuple[Array, bool]:
+    """Convert Y to 2d, returning (Y_2d, was_1d)."""
+    is_1d = Y.ndim == 1
+    if is_1d:
+        Y = Y[:, xp.newaxis]
+    return Y, is_1d
+
+
+def _squeeze_if_1d(xp: Any, coef: Array, is_1d: bool) -> Array:
+    """Squeeze trailing dim if input was 1d."""
+    if is_1d:
+        return xp.squeeze(coef, axis=-1)
+    return coef
+
+
+def _make_cholesky_operator(
+    xp: Any, X: Array, cho_solve: Callable[..., Array]
+) -> Callable[[Array], Array]:
+    """Build a Cholesky-based coefficient operator.
+
+    Args:
+        xp: Array module (numpy or jax.numpy).
+        X: Design matrix, shape (n, p).
+        cho_solve: Cholesky solve function (scipy or jax.scipy).
+
+    Returns:
+        Function mapping Y -> coefficients.
+    """
+    XtX = X.T @ X
+    L = xp.linalg.cholesky(XtX)
+
+    def apply(Y: Array) -> Array:
+        Y = xp.asarray(Y)
+        Y, is_1d = _ensure_2d(xp, Y)
+        coef = cho_solve((L, True), X.T @ Y)
+        return _squeeze_if_1d(xp, coef, is_1d)
+
+    return apply
+
+
+def _make_qr_operator(
+    xp: Any,
+    X: Array,
+    qr_fn: Callable[..., tuple[Array, Array]],
+    solve_tri: Callable[..., Array],
+) -> Callable[[Array], Array]:
+    """Build a QR-based coefficient operator.
+
+    Args:
+        xp: Array module (numpy or jax.numpy).
+        X: Design matrix, shape (n, p).
+        qr_fn: QR decomposition function.
+        solve_tri: Triangular solve function.
+
+    Returns:
+        Function mapping Y -> coefficients.
+    """
+    Q, R = qr_fn(X)
+
+    def apply(Y: Array) -> Array:
+        Y = xp.asarray(Y)
+        Y, is_1d = _ensure_2d(xp, Y)
+        coef = solve_tri(R, Q.T @ Y, lower=False)
+        return _squeeze_if_1d(xp, coef, is_1d)
+
+    return apply
+
+
+# =============================================================================
+# Coefficient Operator (Hat-Matrix Trick)
+# =============================================================================
+
+
+def build_lm_coefficient_operator(
+    X: Array,
+    method: Literal["cholesky", "qr"] = "cholesky",
+) -> Callable[[Array], Array]:
+    """Create reusable linear operator Y -> coefficients.
+
+    Precomputes and caches the expensive factorization. Returns a function
+    that maps Y -> coefficients.
+
+    This implements the "hat-matrix trick" for efficient permutation testing:
+    For fixed design matrix X, coefficients are a linear map from Y -> beta_hat:
+    beta_hat = (X'X)^-1 X'y. We precompute the factorization of (X'X) once
+    and reuse it for all permuted/bootstrapped Y values.
+
+    Args:
+        X: Design matrix, shape (n, p).
+        method: Factorization method.
+            - 'cholesky': Uses Cholesky decomposition of X'X (faster, more stable).
+            - 'qr': Uses QR decomposition of X (handles rank deficiency better).
+
+    Returns:
+        Function that maps Y (shape [n] or [n, m]) to coefficients.
+
+    Examples:
+        ```python
+        X = np.array([[1, 2], [1, 3], [1, 4]])
+        operator = build_lm_coefficient_operator(X)
+        y = np.array([5, 8, 11])
+        coef = operator(y)
+        ```
+    """
+    backend = get_backend()
+
+    if backend == "jax":
+        import jax.numpy as jnp
+        import jax.scipy.linalg as jsp
+
+        X_arr = jnp.asarray(X)
+        xp = jnp
+
+        if method == "cholesky":
+            return _make_cholesky_operator(xp, X_arr, jsp.cho_solve)
+        elif method == "qr":
+            return _make_qr_operator(xp, X_arr, jnp.linalg.qr, jsp.solve_triangular)
+        else:
+            raise ValueError(f"method must be 'cholesky' or 'qr', got {method}")
+
+    else:
+        X_arr = np.asarray(X)
+        xp = np
+
+        if method == "cholesky":
+            return _make_cholesky_operator(xp, X_arr, la.cho_solve)
+        elif method == "qr":
+            return _make_qr_operator(
+                xp,
+                X_arr,
+                lambda x: la.qr(x, mode="economic"),
+                la.solve_triangular,
+            )
+        else:
+            raise ValueError(f"method must be 'cholesky' or 'qr', got {method}")
+
+
+# =============================================================================
+# JIT-Cached OLS from Indices (Bootstrap Fast Path)
+# =============================================================================
+
+_ols_from_indices_cache: dict[str, Any] = {}
+
+
+def _make_ols_from_indices_fn(ops: Any) -> Callable:
+    """Create JIT-compiled OLS solver that indexes into X and y.
+
+    The returned function takes ``(indices, X, y)`` and computes OLS
+    coefficients via Cholesky factorization of ``X[indices]``. This is
+    the inner kernel for ``compute_batched_ols_bootstrap`` — it gets
+    vmap'd over the bootstrap index array.
+
+    On JAX, Cholesky returns NaN (not an exception) for singular input,
+    so rank-deficient bootstrap samples naturally produce NaN coefficients.
+
+    Args:
+        ops: Backend ArrayOps (numpy or JAX).
+
+    Returns:
+        JIT-compiled function ``(indices, X, y) -> coef``.
+    """
+    xp = ops.np
+
+    def _core(indices: Array, X: Array, y: Array) -> Array:
+        X_b = X[indices]
+        y_b = y[indices]
+        XtX = X_b.T @ X_b
+        Xty = X_b.T @ y_b
+        L = xp.linalg.cholesky(XtX)
+        # Solve L @ z = Xty, then L.T @ coef = z
+        z = ops.solve_triangular(L, Xty, lower=True)
+        coef = ops.solve_triangular(L.T, z, lower=False)
+        return coef
+
+    return ops.jit(_core)
+
+
+def _get_ols_from_indices_fn() -> Callable:
+    """Get JIT-compiled OLS-from-indices function for current backend.
+
+    Returns:
+        Cached JIT-compiled function ``(indices, X, y) -> coef``.
+    """
+    from bossanova.internal.maths.backend import get_ops
+
+    backend = get_backend()
+    if backend not in _ols_from_indices_cache:
+        ops = get_ops()
+        _ols_from_indices_cache[backend] = _make_ols_from_indices_fn(ops)
+    return _ols_from_indices_cache[backend]
+
+
+def compute_batched_ols_bootstrap(
+    boot_indices: np.ndarray,
+    X: np.ndarray,
+    y: np.ndarray,
+    chunk_size: int | None = None,
+) -> np.ndarray:
+    """Compute bootstrap OLS coefficients using chunked vmap.
+
+    For LM models (gaussian family, identity link, no random effects),
+    each bootstrap replicate is an independent OLS solve. This function
+    batches those solves via ``ops.vmap`` — on JAX, this compiles to
+    a single XLA program; on NumPy, it degrades to a loop+stack
+    (equivalent to sequential but without Python overhead from
+    ``fit_model``).
+
+    Memory safety is ensured by chunking: ``compute_batch_size``
+    determines how many solves to batch at once based on available
+    memory.
+
+    Args:
+        boot_indices: Bootstrap index array, shape ``(n_boot, n)``.
+        X: Design matrix, shape ``(n, p)``.
+        y: Response vector, shape ``(n,)``.
+        chunk_size: Number of bootstrap samples per chunk. If ``None``,
+            computed automatically from available memory.
+
+    Returns:
+        Coefficient array, shape ``(n_boot, p)``.
+    """
+    from bossanova.internal.maths.backend import get_ops
+    from bossanova.internal.maths.batching import compute_batch_size
+
+    ops = get_ops()
+    n_boot = boot_indices.shape[0]
+    n, p = X.shape
+
+    # Compute chunk size if not provided
+    if chunk_size is None:
+        # Memory per bootstrap sample: X_b(n*p) + y_b(n) + XtX(p*p) + L(p*p) + Xty(p) + coef(p)
+        bytes_per_item = (n * p + n + 2 * p * p + 2 * p) * 8
+        chunk_size = compute_batch_size(n_items=n_boot, bytes_per_item=bytes_per_item)
+
+    # Get JIT-cached OLS function and vmap it
+    ols_fn = _get_ols_from_indices_fn()
+    vmapped_ols = ops.vmap(ols_fn, in_axes=(0, None, None))
+
+    # Convert to backend arrays
+    X_arr = ops.asarray(X)
+    y_arr = ops.asarray(y)
+
+    # Process in chunks
+    results = []
+    for start in range(0, n_boot, chunk_size):
+        end = min(start + chunk_size, n_boot)
+        idx_chunk = ops.asarray(boot_indices[start:end])
+        chunk_coefs = vmapped_ols(idx_chunk, X_arr, y_arr)
+        results.append(np.asarray(chunk_coefs))
+
+    return np.concatenate(results, axis=0)

{bossanova-0.1.0.dev21 → bossanova-0.1.0.dev23}/bossanova/internal/marginal/emm.py

@@ -272,20 +272,15 @@ def _compute_emm_gcomp(
         levels=levels,
     )
 
-    # Compute g-computation estimates and L_matrix for each level
-    n_levels = len(levels)
-    p = X.shape[1]
-    estimates = np.empty(n_levels, dtype=np.float64)
-    L_matrix = np.empty((n_levels, p), dtype=np.float64)
+    # Vectorized g-computation: batch all K counterfactual matmuls
+    X_all = np.stack(cf_matrices)  # (K, N, p)
+    eta_all = np.einsum("knp,p->kn", X_all, fit.coef)  # (K, N)
+    pred_all = np.asarray(apply_link_inverse(spec.link, eta_all))  # (K, N)
+    estimates = pred_all.mean(axis=1)  # (K,)
 
-    for k, X_cf in enumerate(cf_matrices):
-        eta_k = X_cf @ fit.coef  # (N,)
-        pred_k = np.asarray(apply_link_inverse(spec.link, eta_k))  # (N,)
-        estimates[k] = pred_k.mean()
-
-        # Delta method: L_k = mean(dμ/dη(η_k)[:, None] * X_cf_k, axis=0)
-        d_k = np.asarray(apply_link_inverse_deriv(spec.link, eta_k))  # (N,)
-        L_matrix[k] = (d_k[:, np.newaxis] * X_cf).mean(axis=0)
+    # Delta method: L_k = mean(dμ/dη(η_k)[:, None] * X_cf_k, axis=0)
+    d_all = np.asarray(apply_link_inverse_deriv(spec.link, eta_all))  # (K, N)
+    L_matrix = np.einsum("knp,kn->kp", X_all, d_all) / X_all.shape[1]  # (K, p)
 
     # Build grid DataFrame
     grid = _build_emm_grid(focal_var, levels, at_overrides, set_categoricals)
@@ -404,32 +399,26 @@ def compute_conditional_emm(
 
     assert isinstance(varying_offsets, VaryingState)
     n_groups = varying_offsets.n_groups
-    n_levels = len(levels)
 
-    # Compute per-group intercept BLUPs
+    # Compute per-group BLUP offsets: vectorized over all groups at once
     b_intercept = varying_offsets.effects.get("Intercept", np.zeros(n_groups))
-
-    # Build estimates: one row per (level, group) combination
-    # Shape: (n_levels * n_groups,)
-    estimates_link = np.empty(n_levels * n_groups)
-    L_matrix = np.empty((n_levels * n_groups, X_ref.shape[1]))
-
-    for g in range(n_groups):
-        # Total BLUP offset for this group
-        blup_offset = b_intercept[g]
-        for effect_name, blups in varying_offsets.effects.items():
-            if effect_name == "Intercept":
-                continue
-            # Add slope BLUP contributions (evaluated at X_means)
-            X_names_list = list(bundle.X_names)
-            if effect_name in X_names_list:
-                cov_idx = X_names_list.index(effect_name)
-                blup_offset += blups[g] * X_means[cov_idx]
-
-        for lev in range(n_levels):
-            row_idx = lev * n_groups + g
-            estimates_link[row_idx] = X_ref[lev] @ fit.coef + blup_offset
-            L_matrix[row_idx] = X_ref[lev]
+    blup_offsets = b_intercept.copy()  # (G,)
+    X_names_list = list(bundle.X_names)
+    for effect_name, blups in varying_offsets.effects.items():
+        if effect_name == "Intercept":
+            continue
+        if effect_name in X_names_list:
+            cov_idx = X_names_list.index(effect_name)
+            blup_offsets += blups * X_means[cov_idx]  # (G,) broadcast
+
+    # Vectorized per-level, per-group linear predictor
+    base_link = X_ref @ fit.coef  # (K,)
+    # estimates_link[lev, g] = base_link[lev] + blup_offsets[g]
+    estimates_link_2d = base_link[:, np.newaxis] + blup_offsets[np.newaxis, :]  # (K, G)
+    estimates_link = estimates_link_2d.ravel()  # matches lev*n_groups+g indexing
+
+    # L_matrix: each X_ref[lev] repeated n_groups times
+    L_matrix = np.repeat(X_ref, n_groups, axis=0)  # (K*G, p)
 
     estimates = estimates_link
     link_name: str | None = None
@@ -664,8 +653,24 @@ def _compute_emm_crossed_mem(
         L_matrix = d[:, np.newaxis] * X_ref
 
     grid = pl.DataFrame(grid_data)
-    if grid_vars:
-        grid = grid.select(grid_vars + [c for c in grid.columns if c not in grid_vars])
+
+    # Add scalar-pinned conditions as constant columns (set_categoricals
+    # and at_overrides are applied to the design matrix but must also
+    # appear in the output grid so users can see the pinned values).
+    if resolved.set_categoricals:
+        for cond_var, cond_level in resolved.set_categoricals.items():
+            if cond_var not in grid.columns:
+                grid = grid.with_columns(pl.lit(cond_level).alias(cond_var))
+    if resolved.at_overrides:
+        for cond_var, cond_val in resolved.at_overrides.items():
+            if cond_var not in grid.columns:
+                grid = grid.with_columns(pl.lit(cond_val).alias(cond_var))
+
+    # Reorder: grid_vars + pinned conditions first, then focal
+    pin_cols = list(resolved.set_categoricals) + list(resolved.at_overrides)
+    leading = grid_vars + [c for c in pin_cols if c not in grid_vars]
+    if leading:
+        grid = grid.select(leading + [c for c in grid.columns if c not in leading])
 
     cond_str = " ~ " + ", ".join(grid_vars) if grid_vars else ""
     return build_mee_state(
@@ -720,7 +725,6 @@ def _compute_emm_crossed_gcomp(
 
     X = bundle.X.copy()
     X_names_list = list(bundle.X_names)
-    p = X.shape[1]
 
     # Apply scalar at_overrides to X
     for var_name, value in resolved.at_overrides.items():
@@ -746,50 +750,57 @@ def _compute_emm_crossed_gcomp(
         levels=focal_levels,
     )
 
-    n_total = len(focal_levels) * len(grid_combos)
-    estimates = np.empty(n_total, dtype=np.float64)
-    L_matrix = np.empty((n_total, p), dtype=np.float64)
-
+    # Pre-build all C*K patched counterfactual matrices, then batch matmul
+    all_matrices: list[np.ndarray] = []
     grid_data: dict[str, list[object]] = {focal_var: []}
     for gv in grid_vars:
         grid_data[gv] = []
 
-    row_idx = 0
     for combo in grid_combos:
-        # Apply combo-specific overrides on top of the counterfactual matrices
         for k, X_cf_k in enumerate(cf_matrices):
             X_combo = X_cf_k.copy()
 
-            # Apply grid combo conditions
+            # Apply grid combo conditions (O(p) column patching, not O(N*p))
             for i, gv in enumerate(grid_vars):
                 val = combo[i]
                 if isinstance(val, str):
-                    # Categorical: set dummy columns for this condition var
                     for j, name in enumerate(bundle.X_names):
                         info = parse_design_column_name(name)
                         if info.column_type == "categorical" and info.base_term == gv:
                             X_combo[:, j] = 1.0 if info.level == val else 0.0
                 elif gv in X_names_list:
-                    # Numeric: override column
                     X_combo[:, X_names_list.index(gv)] = float(val)
 
-            # G-computation: mean response-scale prediction
-            eta_k = X_combo @ fit.coef
-            pred_k = np.asarray(apply_link_inverse(spec.link, eta_k))
-            estimates[row_idx] = pred_k.mean()
-
-            # Delta method: L_row = mean(dmu/deta[:, None] * X_combo, axis=0)
-            d_k = np.asarray(apply_link_inverse_deriv(spec.link, eta_k))
-            L_matrix[row_idx] = (d_k[:, np.newaxis] * X_combo).mean(axis=0)
-
+            all_matrices.append(X_combo)
             grid_data[focal_var].append(focal_levels[k])
             for i, gv in enumerate(grid_vars):
                 grid_data[gv].append(combo[i])
-            row_idx += 1
+
+    # Batched matmul + link computation over all C*K matrices
+    X_all = np.stack(all_matrices)  # (C*K, N, p)
+    eta_all = np.einsum("mnp,p->mn", X_all, fit.coef)  # (C*K, N)
+    pred_all = np.asarray(apply_link_inverse(spec.link, eta_all))  # (C*K, N)
+    estimates = pred_all.mean(axis=1)  # (C*K,)
+
+    d_all = np.asarray(apply_link_inverse_deriv(spec.link, eta_all))  # (C*K, N)
+    L_matrix = np.einsum("mnp,mn->mp", X_all, d_all) / X_all.shape[1]  # (C*K, p)
 
     grid = pl.DataFrame(grid_data)
-    if grid_vars:
-        grid = grid.select(grid_vars + [c for c in grid.columns if c not in grid_vars])
+
+    # Add scalar-pinned conditions as constant columns (mirrors MEM path)
+    if resolved.set_categoricals:
+        for cond_var, cond_level in resolved.set_categoricals.items():
+            if cond_var not in grid.columns:
+                grid = grid.with_columns(pl.lit(cond_level).alias(cond_var))
+    if resolved.at_overrides:
+        for cond_var, cond_val in resolved.at_overrides.items():
+            if cond_var not in grid.columns:
+                grid = grid.with_columns(pl.lit(cond_val).alias(cond_var))
+
+    pin_cols = list(resolved.set_categoricals) + list(resolved.at_overrides)
+    leading = grid_vars + [c for c in pin_cols if c not in grid_vars]
+    if leading:
+        grid = grid.select(leading + [c for c in grid.columns if c not in leading])
 
     cond_str = " ~ " + ", ".join(grid_vars) if grid_vars else ""
     return build_mee_state(