bossanova 0.1.0.dev11__tar.gz → 0.1.0.dev12__tar.gz

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/.gitignore

@@ -17,7 +17,7 @@ docs/_build/
 docs/performance/
 profile_*.py
 benchmarks/
-!bossanova-docs/reference/benchmarks/
+!bossanova-docs/engineering/benchmarks/
 examples/
 papers/
 paper-summaries/
@@ -64,6 +64,7 @@ dist/
 
 # Logs
 *.log
+logs/
 
 # Refs
 MixedModels.jl/

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/PKG-INFO

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

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/bossanova/__init__.py

@@ -1,10 +1,11 @@
 # Backend API (must be imported first, before any JAX usage)
 from bossanova.internal.maths.backend import backend, get_backend, set_backend
 
-# Configure JAX x64 eagerly if JAX is available
+# Configure JAX x64 eagerly if JAX is available — SINGLE AUTHORITATIVE LOCATION.
 # This MUST happen before any JAX arrays are created anywhere in the codebase.
 # We cannot defer this to lazy loading because submodules (e.g., lmer_core.py)
 # import JAX at module level and would create float32 arrays otherwise.
+# All other modules rely on this having run at import time.
 try:
     import jax
 

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/bossanova/internal/containers/__init__.py

@@ -67,8 +67,11 @@ from bossanova.internal.containers.builders import (
     build_params_dataframe,
     build_prediction_state,
     build_predictions_dataframe,
+    build_mee_resamples,
+    build_params_resamples,
     build_resamples_dataframe,
     build_resamples_state,
+    extract_mee_names,
     build_simulation_inference_state,
     build_simulation_spec,
     build_simulation_spec_from_formula,
@@ -117,10 +120,12 @@ __all__ = [
     "build_inference_state",
     "build_joint_test_dataframe",
     "build_joint_test_state",
+    "build_mee_resamples",
     "build_mee_state",
     "build_model_spec",
     "build_model_spec_from_formula",
     "build_params_dataframe",
+    "build_params_resamples",
     "build_prediction_state",
     "build_predictions_dataframe",
     "build_resamples_dataframe",
@@ -136,5 +141,6 @@ __all__ = [
     "build_varying_spread_dataframe",
     "build_varying_spread_state",
     "build_varying_state",
+    "extract_mee_names",
     "get_varying_random_terms",
 ]

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/bossanova/internal/containers/builders/__init__.py

@@ -26,6 +26,13 @@ from bossanova.internal.containers.builders.specs import (
     build_varying_spec,
     get_varying_random_terms,
 )
+from bossanova.internal.containers.builders.resamples import (
+    build_mee_resamples,
+    build_params_resamples,
+    build_resamples_dataframe,
+    build_resamples_state,
+    extract_mee_names,
+)
 from bossanova.internal.containers.builders.state import (
     build_cv_state,
     build_fit_state,
@@ -33,8 +40,6 @@ from bossanova.internal.containers.builders.state import (
     build_joint_test_state,
     build_mee_state,
     build_prediction_state,
-    build_resamples_dataframe,
-    build_resamples_state,
     build_simulation_inference_state,
     build_varying_spread_state,
     build_varying_state,
@@ -48,10 +53,12 @@ __all__ = [
     "build_inference_state",
     "build_joint_test_dataframe",
     "build_joint_test_state",
+    "build_mee_resamples",
     "build_mee_state",
     "build_model_spec",
     "build_model_spec_from_formula",
     "build_params_dataframe",
+    "build_params_resamples",
     "build_prediction_state",
     "build_predictions_dataframe",
     "build_resamples_dataframe",
@@ -67,5 +74,6 @@ __all__ = [
     "build_varying_spread_dataframe",
     "build_varying_spread_state",
     "build_varying_state",
+    "extract_mee_names",
     "get_varying_random_terms",
 ]

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/bossanova/internal/containers/builders/dataframes.py

@@ -19,7 +19,9 @@ from bossanova.internal.containers.schemas import (
     Col,
     ParamsAsymp,
     ParamsBase,
+    ParamsBoot,
     ParamsCv,
+    ParamsPerm,
     VaryingCorrSchema,
 )
 
@@ -56,33 +58,61 @@ def build_params_dataframe(
 ) -> pl.DataFrame:
     """Build the ``.params`` DataFrame from fit state.
 
+    Column set varies by inference method:
+
+    - **asymp**: all inference columns (p_value last).
+    - **boot**: no p_value; df = n_resamples.
+    - **perm**: no CIs; df = n_valid_resamples.
+    - **cv**: PRE columns only.
+    - **None**: term + estimate only.
+
     Args:
         bundle: Data bundle containing ``X_names`` (coefficient labels).
         fit: Fit state containing ``coef`` (coefficient estimates).
         params_inference: Optional inference state with SE, CI, p-values.
 
     Returns:
-        DataFrame with ``term``, ``estimate``, and optional inference columns.
+        DataFrame with ``term``, ``estimate``, and method-appropriate
+        inference columns.
     """
     terms = list(bundle.X_names)
     estimates = fit.coef.tolist()
 
     data: dict[str, list] = {Col.TERM: terms, Col.ESTIMATE: estimates}
 
-    if params_inference is not None:
-        if params_inference.method == "cv":
-            data[Col.PRE] = params_inference.pre.tolist()
-            data[Col.PRE_SD] = params_inference.pre_sd.tolist()
-            return pl.DataFrame(data, schema=ParamsCv)
+    if params_inference is None:
+        return pl.DataFrame(data, schema=ParamsBase)
+
+    method = params_inference.method
+
+    if method == "cv":
+        data[Col.PRE] = params_inference.pre.tolist()
+        data[Col.PRE_SD] = params_inference.pre_sd.tolist()
+        return pl.DataFrame(data, schema=ParamsCv)
+
+    if method == "boot":
+        data[Col.SE] = params_inference.se.tolist()
+        data[Col.CI_LOWER] = params_inference.ci_lower.tolist()
+        data[Col.CI_UPPER] = params_inference.ci_upper.tolist()
+        data[Col.STATISTIC] = params_inference.statistic.tolist()
+        data[Col.DF] = params_inference.df.tolist()
+        return pl.DataFrame(data, schema=ParamsBoot)
+
+    if method == "perm":
         data[Col.SE] = params_inference.se.tolist()
         data[Col.STATISTIC] = params_inference.statistic.tolist()
         data[Col.DF] = params_inference.df.tolist()
         data[Col.P_VALUE] = params_inference.p_value.tolist()
-        data[Col.CI_LOWER] = params_inference.ci_lower.tolist()
-        data[Col.CI_UPPER] = params_inference.ci_upper.tolist()
-        return pl.DataFrame(data, schema=ParamsAsymp)
+        return pl.DataFrame(data, schema=ParamsPerm)
 
-    return pl.DataFrame(data, schema=ParamsBase)
+    # Default: asymptotic (or any unknown method)
+    data[Col.SE] = params_inference.se.tolist()
+    data[Col.CI_LOWER] = params_inference.ci_lower.tolist()
+    data[Col.CI_UPPER] = params_inference.ci_upper.tolist()
+    data[Col.STATISTIC] = params_inference.statistic.tolist()
+    data[Col.DF] = params_inference.df.tolist()
+    data[Col.P_VALUE] = params_inference.p_value.tolist()
+    return pl.DataFrame(data, schema=ParamsAsymp)
 
 
 def build_varying_offsets_dataframe(
@@ -220,6 +250,7 @@ def build_varying_corr_dataframe(
         groups.append(group)
         corrs.append(value)
 
+    # TODO: are we only limited to 2 effects?
     return pl.DataFrame(
         {
             Col.GROUP: groups,
@@ -246,19 +277,27 @@ def _infer_group_from_spread(varying_spread: VaryingSpreadState) -> str:
     return ""
 
 
-def build_effects_dataframe(mee: MeeState) -> pl.DataFrame:
+def build_effects_dataframe(
+    mee: MeeState,
+    method: str | None = None,
+) -> pl.DataFrame:
     """Build the ``.effects`` DataFrame from marginal effects state.
 
+    Column set varies by inference method: bootstrap excludes ``p_value``,
+    permutation excludes ``ci_lower``/``ci_upper``.
+
     Args:
         mee: MeeState with grid, estimates, and optional inference.
+        method: Inference method (``"asymp"``, ``"boot"``, ``"perm"``,
+            or ``None``). Controls which inference columns are included.
 
     Returns:
-        DataFrame with grid columns, ``estimate``, and optional inference
-        columns.
+        DataFrame with grid columns, ``estimate``, and method-appropriate
+        inference columns.
     """
     result = mee.grid.clone()
     result = result.with_columns(pl.Series(Col.ESTIMATE, mee.estimate))
-    return append_inference_columns(result, mee)
+    return append_inference_columns(result, mee, method=method)
 
 
 def build_joint_test_dataframe(state: JointTestState) -> pl.DataFrame:

bossanova-0.1.0.dev12/bossanova/internal/containers/builders/resamples.py

@@ -0,0 +1,226 @@
+"""Builder functions for resamples-related containers.
+
+Provides constructors for ResamplesState and helpers for building
+resamples from inference results. Moved from state.py to keep modules
+under the 800-line limit.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from bossanova.internal.containers.schemas import (
+    Col,
+    ResamplesRawSchema,
+)
+from bossanova.internal.containers.structs.state import (
+    ResamplesState,
+)
+
+if TYPE_CHECKING:
+    import polars as pl
+    from numpy.typing import NDArray
+
+    from bossanova.internal.containers.structs.state import (
+        InferenceState,
+        MeeState,
+    )
+
+__all__ = [
+    "build_mee_resamples",
+    "build_params_resamples",
+    "build_resamples_dataframe",
+    "build_resamples_state",
+    "extract_mee_names",
+]
+
+
+def build_resamples_state(
+    *,
+    samples: NDArray[np.floating],
+    observed: NDArray[np.floating],
+    names: tuple[str, ...] | list[str],
+    method: str,
+    n_resamples: int,
+    context: str,
+) -> ResamplesState:
+    """Build a ResamplesState from resampling results.
+
+    Args:
+        samples: Resampled statistics array, shape (n_resamples, k).
+        observed: Observed statistics, shape (k,).
+        names: Term/effect names corresponding to columns of samples.
+        method: Resampling method (``"boot"`` or ``"perm"``).
+        n_resamples: Number of resamples.
+        context: What was resampled (``"params"`` or ``"effects"``).
+
+    Returns:
+        Frozen ResamplesState instance.
+
+    Examples:
+        >>> state = build_resamples_state(
+        ...     samples=np.random.randn(100, 2),
+        ...     observed=np.array([1.0, 2.0]),
+        ...     names=("Intercept", "x"),
+        ...     method="boot",
+        ...     n_resamples=100,
+        ...     context="params",
+        ... )
+        >>> state.method
+        'boot'
+    """
+    if isinstance(names, list):
+        names = tuple(names)
+
+    return ResamplesState(
+        samples=samples,
+        observed=observed,
+        names=names,
+        method=method,
+        n_resamples=n_resamples,
+        context=context,
+    )
+
+
+def build_resamples_dataframe(rs: ResamplesState) -> pl.DataFrame:
+    """Build a long-format DataFrame of raw resampled values.
+
+    Returns one row per (resample, term) combination with the raw
+    resampled value — coefficient estimates for bootstrap, null
+    t-statistics for permutation.
+
+    Columns: ``resample`` (int), ``term`` (str), ``value`` (float).
+
+    Args:
+        rs: Frozen ResamplesState from bootstrap or permutation inference.
+
+    Returns:
+        Polars DataFrame with ``n_resamples × k`` rows, where *k* is the
+        number of terms/effects.
+
+    Examples:
+        >>> df = build_resamples_dataframe(rs)
+        >>> df.columns
+        ['resample', 'term', 'value']
+        >>> df.shape
+        (1000, 3)  # 100 resamples × 10 terms
+    """
+    import polars as pl
+
+    samples = np.asarray(rs.samples)
+    n_resamples, k = samples.shape
+    names = list(rs.names)
+
+    # Repeat resample indices for each term, repeat term names for each resample
+    resample_idx = np.repeat(np.arange(n_resamples), k)
+    term_names = names * n_resamples
+    values = samples.ravel()
+
+    return pl.DataFrame(
+        {
+            Col.RESAMPLE: resample_idx,
+            Col.TERM: term_names,
+            Col.VALUE: values.tolist(),
+        },
+        schema=ResamplesRawSchema,
+    )
+
+
+def extract_mee_names(mee: MeeState) -> tuple[str, ...]:
+    """Extract human-readable names from a MeeState.
+
+    Args:
+        mee: The MeeState to extract names from.
+
+    Returns:
+        Tuple of effect names for each estimate.
+    """
+    grid = mee.grid
+    focal = mee.focal_var
+
+    if mee.type == "contrasts" and "contrast" in grid.columns:
+        return tuple(grid["contrast"].to_list())
+
+    if focal in grid.columns:
+        return tuple(str(v) for v in grid[focal].to_list())
+
+    # Fallback: numbered estimates
+    return tuple(f"estimate_{i}" for i in range(len(mee.estimate)))
+
+
+def build_params_resamples(
+    inference: InferenceState | None,
+    fit_coef: np.ndarray,
+    x_names: tuple[str, ...],
+    how: str,
+) -> ResamplesState | None:
+    """Build ResamplesState from params inference if samples are available.
+
+    Args:
+        inference: The InferenceState from params inference, or None.
+        fit_coef: Coefficient estimates from the FitState.
+        x_names: Design matrix column names from the DataBundle.
+        how: Inference method used (``"boot"``, ``"perm"``, etc.).
+
+    Returns:
+        ResamplesState if boot/perm samples were saved, else None.
+    """
+    if inference is None:
+        return None
+
+    if how == "boot" and inference.boot_samples is not None:
+        return build_resamples_state(
+            samples=inference.boot_samples,
+            observed=fit_coef,
+            names=x_names,
+            method="boot",
+            n_resamples=inference.n_resamples,
+            context="params",
+        )
+
+    if how == "perm" and inference.perm_samples is not None:
+        return build_resamples_state(
+            samples=inference.perm_samples,
+            observed=inference.statistic,
+            names=x_names,
+            method="perm",
+            n_resamples=inference.n_resamples,
+            context="params",
+        )
+
+    return None
+
+
+def build_mee_resamples(
+    mee: MeeState | None,
+    samples: np.ndarray | None,
+    how: str,
+) -> ResamplesState | None:
+    """Build ResamplesState from MEE inference if samples are available.
+
+    Args:
+        mee: The MeeState from explore, or None.
+        samples: Raw resample array from dispatch_mee_inference, or None.
+        how: Inference method used (``"boot"``, ``"perm"``, etc.).
+
+    Returns:
+        ResamplesState if boot/perm samples were saved, else None.
+    """
+    if samples is None or mee is None:
+        return None
+
+    names = extract_mee_names(mee)
+
+    method = "boot" if how == "boot" else "perm"
+    observed = mee.estimate if how == "boot" else mee.statistic
+
+    return build_resamples_state(
+        samples=samples,
+        observed=observed,
+        names=names,
+        method=method,
+        n_resamples=samples.shape[0],
+        context="effects",
+    )

{bossanova-0.1.0.dev11 → bossanova-0.1.0.dev12}/bossanova/internal/containers/builders/results.py

@@ -17,29 +17,43 @@ __all__ = [
 ]
 
 # Standard inference column names, in canonical display order.
+# Estimation group (se, CIs) first, then testing group (statistic, df, p_value).
 _INFERENCE_COLS: tuple[str, ...] = (
     Col.SE,
+    Col.CI_LOWER,
+    Col.CI_UPPER,
     Col.STATISTIC,
     Col.DF,
     Col.P_VALUE,
-    Col.CI_LOWER,
-    Col.CI_UPPER,
 )
 
+# Columns excluded per inference method.
+_BOOT_EXCLUDE: frozenset[str] = frozenset({Col.P_VALUE})
+_PERM_EXCLUDE: frozenset[str] = frozenset({Col.CI_LOWER, Col.CI_UPPER})
 
-def append_inference_columns(df: pl.DataFrame, state: object) -> pl.DataFrame:
+
+def append_inference_columns(
+    df: pl.DataFrame,
+    state: object,
+    method: str | None = None,
+) -> pl.DataFrame:
     """Append standard inference columns to a DataFrame if available.
 
     Checks ``state.has_inference`` and, when True, adds each inference
-    column (se, statistic, df, p_value, ci_lower, ci_upper) that is not
-    None on *state*. Columns are added in canonical order. Attributes
-    that do not exist on *state* are silently skipped.
+    column that is not None on *state*. Columns are added in canonical
+    order: se, ci_lower, ci_upper, statistic, df, p_value.
+
+    When *method* is ``"boot"``, the ``p_value`` column is excluded.
+    When *method* is ``"perm"``, the ``ci_lower`` and ``ci_upper``
+    columns are excluded.
 
     Args:
         df: Base DataFrame to augment (not mutated).
         state: Object with a ``has_inference`` bool and optional ``se``,
             ``statistic``, ``df``, ``p_value``, ``ci_lower``,
             ``ci_upper`` array attributes.
+        method: Inference method (``"asymp"``, ``"boot"``, ``"perm"``,
+            or ``None``). Controls which columns are included.
 
     Returns:
         A new DataFrame with inference columns appended (or the
@@ -48,7 +62,16 @@ def append_inference_columns(df: pl.DataFrame, state: object) -> pl.DataFrame:
     if not getattr(state, "has_inference", False):
         return df
 
+    if method == "boot":
+        exclude = _BOOT_EXCLUDE
+    elif method == "perm":
+        exclude = _PERM_EXCLUDE
+    else:
+        exclude = frozenset()
+
     for col_name in _INFERENCE_COLS:
+        if col_name in exclude:
+            continue
         values = getattr(state, col_name, None)
         if values is not None:
             df = df.with_columns(pl.Series(col_name, values))