bossanova 0.1.0.dev15__tar.gz → 0.1.0.dev16__tar.gz

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/PKG-INFO

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

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/containers/builders/state.py

@@ -197,6 +197,7 @@ def build_mee_state(
     mee_type: str,
     *,
     units: str = "link",
+    weights: str = "equal",
     L_matrix: np.ndarray | None = None,
     contrast_method: str | None = None,
     n_contrast_levels: int | None = None,
@@ -219,6 +220,8 @@ def build_mee_state(
         focal_var: The primary variable being explored.
         mee_type: Type of effect ("means", "slopes", "contrasts").
         units: Scale of the estimates: "link" (linear predictor) or "data" (response).
+        weights: Averaging philosophy: ``"equal"`` (balanced reference grid,
+            emmeans-style) or ``"observed"`` (g-computation over observed data).
         L_matrix: Design matrix for delta method inference (optional).
             Shape (n_estimates, n_coef). For EMMs this is X_ref.
         contrast_method: Original contrast type for multiplicity adjustment
@@ -258,6 +261,7 @@ def build_mee_state(
         focal_var=focal_var,
         type=mee_type,
         units=units,
+        weights=weights,
         L_matrix=L_matrix,
         contrast_method=contrast_method,
         n_contrast_levels=n_contrast_levels,

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/containers/structs/state.py

@@ -258,7 +258,9 @@ class MeeState:
     focal_var: str = field(validator=validators.instance_of(str))
     type: str = field(validator=validators.in_(("means", "slopes", "contrasts")))
     units: str = field(default="link", validator=validators.in_(("link", "data")))
-    weights: str = field(default="equal", validator=validators.in_(("equal", "observed")))
+    weights: str = field(
+        default="equal", validator=validators.in_(("equal", "observed"))
+    )
 
     # Design matrix for delta method inference (optional)
     # Shape: (n_estimates, n_coef). For EMMs, this is X_ref; for slopes, a selector row.

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/maths/design/reference.py

@@ -9,6 +9,7 @@ from bossanova.internal.maths.design.names import (
 
 __all__ = [
     "build_continuous_reference_matrix",
+    "build_counterfactual_design_matrices",
     "build_reference_design_matrix",
     "build_reference_row",
 ]
@@ -248,6 +249,98 @@ def build_continuous_reference_matrix(
     return X_ref
 
 
+def build_counterfactual_design_matrices(
+    X: np.ndarray,
+    X_names: tuple[str, ...] | list[str],
+    focal_var: str,
+    levels: list[str],
+) -> list[np.ndarray]:
+    """Build counterfactual design matrices for g-computation.
+
+    For each focal level, creates a modified copy of the full design matrix
+    ``X`` where the focal variable's indicator columns are set to match that
+    level and interaction columns involving the focal variable are recomputed
+    from component values.  Non-focal columns are left unchanged, preserving
+    the observed covariate distribution.
+
+    This is the core building block for ``weights="observed"``
+    (g-computation / counterfactual prediction).  Each returned matrix answers:
+    "What would the design matrix look like if every observation were assigned
+    to this focal level?"
+
+    Args:
+        X: Original design matrix, shape ``(N, p)``.
+        X_names: Column names from the design matrix, in order.
+        focal_var: Name of the categorical focal variable.
+        levels: List of levels to compute counterfactuals for.
+
+    Returns:
+        List of counterfactual design matrices (one per level), each shape
+        ``(N, p)``.  Order matches ``levels``.
+
+    Examples:
+        For a model ``y ~ treatment + x + treatment:x`` with
+        ``X_names = ("Intercept", "x", "treatment[B]", "x:treatment[B]")``::
+
+            mats = build_counterfactual_design_matrices(X, X_names, "treatment", ["ref", "B"])
+            # mats[0]: treatment set to ref for all rows (treatment[B]=0, x:treatment[B]=0)
+            # mats[1]: treatment set to B for all rows (treatment[B]=1, x:treatment[B]=x_i)
+    """
+    X_names_list = list(X_names)
+    N, p = X.shape
+
+    # Pre-parse column metadata once
+    col_infos = [parse_design_column_name(name) for name in X_names_list]
+
+    # Identify focal dummy columns and their levels
+    focal_col_indices: list[int] = []
+    focal_col_levels: list[str | None] = []
+    for j, info in enumerate(col_infos):
+        if info.column_type == "categorical" and info.base_term == focal_var:
+            focal_col_indices.append(j)
+            focal_col_levels.append(info.level)
+
+    # Identify interaction columns involving the focal variable
+    interaction_cols: list[int] = []
+    for j, info in enumerate(col_infos):
+        if info.is_interaction:
+            parts = X_names_list[j].split(":")
+            if any(
+                parse_design_column_name(part).base_term == focal_var for part in parts
+            ):
+                interaction_cols.append(j)
+
+    result: list[np.ndarray] = []
+    for level in levels:
+        X_cf = X.copy()
+
+        # Set focal dummy columns
+        for j, lvl in zip(focal_col_indices, focal_col_levels):
+            X_cf[:, j] = 1.0 if lvl == level else 0.0
+
+        # Recompute interaction columns involving the focal variable
+        for j in interaction_cols:
+            parts = X_names_list[j].split(":")
+            col_product = np.ones(N, dtype=np.float64)
+            for part in parts:
+                pinfo = parse_design_column_name(part)
+                if pinfo.base_term == focal_var:
+                    # Focal component: use counterfactual indicator
+                    if pinfo.column_type == "categorical":
+                        col_product *= 1.0 if pinfo.level == level else 0.0
+                    elif part in X_names_list:
+                        # Continuous focal in interaction (rare but possible)
+                        col_product *= X[:, X_names_list.index(part)]
+                elif part in X_names_list:
+                    # Non-focal component: use original observed values
+                    col_product *= X[:, X_names_list.index(part)]
+            X_cf[:, j] = col_product
+
+        result.append(X_cf)
+
+    return result
+
+
 # ---------------------------------------------------------------------------
 # Internal helpers
 # ---------------------------------------------------------------------------

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/operations/marginal/compute.py

@@ -189,6 +189,19 @@ def _dispatch_marginal_core(
     if varying not in ("exclude", "include"):
         raise ValueError(f"varying must be 'exclude' or 'include', got {varying!r}")
 
+    # Extract weights= kwarg with smart default
+    weights_raw = kwargs.pop("weights", None)
+    if weights_raw is not None:
+        weights = str(weights_raw)
+        if weights not in ("equal", "observed"):
+            raise ValueError(f"weights must be 'equal' or 'observed', got {weights!r}")
+    else:
+        # Smart default: "observed" for GLMs (non-identity link), "equal" otherwise
+        if spec is not None and spec.link != "identity":
+            weights = "observed"
+        else:
+            weights = "equal"
+
     # Extract inverse_transforms= kwarg (default True)
     inverse_transforms = bool(kwargs.pop("inverse_transforms", True))
 
@@ -253,6 +266,7 @@ def _dispatch_marginal_core(
             data=data,
             spec=spec,
             units=units,
+            weights=weights,
             resolved=resolved,
         )
 
@@ -267,6 +281,7 @@ def _dispatch_marginal_core(
             data,
             spec=spec,
             units=units,
+            weights=weights,
             resolved=resolved,
             focal_at_values=parsed.focal_at_values,
             contrast_ref=parsed.contrast_ref,
@@ -290,6 +305,7 @@ def _dispatch_marginal_core(
                 set_categoricals=resolved.set_categoricals or None,
                 spec=spec,
                 units=units,
+                weights=weights,
             )
         # Continuous at-values (e.g. "Days@[0, 3, 6, 9]"): EMMs at
         # specific values.  Forward-transform the at-values through the
@@ -341,6 +357,7 @@ def _dispatch_marginal_core(
             set_categoricals=resolved.set_categoricals or None,
             spec=spec,
             units=units,
+            weights=weights,
         )
 
     # Continuous focal variable → slopes
@@ -390,6 +407,7 @@ def _dispatch_marginal_core(
         spec=spec,
         formula_spec=fspec,
         units=units,
+        weights=weights,
     )
 
 
@@ -549,6 +567,7 @@ def _compute_emm_categorical(
     set_categoricals: dict[str, str] | None = None,
     spec: ModelSpec | None = None,
     units: str = "link",
+    weights: str = "equal",
 ) -> MeeState:
     """Compute EMMs for a categorical focal variable.
 
@@ -563,6 +582,7 @@ def _compute_emm_categorical(
             specific levels (e.g. ``{"Ethnicity": "Asian"}``).
         spec: ModelSpec with link info (for units="data").
         units: Scale of estimates.
+        weights: Averaging philosophy: ``"equal"`` or ``"observed"``.
 
     Returns:
         MeeState with grid of levels and their estimated means.
@@ -577,6 +597,7 @@ def _compute_emm_categorical(
         set_categoricals=set_categoricals,
         spec=spec,
         units=units,
+        weights=weights,
     )
 
 
@@ -779,6 +800,7 @@ def _compute_marginal_slope(
     spec: ModelSpec | None = None,
     formula_spec: FormulaSpec | None = None,
     units: str = "link",
+    weights: str = "equal",
 ) -> MeeState:
     """Compute marginal slope for a continuous focal variable.
 
@@ -797,6 +819,7 @@ def _compute_marginal_slope(
         spec: ModelSpec with family/link info.
         formula_spec: FormulaSpec with learned encoding (for finite-diff).
         units: Scale of estimates: ``"link"`` or ``"data"``.
+        weights: Averaging philosophy: ``"equal"`` or ``"observed"``.
 
     Returns:
         MeeState with marginal slope estimate.
@@ -825,6 +848,7 @@ def _compute_marginal_slope(
             formula_spec=formula_spec,
             data=data,
             units=units,
+            weights=weights,
         )
 
     # Fast path: coefficient extraction — use resolved name
@@ -1021,6 +1045,7 @@ def _compute_contrasts(
     *,
     spec: ModelSpec | None = None,
     units: str = "link",
+    weights: str = "equal",
     resolved: ResolvedConditions | None = None,
     focal_at_values: tuple[float | str, ...] | None = None,
     contrast_ref: str | None = None,
@@ -1030,7 +1055,7 @@ def _compute_contrasts(
     First computes EMMs, then applies the requested contrast matrix.
     When ``resolved`` contains grid conditions, computes crossed EMMs and
     applies grouped contrasts.  When ``focal_at_values`` is provided
-    (e.g. from ``pairwise(cyl[4, 8])``), contrasts are computed only
+    (e.g. from ``pairwise(cyl@[4, 8])``), contrasts are computed only
     over the requested subset of levels.
 
     Args:
@@ -1043,6 +1068,7 @@ def _compute_contrasts(
         data: Model data for level extraction.
         spec: ModelSpec with link info (for units="data").
         units: Scale of estimates.
+        weights: Averaging philosophy: ``"equal"`` or ``"observed"``.
         resolved: Resolved conditions for conditioning.
         focal_at_values: Optional subset of levels from at-spec syntax
             (e.g. ``pairwise(cyl@[4, 8])``).
@@ -1113,6 +1139,7 @@ def _compute_contrasts(
         set_categoricals=set_cats,
         spec=spec,
         units=units,
+        weights=weights,
     )
     if contrast_ref is not None:
         ref_idx = _resolve_ref_idx(emm_state.grid)
@@ -1130,6 +1157,7 @@ def _compute_bracket_contrasts(
     data: pl.DataFrame,
     spec: ModelSpec | None = None,
     units: str = "link",
+    weights: str = "equal",
     resolved: ResolvedConditions | None = None,
 ) -> MeeState:
     """Compute bracket contrast expression for a categorical focal variable.
@@ -1146,6 +1174,7 @@ def _compute_bracket_contrasts(
         data: Model data DataFrame.
         spec: ModelSpec with link info (for units="data").
         units: Scale of estimates.
+        weights: Averaging philosophy: ``"equal"`` or ``"observed"``.
         resolved: Resolved conditions for conditioning.
 
     Returns:
@@ -1179,5 +1208,6 @@ def _compute_bracket_contrasts(
         set_categoricals=set_cats,
         spec=spec,
         units=units,
+        weights=weights,
     )
     return apply_bracket_contrasts(emm_state, contrast_expr)

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/operations/marginal/emm.py

@@ -16,7 +16,10 @@ import polars as pl
 
 from bossanova.internal.containers.builders.state import build_mee_state
 from bossanova.internal.containers.structs.state import MeeState
-from bossanova.internal.maths.design.reference import build_reference_design_matrix
+from bossanova.internal.maths.design.reference import (
+    build_counterfactual_design_matrices,
+    build_reference_design_matrix,
+)
 from bossanova.internal.operations.common.factors import get_factor_levels
 from bossanova.internal.operations.marginal.validation import validate_focal_var
 
@@ -41,18 +44,21 @@ def compute_emm(
     set_categoricals: dict[str, str] | None = None,
     spec: object | None = None,
     units: str = "link",
+    weights: str = "equal",
 ) -> MeeState:
     """Compute estimated marginal means for a categorical focal variable.
 
-    Computes predictions at each level of the focal variable, with other
-    covariates set to their reference values (means for continuous, 0 for
-    non-focal categorical).
+    Supports two averaging philosophies via the ``weights`` parameter:
 
-    This matches the emmeans package behavior: EMMs are model predictions
-    at a reference grid where:
-    - The focal variable varies across its levels
-    - Other continuous covariates are set to their overall means
-    - Other categorical covariates are set to the reference level (0 in dummy coding)
+    - ``"equal"`` (default): Balanced reference grid (emmeans-style).
+      Predictions at a grid where covariates are at their means.
+    - ``"observed"``: G-computation / counterfactual prediction.  For each
+      focal level, sets every observation to that level and averages the
+      resulting predictions.  Preserves the observed covariate distribution.
+
+    For linear models (identity link), both approaches give identical results.
+    For GLMs (non-identity link), they diverge because
+    ``mean(g⁻¹(Xᵢβ)) ≠ g⁻¹(mean(Xᵢ) · β)``.
 
     Args:
         bundle: DataBundle with model data and metadata. Used to extract:
@@ -70,6 +76,11 @@ def compute_emm(
         set_categoricals: Optional dict mapping non-focal categorical variable
             names to specific levels to pin them at (instead of marginalizing
             at column means). E.g. ``{"Ethnicity": "Asian"}``.
+        spec: ModelSpec with link/family info (needed for units="data").
+        units: Scale of estimates: ``"link"`` or ``"data"``.
+        weights: Averaging philosophy: ``"equal"`` for balanced reference
+            grid (emmeans-style), ``"observed"`` for g-computation over
+            observed data.
 
     Returns:
         MeeState with grid of levels and their estimated means.
@@ -78,17 +89,20 @@ def compute_emm(
         ValueError: If focal_var not found in bundle.factor_levels.
 
     Examples:
-        Compute EMMs for treatment levels::
+        Compute EMMs with default (reference grid) averaging::
 
-            from bossanova.internal.operations.marginal import compute_emm
             mee = compute_emm(bundle, fit, "treatment", "treatment")
-            # mee.estimate: array([2.1, 3.4, 2.8]) for each level
-            # mee.grid: DataFrame with "treatment" column
+
+        G-computation averaging for a logistic regression::
+
+            mee = compute_emm(bundle, fit, "treatment", "treatment",
+                              spec=spec, units="data", weights="observed")
 
     Note:
-        The X_ref matrix is constructed internally using bundle.X_names
-        to identify which columns correspond to the focal variable dummies.
-        This works for standard treatment/dummy coding.
+        For ``weights="observed"`` with ``units="data"`` on a non-identity
+        link, confidence intervals use the response-scale delta method
+        (symmetric CIs) rather than the link-scale back-transformation
+        (asymmetric CIs) used by ``weights="equal"``.
     """
     validate_focal_var(bundle, focal_var)
 
@@ -96,10 +110,8 @@ def compute_emm(
     if levels is None:
         levels = get_factor_levels(bundle, focal_var)
 
-    # Build reference X matrix using shared utility
-    X_means = bundle.X.mean(axis=0).copy()
-
     # Apply at-value overrides for conditioning
+    X_means = bundle.X.mean(axis=0).copy()
     if at_overrides:
         X_names_list = list(bundle.X_names)
         for var_name, value in at_overrides.items():
@@ -117,6 +129,29 @@ def compute_emm(
                 stacklevel=4,
             )
 
+    # G-computation path: weights="observed" with non-identity link on data scale.
+    # For link scale or identity link, the two approaches are mathematically
+    # equivalent, so we fall through to the reference grid path for efficiency.
+    use_gcomp = (
+        weights == "observed"
+        and units == "data"
+        and spec is not None
+        and spec.link != "identity"
+    )
+
+    if use_gcomp:
+        return _compute_emm_gcomp(
+            bundle=bundle,
+            fit=fit,
+            focal_var=focal_var,
+            explore_formula=explore_formula,
+            levels=levels,
+            spec=spec,
+            at_overrides=at_overrides,
+            set_categoricals=set_categoricals,
+        )
+
+    # Reference grid path (weights="equal" or equivalent cases)
     X_ref = build_reference_design_matrix(
         X_names=bundle.X_names,
         focal_var=focal_var,
@@ -149,9 +184,143 @@ def compute_emm(
         L_matrix_link = X_ref
 
     # Create grid DataFrame with conditioning columns first
+    grid = _build_emm_grid(focal_var, levels, at_overrides, set_categoricals)
+
+    return build_mee_state(
+        grid=grid,
+        estimate=estimates,
+        explore_formula=explore_formula,
+        focal_var=focal_var,
+        mee_type="means",
+        units=units,
+        weights="equal",
+        L_matrix=L_matrix,
+        link=link_name,
+        L_matrix_link=L_matrix_link,
+    )
+
+
+def _compute_emm_gcomp(
+    bundle: "DataBundle",
+    fit: "FitState",
+    focal_var: str,
+    explore_formula: str,
+    levels: list[str],
+    spec: object,
+    *,
+    at_overrides: dict[str, float] | None = None,
+    set_categoricals: dict[str, str] | None = None,
+) -> MeeState:
+    """G-computation EMMs: average counterfactual predictions over observed data.
+
+    For each focal level k:
+        1. Build counterfactual X_cf_k (all N rows, focal set to level k)
+        2. η_k = X_cf_k @ β  (N linear predictors)
+        3. pred_k = g⁻¹(η_k)  (N response-scale predictions)
+        4. emm_k = mean(pred_k)
+        5. L_k = mean(dμ/dη(η_k)[:, None] * X_cf_k, axis=0)  (delta method)
+
+    Args:
+        bundle: DataBundle with model data and design matrix.
+        fit: FitState with fitted coefficients.
+        focal_var: Name of the categorical focal variable.
+        explore_formula: Explore formula string for metadata.
+        levels: Focal variable levels to compute EMMs for.
+        spec: ModelSpec with link function info.
+        at_overrides: Optional covariate overrides (applied before
+            counterfactual construction).
+        set_categoricals: Optional non-focal categorical pins.
+
+    Returns:
+        MeeState with g-computation estimates and response-scale L_matrix.
+    """
+    from bossanova.internal.maths.family.links import (
+        apply_link_inverse,
+        apply_link_inverse_deriv,
+    )
+
+    X = bundle.X
+
+    # Apply at_overrides: replace column values in X before counterfactual
+    if at_overrides:
+        X = X.copy()
+        X_names_list = list(bundle.X_names)
+        for var_name, value in at_overrides.items():
+            if var_name in X_names_list:
+                idx = X_names_list.index(var_name)
+                X[:, idx] = value
+
+    # Apply set_categoricals: pin non-focal categoricals in X
+    if set_categoricals:
+        from bossanova.internal.maths.design.names import parse_design_column_name
+
+        X = X if at_overrides else X.copy()  # ensure we have a copy
+        for j, name in enumerate(bundle.X_names):
+            info = parse_design_column_name(name)
+            if info.column_type == "categorical" and info.base_term in set_categoricals:
+                X[:, j] = 1.0 if info.level == set_categoricals[info.base_term] else 0.0
+
+    # Build counterfactual design matrices
+    cf_matrices = build_counterfactual_design_matrices(
+        X=X,
+        X_names=bundle.X_names,
+        focal_var=focal_var,
+        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)
+
+    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)
+
+    # Build grid DataFrame
+    grid = _build_emm_grid(focal_var, levels, at_overrides, set_categoricals)
+
+    # G-computation: no link/L_matrix_link — inference uses response-scale
+    # delta method directly (symmetric CIs on response scale).
+    return build_mee_state(
+        grid=grid,
+        estimate=estimates,
+        explore_formula=explore_formula,
+        focal_var=focal_var,
+        mee_type="means",
+        units="data",
+        weights="observed",
+        L_matrix=L_matrix,
+        link=None,
+        L_matrix_link=None,
+    )
+
+
+def _build_emm_grid(
+    focal_var: str,
+    levels: list[str],
+    at_overrides: dict[str, float] | None,
+    set_categoricals: dict[str, str] | None,
+) -> pl.DataFrame:
+    """Build the grid DataFrame for EMM results.
+
+    Args:
+        focal_var: Name of the focal variable.
+        levels: Focal variable levels.
+        at_overrides: Optional covariate overrides to include as columns.
+        set_categoricals: Optional categorical pins to include as columns.
+
+    Returns:
+        Polars DataFrame with condition columns first, then focal column.
+    """
     grid = pl.DataFrame({focal_var: levels})
 
-    # Add set_categoricals as condition columns
     if set_categoricals:
         for cond_var, cond_level in set_categoricals.items():
             grid = grid.with_columns(pl.lit(cond_level).alias(cond_var))
@@ -165,17 +334,7 @@ def compute_emm(
     if cond_cols:
         grid = grid.select(cond_cols + [c for c in grid.columns if c not in cond_cols])
 
-    return build_mee_state(
-        grid=grid,
-        estimate=estimates,
-        explore_formula=explore_formula,
-        focal_var=focal_var,
-        mee_type="means",
-        units=units,
-        L_matrix=L_matrix,
-        link=link_name,
-        L_matrix_link=L_matrix_link,
-    )
+    return grid
 
 
 def compute_emm_with_xref(

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/internal/operations/marginal/slopes.py

@@ -251,22 +251,24 @@ def compute_slopes_finite_diff(
     formula_spec: FormulaSpec,
     data: pl.DataFrame,
     units: str = "link",
+    weights: str = "equal",
     delta_frac: float = 0.001,
 ) -> MeeState:
     """Compute marginal slopes via centered finite differences (AME).
 
-    Builds a reference grid (Cartesian product of categorical levels,
-    continuous covariates at means), perturbs the focal variable by
-    ``±delta/2``, constructs design matrices via ``evaluate_newdata``,
-    and differences them to obtain the Average Marginal Effect.
+    Supports two averaging philosophies via the ``weights`` parameter:
 
-    For ``units="data"`` on GLMs the response-scale derivative
-    ``dμ/dη`` is applied row-wise before averaging, giving a true AME
-    (not MEM).
+    - ``"equal"``: Average over a balanced reference grid (Cartesian product
+      of categorical levels, continuous covariates at means).  Matches R's
+      ``emmeans::emtrends``.
+    - ``"observed"``: Average over actual data rows.  Gives a true Average
+      Marginal Effect (AME) that preserves the observed covariate distribution.
 
-    Algorithm (matches R ``emmeans::emtrends`` conventions):
+    For linear models (identity link), both approaches give identical results.
+
+    Algorithm:
         1. ``delta = delta_frac × range(focal_var)``
-        2. Build reference grid (factor Cartesian product, numerics at means)
+        2. Build evaluation grid (balanced or observed data)
         3. Perturb: ``grid_plus = grid[focal + delta/2]``,
            ``grid_minus = grid[focal − delta/2]``
         4. ``X_plus, X_minus = evaluate_newdata(formula_spec, grid_*)``
@@ -283,6 +285,8 @@ def compute_slopes_finite_diff(
         formula_spec: FormulaSpec with learned encoding for ``evaluate_newdata``.
         data: Raw model data (for computing ranges and covariate means).
         units: ``"link"`` (linear predictor scale) or ``"data"`` (response).
+        weights: ``"equal"`` for balanced reference grid, ``"observed"``
+            for actual data rows.
         delta_frac: Fraction of the focal variable's range used as the
             finite-difference step size.  Default ``0.001`` matches R's
             ``emmeans``.
@@ -309,8 +313,15 @@ def compute_slopes_finite_diff(
         )
     delta = delta_frac * var_range
 
-    # --- reference grid ---------------------------------------------------
-    grid = _build_slopes_data_grid(bundle, data)
+    # --- evaluation grid --------------------------------------------------
+    # weights="observed": use actual data rows (true AME)
+    # weights="equal": use balanced reference grid (emtrends-style)
+    if weights == "observed":
+        # Use actual data rows, dropping the response column
+        grid_cols = [c for c in data.columns if c != bundle.y_name]
+        grid = data.select(grid_cols)
+    else:
+        grid = _build_slopes_data_grid(bundle, data)
 
     # --- perturbed grids --------------------------------------------------
     grid_plus = grid.with_columns(
@@ -351,6 +362,7 @@ def compute_slopes_finite_diff(
         focal_var=focal_var,
         mee_type="slopes",
         units=units,
+        weights=weights,
         L_matrix=L_avg,
     )
 

{bossanova-0.1.0.dev15 → bossanova-0.1.0.dev16}/bossanova/model/core.py

@@ -1247,6 +1247,10 @@ class model:
                   use transformed-scale names/values directly.
                 - ``units`` (str): ``"link"`` or ``"data"``.
                 - ``varying`` (str): ``"exclude"`` or ``"include"``.
+                - ``weights`` (str): ``"equal"`` for emmeans-style balanced
+                  reference grid, ``"observed"`` for g-computation over
+                  observed data.  Default is smart: ``"observed"`` for GLMs
+                  (non-identity link), ``"equal"`` for linear models.
 
         Returns:
             self: For method chaining. Results in ``.effects``.