bossanova 0.1.0.dev22__tar.gz → 0.1.0.dev24__tar.gz

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/PKG-INFO

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

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/marginal/emm.py

@@ -653,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(
@@ -770,8 +786,21 @@ def _compute_emm_crossed_gcomp(
     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(

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/simulation/model_sim.py

@@ -96,13 +96,29 @@ def generate_data_from_spec(
         if isinstance(v, np.ndarray) and v.dtype.kind == "f"
     ]
 
+    # Build interaction columns from coef keys (e.g. "x1:x2" → x1 * x2)
     obs_n = sim_spec.n
-    if continuous_vars:
-        X = np.column_stack([np.ones(obs_n)] + [data_dict[v] for v in continuous_vars])
-    else:
-        X = np.ones((obs_n, 1))
-
-    coef_names = ["Intercept"] + continuous_vars
+    interaction_cols: dict[str, np.ndarray] = {}
+    for name in sim_spec.coef:
+        if ":" not in name:
+            continue
+        parts = name.split(":")
+        # All parts must be sampled continuous variables
+        if not all(p in data_dict and data_dict[p].dtype.kind == "f" for p in parts):
+            continue
+        col = np.ones(obs_n)
+        for p in parts:
+            col = col * data_dict[p]
+        interaction_cols[name] = col
+
+    # Build X: intercept + main effects + interactions
+    interaction_names = list(interaction_cols.keys())
+    columns = [np.ones(obs_n)]
+    columns.extend(data_dict[v] for v in continuous_vars)
+    columns.extend(interaction_cols[name] for name in interaction_names)
+    X = np.column_stack(columns) if columns else np.ones((obs_n, 1))
+
+    coef_names = ["Intercept"] + continuous_vars + interaction_names
     coef_values = np.array([sim_spec.coef.get(name, 0.0) for name in coef_names])
 
     # Linear predictor

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/viz/README.md

@@ -35,7 +35,6 @@ m.plot_resid()
 | `plot_resid` | Residual diagnostic grid (4 panels) |
 | `plot_predict` | Marginal predictions across predictor range |
 | `plot_explore` | Marginal effects/means visualization |
-| `plot_fit` | Composite diagnostic panel |
 | `plot_compare` | Multi-model coefficient comparison |
 | `plot_dag` | Causal DAG visualization |
 | `plot_lattice` | Model lattice (Hasse diagram) |

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/viz/__init__.py

@@ -6,7 +6,6 @@ Post-estimation plots (require .fit()):
     plot_params: Coefficient forest plot with CIs
     plot_resid: Residual diagnostic panels (QQ, fitted vs residuals)
     plot_predict: Prediction plots with confidence bands
-    plot_fit: Observed vs fitted scatter
     plot_explore: Marginal effects visualization
     plot_ranef: Random effects caterpillar plot (mixed models)
     plot_resamples: Bootstrap/permutation distribution plots
@@ -44,7 +43,6 @@ from bossanova.internal.viz.params import plot_params
 from bossanova.internal.viz.ranef import plot_ranef
 from bossanova.internal.viz.resid import plot_resid
 from bossanova.internal.viz.predict import plot_predict
-from bossanova.internal.viz.fit import plot_fit
 from bossanova.internal.viz.mem import plot_explore
 from bossanova.internal.viz.compare import plot_compare
 from bossanova.internal.viz.resamples import plot_resamples
@@ -75,7 +73,6 @@ __all__ = [
     "plot_dag",
     "plot_design",
     "plot_explore",
-    "plot_fit",
     "plot_lattice",
     "plot_params",
     "plot_predict",

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/viz/helpers.py

@@ -18,7 +18,6 @@ if TYPE_CHECKING:
     from matplotlib.axes import Axes
 
 __all__ = [
-    "detect_blup_mode",
     "detect_ci_columns",
     "fill_ci_band",
     "get_grouping_factors",
@@ -59,8 +58,7 @@ def fill_ci_band(
 ) -> None:
     """Render sorted CI bands on a matplotlib axes.
 
-    This is the shared rendering core used by both ``fit_layers.add_ci_band``
-    and ``predict.add_ci_band``. It sorts data along x, then calls
+    Shared rendering core for CI bands. Sorts data along x, then calls
     ``ax.fill_between()`` once per hue level (or once if no hue).
 
     Args:
@@ -170,43 +168,3 @@ def get_grouping_factors(model: "Any") -> list[str]:
         return groups
 
     return []
-
-
-def detect_blup_mode(
-    model: "Any",
-    hue: str | None,
-    col: str | None,
-    row: str | None,
-    blups: bool | None,
-) -> bool:
-    """Determine if BLUP mode should be enabled for a plot.
-
-    Returns True if:
-    - ``blups=True`` explicitly, OR
-    - ``blups=None`` AND model is mixed AND any of ``hue``/``col``/``row``
-      is a grouping factor (auto-detection).
-
-    Args:
-        model: A fitted bossanova model.
-        hue: Hue variable name (or None).
-        col: Column facet variable name (or None).
-        row: Row facet variable name (or None).
-        blups: Explicit BLUP mode setting. None triggers auto-detection.
-
-    Returns:
-        True if BLUP mode should be enabled.
-    """
-    if blups is not None:
-        return blups
-
-    # Only mixed models have grouping factors
-    grouping_factors = get_grouping_factors(model)
-    if not grouping_factors:
-        return False
-
-    # Check if any faceting/hue variable is a grouping factor
-    for var in (hue, col, row):
-        if var and var in grouping_factors:
-            return True
-
-    return False

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/internal/viz/predict.py

@@ -36,6 +36,142 @@ if TYPE_CHECKING:
 __all__ = ["plot_predict"]
 
 
+# ---------------------------------------------------------------------------
+# Formula parsing helpers
+# ---------------------------------------------------------------------------
+
+_FORMULA_CHARS = {"~", "@", "(", "["}
+
+
+def _is_formula(term: str) -> bool:
+    """Return True if *term* looks like formula syntax rather than a bare column name.
+
+    Args:
+        term: The user-supplied term string.
+
+    Returns:
+        True when the string contains ``~``, ``@``, or ``(``.
+    """
+    return bool(_FORMULA_CHARS & set(term))
+
+
+def _parse_predict_formula(
+    formula: str,
+    model: "Any",
+) -> dict[str, Any]:
+    """Translate an explore-style formula into :func:`plot_predict` kwargs.
+
+    Parses ``formula`` using the existing :func:`parse_explore_formula`
+    machinery, then maps the resulting :class:`ExploreFormulaSpec` fields
+    onto the keyword arguments that :func:`plot_predict` understands.
+
+    Contrast formulas (``pairwise(treatment)``, ``Drug[A - B]``) are
+    rejected with a ``ValueError`` — use :func:`plot_explore` instead.
+
+    Args:
+        formula: An explore-style formula string (e.g. ``"age ~ sex"``).
+        model: A fitted bossanova model (used for data access and term
+            validation).
+
+    Returns:
+        Dict with keys: ``term``, and optionally ``hue``, ``col``, ``row``,
+        ``at``, ``n_points``, ``focal_values``.
+
+    Raises:
+        ValueError: If the formula contains contrasts.
+    """
+    from bossanova.internal.marginal import parse_explore_formula
+
+    spec = parse_explore_formula(formula)
+
+    # Reject contrast formulas — they belong in plot_explore
+    if spec.has_contrast:
+        raise ValueError(
+            f"Contrast formula {formula!r} is not supported by plot_predict(). "
+            "Use plot_explore() for contrasts."
+        )
+    # Also reject RHS bracket contrasts
+    if spec.has_rhs_contrasts:
+        raise ValueError(
+            f"Contrast formula {formula!r} is not supported by plot_predict(). "
+            "Use plot_explore() for contrasts."
+        )
+
+    result: dict[str, Any] = {"term": spec.focal_var}
+    data = model.data
+
+    # Handle focal variable modifiers (@range, @[values], @quantile)
+    if spec.focal_at_range is not None:
+        result["n_points"] = spec.focal_at_range
+    elif spec.focal_at_values is not None:
+        result["focal_values"] = list(spec.focal_at_values)
+    elif spec.focal_at_quantile is not None:
+        col_data = data[spec.focal_var].drop_nulls().to_numpy()
+        quantiles = np.linspace(0, 1, spec.focal_at_quantile + 2)[1:-1]
+        result["focal_values"] = list(np.quantile(col_data, quantiles))
+
+    # Classify conditions into faceting roles (hue/col/row) or at-pins
+    at: dict[str, Any] = {}
+    facet_slots: list[str] = []  # hue, col, row — assigned in order
+
+    for cond in spec.conditions:
+        resolved = _resolve_condition_values(cond, data)
+        if resolved is not None and len(resolved) == 1:
+            # Single pinned value → goes into at dict
+            at[cond.var] = resolved[0]
+        else:
+            # Bare or multi-value condition → assigned to facet slot
+            if len(facet_slots) >= 3:
+                raise ValueError(
+                    f"Too many unpinned conditions in formula {formula!r}. "
+                    "plot_predict supports at most 3 faceting variables "
+                    "(hue, col, row). Pin additional conditions with @value."
+                )
+            facet_slots.append(cond.var)
+            # If condition has explicit multi-values, inject via at as list
+            if resolved is not None:
+                at[cond.var] = resolved
+
+    # Map facet slots to hue → col → row
+    facet_keys = ["hue", "col", "row"]
+    for key, var in zip(facet_keys, facet_slots):
+        result[key] = var
+
+    if at:
+        result["at"] = at
+
+    return result
+
+
+def _resolve_condition_values(
+    cond: Any,
+    data: pl.DataFrame,
+) -> list | None:
+    """Resolve a :class:`Condition` to concrete values or ``None``.
+
+    Args:
+        cond: A ``Condition`` from :func:`parse_explore_formula`.
+        data: The model's training data.
+
+    Returns:
+        A list of concrete values if the condition specifies explicit
+        values (``at_values``, ``at_range``, ``at_quantile``), or
+        ``None`` for bare conditions (use all unique levels).
+    """
+    if cond.at_values is not None:
+        return list(cond.at_values)
+    if cond.at_range is not None:
+        col_data = data[cond.var].drop_nulls().to_numpy()
+        return list(
+            np.linspace(np.nanmin(col_data), np.nanmax(col_data), cond.at_range)
+        )
+    if cond.at_quantile is not None:
+        col_data = data[cond.var].drop_nulls().to_numpy()
+        quantiles = np.linspace(0, 1, cond.at_quantile + 2)[1:-1]
+        return list(np.quantile(col_data, quantiles))
+    return None
+
+
 def plot_predict(
     model: "Any",
     term: str,
@@ -65,9 +201,21 @@ def plot_predict(
     holding others at reference values (means for continuous, reference level
     for categorical).
 
+    Accepts either a bare column name or an explore-style formula::
+
+        plot_predict(m, "age")              # bare column
+        plot_predict(m, "age ~ sex")        # hue by sex
+        plot_predict(m, "age ~ sex@Female") # pin sex=Female
+        plot_predict(m, "age@range(5)")     # 5-point grid
+        plot_predict(m, "age@[25,50,75]")   # explicit grid values
+
+    Contrast formulas (``pairwise()``, ``Drug[A - B]``) are not supported —
+    use :func:`plot_explore` for those.
+
     Args:
         model: A fitted bossanova model.
-        term: The predictor variable to vary. Can be continuous or categorical.
+        term: The predictor variable to vary, or an explore-style formula
+            string. Can be continuous or categorical.
         hue: Column name for color encoding (creates separate lines per level).
         col: Column name for faceting into columns.
         row: Column name for faceting into rows.
@@ -121,7 +269,6 @@ def plot_predict(
             viz.plot_predict(m, "Days", show_blups=True)
 
     See Also:
-        plot_fit: Observed vs Predicted scatter plot.
         plot_explore: Marginal effects/means visualization.
     """
     import seaborn as sns
@@ -129,6 +276,22 @@ def plot_predict(
     if not model._is_fitted:
         raise RuntimeError("Model must be fitted before plotting predictions")
 
+    # Formula detection: if term contains ~, @, or ( → parse as formula
+    focal_values: list[float | str] | None = None
+    if _is_formula(term):
+        parsed = _parse_predict_formula(term, model)
+        term = parsed["term"]
+        # Override kwargs only if formula provided them and user didn't
+        hue = hue or parsed.get("hue")
+        col = col or parsed.get("col")
+        row = row or parsed.get("row")
+        if parsed.get("n_points") is not None:
+            n_points = parsed["n_points"]
+        focal_values = parsed.get("focal_values")
+        # Merge at dicts: user-supplied at takes precedence
+        if parsed.get("at"):
+            at = {**parsed["at"], **(at or {})}
+
     style = BOSSANOVA_STYLE
     if palette is None:
         palette = style["palette"]
@@ -159,6 +322,7 @@ def plot_predict(
         conf_int=conf_int,
         show_blups=show_blups,
         groups=groups,
+        focal_values=focal_values,
     )
 
     # Create the plot
@@ -209,8 +373,32 @@ def build_prediction_data(
     conf_int: float,
     show_blups: bool,
     groups: list[str] | None,
+    focal_values: list[float | str] | None = None,
 ) -> pl.DataFrame:
-    """Build prediction DataFrame with fitted values and intervals."""
+    """Build prediction DataFrame with fitted values and intervals.
+
+    Args:
+        model: A fitted bossanova model.
+        data: The model's training data.
+        term: The focal predictor variable to vary.
+        hue: Column for color encoding.
+        col: Column for column faceting.
+        row: Column for row faceting.
+        at: Dictionary of predictor values to hold fixed. Values can be
+            scalars (pinned to one value) or lists (expanded into the grid).
+        n_points: Number of points for continuous predictors.
+        is_categorical: Whether term is categorical.
+        pred_type: Prediction type ("data" or "link").
+        interval: Interval type ("confidence", "prediction", or None).
+        conf_int: Confidence level.
+        show_blups: Whether to show group-specific BLUP lines.
+        groups: Which group levels to show for BLUPs.
+        focal_values: Explicit values for the focal variable's grid.
+            When set, overrides the default linspace grid.
+
+    Returns:
+        Polars DataFrame with prediction grid and fitted values.
+    """
     response_col = get_model_formula(model).split("~")[0].strip()
     grouping_factors = get_grouping_factors(model)
     is_mixed = is_mixed_model(model)
@@ -239,8 +427,10 @@ def build_prediction_data(
             continue  # Skip if not showing BLUPs
 
         if col_name == term:
-            # Vary this term
-            if is_categorical:
+            # Vary this term — use focal_values if provided
+            if focal_values is not None:
+                grid_data[col_name] = list(focal_values)
+            elif is_categorical:
                 grid_data[col_name] = data[col_name].unique().sort().to_list()
             else:
                 col_data = data[col_name].to_numpy()
@@ -249,7 +439,9 @@ def build_prediction_data(
                     np.linspace(np.nanmin(col_data), np.nanmax(col_data), n_points)
                 )
         elif at and col_name in at:
-            grid_data[col_name] = [at[col_name]]
+            # Support list-valued at entries (expand into grid)
+            val = at[col_name]
+            grid_data[col_name] = list(val) if isinstance(val, list) else [val]
         elif col_name in [hue, col, row] and col_name is not None:
             # Vary grouping variables
             grid_data[col_name] = data[col_name].unique().sort().to_list()
@@ -472,11 +664,6 @@ def add_ci_band(
     confidence bounds. Supports both ``ci_lower``/``ci_upper``
     (preferred) and ``lwr``/``upr`` (legacy) column naming.
 
-    Unlike ``fit_layers.add_ci_band``, this version receives
-    pre-filtered data from ``FacetGrid.map_dataframe()`` and generates
-    hue colors dynamically via ``plt.cm.tab10`` rather than using a
-    pre-computed color map.
-
     Delegates rendering to :func:`~bossanova.internal.viz.helpers.fill_ci_band`.
 
     Args:

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/bossanova/model/core.py

@@ -1968,8 +1968,16 @@ class model:
     def plot_predict(self, term: str, **kwargs: object) -> object:
         """Plot marginal predictions across a predictor range.
 
+        Accepts either a bare column name or an explore-style formula::
+
+            m.plot_predict("age")              # bare column
+            m.plot_predict("age ~ sex")        # hue by sex
+            m.plot_predict("age ~ sex@Female") # pin sex=Female
+            m.plot_predict("age@range(5)")     # 5-point grid
+            m.plot_predict("age@[25,50,75]")   # explicit grid values
+
         Args:
-            term: Predictor variable to plot predictions for.
+            term: Predictor variable or explore-style formula string.
             **kwargs: Forwarded to ``bossanova.viz.plot_predict``. Key
                 options: ``hue``, ``col``, ``at``, ``interval``,
                 ``show_data``, ``show_rug``, ``height``, ``aspect``.
@@ -1979,20 +1987,6 @@ class model:
         """
         return self._plot("plot_predict", term, **kwargs)
 
-    def plot_fit(self, x: str, **kwargs: object) -> object:
-        """Plot fitted model against observed data.
-
-        Args:
-            x: Predictor variable for the x-axis.
-            **kwargs: Forwarded to ``bossanova.viz.plot_fit``. Key
-                options: ``hue``, ``col``, ``scatter``, ``fit_line``,
-                ``ci``, ``blups``, ``groups``, ``height``.
-
-        Returns:
-            object: Matplotlib figure.
-        """
-        return self._plot("plot_fit", x, **kwargs)
-
     def plot_explore(self, specs: str, **kwargs: object) -> object:
         """Plot marginal effects or estimated marginal means.
 

{bossanova-0.1.0.dev22 → bossanova-0.1.0.dev24}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "bossanova"
 requires-python = ">=3.11"
-version = "0.1.0.dev22"
+version = "0.1.0.dev24"
 authors = [
     {name = "Eshin Jolly"},
 ]