bossanova 0.1.0.dev25__tar.gz → 0.1.0.dev27__tar.gz

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/PKG-INFO

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

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/containers/builders/dataframes.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import polars as pl
 
 from bossanova.internal.containers.builders.results import (
@@ -183,35 +184,140 @@ def build_varying_params_dataframe(
 def build_varying_spread_dataframe(
     varying_spread: VaryingSpreadState,
 ) -> pl.DataFrame:
-    """Build the ``.varying_spread`` DataFrame from variance components.
+    """Build long-form ``.varying_spread`` DataFrame from variance components.
+
+    Produces a tidy DataFrame with one row per (term, statistic) pair.
+    Variance terms (sigma2, tau2) get both ``variance`` and ``std_dev``
+    rows. Correlations get a ``corr`` row. ICC gets an ``icc`` row.
 
     Args:
         varying_spread: Variance component state with components DataFrame
             and optional CI information.
 
     Returns:
-        DataFrame with ``component``, ``estimate``, and optional
-        ``ci_lower``, ``ci_upper``, ``ci_method`` columns.
+        DataFrame with columns ``term``, ``statistic``, ``estimate``,
+        and optional ``ci_lower``, ``ci_upper``, ``ci_method``.
     """
-    result = varying_spread.components.clone()
+    has_ci = varying_spread.has_inference
+    terms: list[str] = []
+    statistics: list[str] = []
+    estimates: list[float] = []
+    ci_lowers: list[float] = []
+    ci_uppers: list[float] = []
+    ci_methods: list[str] = []
+
+    # Process components with sigma2 deferred to the end so Residual rows
+    # appear last, after all random effect terms.
+    comp_names = varying_spread.components[Col.COMPONENT].to_list()
+    deferred_sigma2: float | None = None
+
+    for comp_name in comp_names:
+        estimate = varying_spread.components.filter(pl.col(Col.COMPONENT) == comp_name)[
+            Col.ESTIMATE
+        ].item()
+
+        if comp_name == Col.SIGMA2:
+            deferred_sigma2 = estimate
+        elif comp_name.startswith(Col.TAU2_PREFIX):
+            # "tau2_Subject:Intercept" → "Subject | Intercept"
+            term = comp_name[len(Col.TAU2_PREFIX) :].replace(":", " | ")
+            _add_variance_rows(
+                term,
+                estimate,
+                comp_name,
+                varying_spread,
+                has_ci,
+                terms,
+                statistics,
+                estimates,
+                ci_lowers,
+                ci_uppers,
+                ci_methods,
+            )
+        elif comp_name.startswith(Col.RHO_PREFIX):
+            # "rho_Intercept:Days" → "Intercept:Days"
+            term = comp_name[len(Col.RHO_PREFIX) :]
+            terms.append(term)
+            statistics.append("corr")
+            estimates.append(estimate)
+            if has_ci:
+                ci_lowers.append(varying_spread.ci_lower.get(comp_name, float("nan")))
+                ci_uppers.append(varying_spread.ci_upper.get(comp_name, float("nan")))
+                ci_methods.append(varying_spread.ci_method or "")
+        elif comp_name == Col.ICC:
+            terms.append("ICC")
+            statistics.append("icc")
+            estimates.append(estimate)
+            if has_ci:
+                ci_lowers.append(varying_spread.ci_lower.get(comp_name, float("nan")))
+                ci_uppers.append(varying_spread.ci_upper.get(comp_name, float("nan")))
+                ci_methods.append(varying_spread.ci_method or "")
+
+    # Append Residual rows last
+    if deferred_sigma2 is not None:
+        _add_variance_rows(
+            "Residual",
+            deferred_sigma2,
+            Col.SIGMA2,
+            varying_spread,
+            has_ci,
+            terms,
+            statistics,
+            estimates,
+            ci_lowers,
+            ci_uppers,
+            ci_methods,
+        )
 
-    if varying_spread.has_inference:
-        ci_lower_vals = []
-        ci_upper_vals = []
-        ci_method_vals = []
+    data: dict[str, list] = {
+        Col.TERM: terms,
+        Col.METRIC: statistics,
+        Col.ESTIMATE: estimates,
+    }
+    if has_ci:
+        data[Col.CI_LOWER] = ci_lowers
+        data[Col.CI_UPPER] = ci_uppers
+        data[Col.CI_METHOD] = ci_methods
 
-        for comp_name in result[Col.COMPONENT].to_list():
-            ci_lower_vals.append(varying_spread.ci_lower.get(comp_name, float("nan")))
-            ci_upper_vals.append(varying_spread.ci_upper.get(comp_name, float("nan")))
-            ci_method_vals.append(varying_spread.ci_method or "")
+    return pl.DataFrame(data)
 
-        result = result.with_columns(
-            pl.Series(Col.CI_LOWER, ci_lower_vals),
-            pl.Series(Col.CI_UPPER, ci_upper_vals),
-            pl.Series(Col.CI_METHOD, ci_method_vals),
-        )
 
-    return result
+def _add_variance_rows(
+    term: str,
+    variance: float,
+    comp_name: str,
+    varying_spread: VaryingSpreadState,
+    has_ci: bool,
+    terms: list[str],
+    statistics: list[str],
+    estimates: list[float],
+    ci_lowers: list[float],
+    ci_uppers: list[float],
+    ci_methods: list[str],
+) -> None:
+    """Append variance and std_dev rows for a single term."""
+    sd = np.sqrt(variance) if variance >= 0 else 0.0
+
+    # Variance row
+    terms.append(term)
+    statistics.append("variance")
+    estimates.append(variance)
+    if has_ci:
+        ci_lowers.append(varying_spread.ci_lower.get(comp_name, float("nan")))
+        ci_uppers.append(varying_spread.ci_upper.get(comp_name, float("nan")))
+        ci_methods.append(varying_spread.ci_method or "")
+
+    # Std dev row
+    terms.append(term)
+    statistics.append("std_dev")
+    estimates.append(sd)
+    if has_ci:
+        # Convert variance CIs to SD CIs
+        var_lo = varying_spread.ci_lower.get(comp_name, float("nan"))
+        var_hi = varying_spread.ci_upper.get(comp_name, float("nan"))
+        ci_lowers.append(np.sqrt(var_lo) if var_lo >= 0 else float("nan"))
+        ci_uppers.append(np.sqrt(var_hi) if var_hi >= 0 else float("nan"))
+        ci_methods.append(varying_spread.ci_method or "")
 
 
 def build_varying_corr_dataframe(

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/containers/builders/state.py

@@ -452,6 +452,7 @@ def build_prediction_state(
 
 def build_cv_state(
     k: int,
+    mse: float,
     rmse: float,
     mae: float,
     r_squared: float,
@@ -471,6 +472,7 @@ def build_cv_state(
 
     Args:
         k: Number of folds used.
+        mse: Mean squared error (out-of-sample).
         rmse: Root mean squared error.
         mae: Mean absolute error.
         r_squared: Coefficient of determination.
@@ -491,6 +493,7 @@ def build_cv_state(
     Examples:
         >>> state = build_cv_state(
         ...     k=10,
+        ...     mse=0.274,
         ...     rmse=0.523,
         ...     mae=0.412,
         ...     r_squared=0.891,
@@ -498,6 +501,7 @@ def build_cv_state(
     """
     return CVState(
         k=k,
+        mse=mse,
         rmse=rmse,
         mae=mae,
         r_squared=r_squared,

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/containers/schemas.py

@@ -98,6 +98,7 @@ class Col:
     NULL_DEVIANCE: str = "null_deviance"
     DEVIANCE: str = "deviance"
     DISPERSION: str = "dispersion"
+    MSE: str = "mse"
     PSEUDO_RSQUARED: str = "pseudo_rsquared"
     RSQUARED_MARGINAL: str = "rsquared_marginal"
     RSQUARED_CONDITIONAL: str = "rsquared_conditional"
@@ -112,6 +113,7 @@ class Col:
 
     # -- Varying (random effects) --
     COMPONENT: str = "component"
+    METRIC: str = "metric"
     GROUP: str = "group"
     LEVEL: str = "level"
     EFFECT1: str = "effect1"
@@ -177,6 +179,8 @@ class Col:
     N: str = "n"
 
     # -- CV diagnostics --
+    CV_MSE: str = "cv_mse"
+    CV_MSE_SD: str = "cv_mse_sd"
     CV_RMSE: str = "cv_rmse"
     CV_MAE: str = "cv_mae"
     CV_RSQUARED: str = "cv_rsquared"
@@ -408,6 +412,7 @@ DiagnosticsGaussian = Schema(
         Col.NULL_DEVIANCE: Float64,
         Col.DEVIANCE: Float64,
         Col.DISPERSION: Float64,
+        Col.MSE: Float64,
         Col.PSEUDO_RSQUARED: Float64,
         Col.AIC: Float64,
         Col.BIC: Float64,
@@ -423,6 +428,7 @@ DiagnosticsGlm = Schema(
         Col.NULL_DEVIANCE: Float64,
         Col.DEVIANCE: Float64,
         Col.DISPERSION: Float64,
+        Col.MSE: Float64,
         Col.PSEUDO_RSQUARED: Float64,
         Col.AIC: Float64,
         Col.BIC: Float64,
@@ -437,6 +443,7 @@ DiagnosticsLmer = Schema(
         Col.DF_RESID: Float64,
         Col.SIGMA: Float64,
         Col.DEVIANCE: Float64,
+        Col.MSE: Float64,
         Col.RSQUARED_MARGINAL: Float64,
         Col.RSQUARED_CONDITIONAL: Float64,
         Col.ICC: Float64,
@@ -447,10 +454,12 @@ DiagnosticsLmer = Schema(
 )
 
 # CV-augmented columns (added by .infer(how='cv') after .fit())
-DiagnosticsCvCols = (Col.CV_RMSE, Col.CV_MAE, Col.CV_RSQUARED, Col.CV_K)
+DiagnosticsCvCols = (Col.CV_MSE, Col.CV_RMSE, Col.CV_MAE, Col.CV_RSQUARED, Col.CV_K)
 
 # Prediction CV columns (added by .predict().infer(how='cv'))
 DiagnosticsPredCvCols = (
+    Col.CV_MSE,
+    Col.CV_MSE_SD,
     Col.CV_RMSE,
     Col.CV_RMSE_SD,
     Col.CV_MAE,

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/containers/structs/inference.py

@@ -98,6 +98,9 @@ class CVState:
     """
 
     k: int = field(validator=[validators.instance_of(int), validators.gt(0)])
+    mse: float = field(
+        validator=[validators.instance_of((int, float)), validators.ge(0)]
+    )
     rmse: float = field(
         validator=[validators.instance_of((int, float)), validators.ge(0)]
     )

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/fit/diagnostics.py

@@ -267,6 +267,8 @@ def _lm_diagnostics(
             fit.null_deviance if fit.null_deviance is not None else float(np.sum(y**2))
         )
 
+    mse = ss_res / n
+
     return {
         Col.RSQUARED: [rsquared],
         Col.RSQUARED_ADJ: [rsquared_adj],
@@ -276,6 +278,7 @@ def _lm_diagnostics(
         Col.NULL_DEVIANCE: [null_deviance],
         Col.DEVIANCE: [deviance],
         Col.DISPERSION: [sigma**2],
+        Col.MSE: [mse],
         Col.PSEUDO_RSQUARED: [
             1.0 - deviance / null_deviance if null_deviance > 0 else float("nan")
         ],
@@ -313,10 +316,17 @@ def _glm_diagnostics(
         else float("nan")
     )
 
+    import numpy as np
+
+    mse = (
+        float(np.mean(fit.residuals**2)) if fit.residuals is not None else float("nan")
+    )
+
     return {
         Col.NULL_DEVIANCE: [null_deviance],
         Col.DEVIANCE: [deviance],
         Col.DISPERSION: [dispersion],
+        Col.MSE: [mse],
         Col.PSEUDO_RSQUARED: [pseudo_rsquared],
     }
 
@@ -385,9 +395,14 @@ def _mixed_diagnostics(
         var_group_resid = var_random + var_resid
         icc = var_random / var_group_resid if var_group_resid > 0 else 0.0
 
+    mse = (
+        float(np.mean(fit.residuals**2)) if fit.residuals is not None else float("nan")
+    )
+
     return {
         Col.SIGMA: [sigma],
         Col.DEVIANCE: [deviance],
+        Col.MSE: [mse],
         Col.RSQUARED_MARGINAL: [rsquared_marginal],
         Col.RSQUARED_CONDITIONAL: [rsquared_conditional],
         Col.ICC: [icc],
@@ -403,12 +418,15 @@ def _add_cv_metrics(data: dict[str, list], cv: CVState) -> None:
     """
     import numpy as np
 
+    data[Col.CV_MSE] = [cv.mse]
     data[Col.CV_RMSE] = [cv.rmse]
     data[Col.CV_MAE] = [cv.mae]
     data[Col.CV_RSQUARED] = [cv.r_squared]
     data[Col.CV_K] = [cv.k]
 
     if cv.fold_metrics:
+        if "mse" in cv.fold_metrics:
+            data[Col.CV_MSE_SD] = [float(np.std(cv.fold_metrics["mse"]))]
         if "rmse" in cv.fold_metrics:
             data[Col.CV_RMSE_SD] = [float(np.std(cv.fold_metrics["rmse"]))]
         if "mae" in cv.fold_metrics:

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/fit/varying.py

@@ -253,10 +253,112 @@ def compute_varying_state(
     )
 
 
+def _get_predictor_column(
+    re_name: str,
+    X: NDArray[np.floating] | None,
+    X_names: tuple[str, ...] | None,
+) -> NDArray[np.floating] | None:
+    """Look up the data column for a random effect predictor.
+
+    Returns None for ``"Intercept"`` (caller uses all-ones logic) or when
+    the predictor cannot be found in the design matrix.
+    """
+    if re_name == "Intercept" or X is None or X_names is None:
+        return None
+    if re_name in X_names:
+        return X[:, X_names.index(re_name)]
+    return None
+
+
+def _compute_johnson_icc(
+    tau2: dict[str, float],
+    rho: dict[str, float],
+    sigma2: float,
+    re_meta: REInfo,
+    X: NDArray[np.floating] | None,
+    X_names: tuple[str, ...] | None,
+) -> float | None:
+    """Compute ICC using Johnson (2014, PLoS ONE, eq. 10).
+
+    For each grouping factor, computes the mean random effect variance
+    accounting for predictor distributions when random slopes are present:
+
+        contribution_g = Σ_i Σ_j Cov(b_i, b_j) · E[X_i · X_j]
+
+    For intercept-only factors this simplifies to τ00. For slope models
+    it incorporates E[X] and E[X²] from the data.
+
+    Falls back to sum(τ00) / (sum(τ00) + σ²) when data is unavailable.
+    """
+    group_names = list(re_meta.grouping_vars)
+    if not group_names:
+        return None
+
+    # Get per-factor RE names
+    re_structures_list = re_meta.metadata.get("re_structures_list", None)
+    if re_structures_list is not None and len(re_structures_list) == len(group_names):
+        _, names_per_group = per_factor_re_info(re_meta, group_names)
+    else:
+        # Single-factor model: all random names belong to the one group
+        names_per_group = {group_names[0]: re_meta.random_names or ["Intercept"]}
+
+    var_random = 0.0
+
+    for group_name in group_names:
+        factor_names = names_per_group[group_name]
+        k = len(factor_names)
+
+        # Build covariance matrix for this factor
+        cov = np.zeros((k, k))
+        for i, name_i in enumerate(factor_names):
+            key_i = f"{group_name}:{name_i}"
+            cov[i, i] = tau2.get(key_i, 0.0)
+
+        # Fill off-diagonal from rho
+        # NOTE: rho keys are "RE1:RE2" without group prefix. When two groups
+        # share the same slope predictor names, the last group's rho wins.
+        # This is a known limitation of the current rho dict structure.
+        if k > 1:
+            for i in range(k):
+                for j in range(i + 1, k):
+                    rho_key = f"{factor_names[i]}:{factor_names[j]}"
+                    if rho_key in rho:
+                        cov_ij = rho[rho_key] * np.sqrt(cov[i, i] * cov[j, j])
+                        cov[i, j] = cov_ij
+                        cov[j, i] = cov_ij
+
+        # Build E[X_i · X_j] matrix
+        e_xx = np.zeros((k, k))
+        cols: list[NDArray[np.floating] | None] = [
+            _get_predictor_column(name, X, X_names) for name in factor_names
+        ]
+        for i in range(k):
+            for j in range(k):
+                ci = cols[i]
+                cj = cols[j]
+                if ci is not None and cj is not None:
+                    e_xx[i, j] = np.mean(ci * cj)
+                elif ci is None and factor_names[i] == "Intercept":
+                    # Intercept is all 1s: E[1 · X_j]
+                    e_xx[i, j] = np.mean(cj) if cj is not None else 1.0
+                elif cj is None and factor_names[j] == "Intercept":
+                    # E[X_i · 1]
+                    e_xx[i, j] = np.mean(ci) if ci is not None else 1.0
+
+        var_random += float(np.sum(cov * e_xx))
+
+    if var_random <= 0:
+        return None
+    return var_random / (var_random + sigma2)
+
+
 def compute_varying_spread_state(
     theta: NDArray[np.floating],
     sigma: float,
     re_meta: REInfo,
+    *,
+    X: NDArray[np.floating] | None = None,
+    X_names: tuple[str, ...] | None = None,
 ) -> VaryingSpreadState:
     """Compute VaryingSpreadState (variance components) from theta parameters.
 
@@ -264,10 +366,17 @@ def compute_varying_spread_state(
     correlations (rho), and intraclass correlation (ICC) from the fitted
     theta vector using the random effects structure.
 
+    ICC is computed using Johnson (2014, PLoS ONE, eq. 10) which accounts
+    for predictor distributions when random slopes are present. When ``X``
+    and ``X_names`` are provided, the formula uses E[X] and E[X²] from
+    the data; otherwise falls back to summing intercept variances only.
+
     Args:
         theta: Variance component parameters from the fitted model.
         sigma: Residual standard deviation from the fitted model.
         re_meta: Random effects metadata (grouping vars, structure, etc.).
+        X: Fixed-effects design matrix (n × p), used for ICC computation.
+        X_names: Column names for ``X``, used to look up slope predictors.
 
     Returns:
         VaryingSpreadState container with components DataFrame and
@@ -313,12 +422,8 @@ def compute_varying_spread_state(
             tau2_key = name.removesuffix("_sd")
             tau2[tau2_key] = value**2
 
-    # Compute ICC (for intercept-only models)
-    # ICC = tau² / (tau² + sigma²)
-    icc: float | None = None
-    total_tau2 = sum(tau2.values())
-    if total_tau2 > 0:
-        icc = total_tau2 / (total_tau2 + sigma2)
+    # Compute ICC using Johnson (2014) formula
+    icc = _compute_johnson_icc(tau2, rho, sigma2, re_meta, X, X_names)
 
     # Build components DataFrame
     components_data: dict[str, list] = {Col.COMPONENT: [], Col.ESTIMATE: []}
@@ -403,6 +508,8 @@ def build_mixed_post_fit_state(
             theta=fit.theta,
             sigma=sigma,
             re_meta=bundle.re_metadata,
+            X=bundle.X,
+            X_names=bundle.X_names,
         )
 
     # Emit convergence warnings if any issues detected

{bossanova-0.1.0.dev25 → bossanova-0.1.0.dev27}/bossanova/internal/infer/cv.py

@@ -485,8 +485,10 @@ def _compute_fold_metrics(
         'sensitivity', 'specificity', 'f1', 'auc'.
     """
     residuals = y_test - y_pred
+    mse = float(np.mean(residuals**2))
     metrics: dict[str, float] = {
-        "rmse": float(np.sqrt(np.mean(residuals**2))),
+        "mse": mse,
+        "rmse": float(np.sqrt(mse)),
         "mae": float(np.mean(np.abs(residuals))),
     }
 
@@ -644,7 +646,8 @@ def compute_cv_metrics(
     resid_valid = oos_residuals[valid_mask]
 
     # Overall out-of-sample metrics
-    rmse = float(np.sqrt(np.mean(resid_valid**2)))
+    mse = float(np.mean(resid_valid**2))
+    rmse = float(np.sqrt(mse))
     mae = float(np.mean(np.abs(resid_valid)))
 
     # R² computed on out-of-sample predictions
@@ -666,6 +669,7 @@ def compute_cv_metrics(
 
     return build_cv_state(
         k=k_int,
+        mse=mse,
         rmse=rmse,
         mae=mae,
         r_squared=r_squared,