marginal

Marginal effects and estimated marginal means computation.

Call chain:

model.explore() -> build_reference_grid() -> compute_emm() / compute_slopes() -> compute_contrasts()

Classes:

Functions:

Modules:

Classes¶

Condition¶

A conditioning specification in explore formula.

Represents a variable to condition on, optionally with specific values, a method for generating values, or a bracket contrast expression.

Created by: parse_explore_formula(), ExploreParser Consumed by: resolve_all_conditions(), dispatch_marginal_computation() Augmented by: Never

Attributes:

Attributes¶

at_quantile¶

at_quantile: int | None = field(default=None, validator=is_optional_positive_int)

at_range¶

at_range: int | None = field(default=None, validator=is_optional_positive_int)

at_values¶

at_values: tuple | None = field(default=None)

contrast_expr¶

contrast_expr: ContrastExpr | None = field(default=None, validator=(validators.optional(validators.instance_of(ContrastExpr))))

var¶

var: str = field(validator=(validators.instance_of(str)))

ExploreFormulaError¶

ExploreFormulaError(message: str, formula: str, position: int | None = None) -> None

Bases: ValueError

Error in explore formula syntax.

Provides helpful error messages with position indicators for syntax errors.

Parameters:

Attributes:

Parameters:

Attributes¶

formula¶

formula = formula

position¶

position = position

ExploreFormulaSpec¶

Parsed explore formula.

Represents a parsed explore formula with focal variable, optional contrast type, and conditioning specifications.

Created by: parse_explore_formula() Consumed by: dispatch_marginal_computation(), plot_predict(), plot_explore() Augmented by: attrs.evolve() in resolve_focal_at_spec() materializes focal at-values

Attributes:

Attributes¶

conditions¶

conditions: tuple[Condition, ...] = field(factory=tuple, converter=tuple)

contrast_degree¶

contrast_degree: int | None = field(default=None, validator=is_optional_positive_int)

contrast_expr¶

contrast_expr: ContrastExpr | None = field(default=None)

contrast_level_ordering¶

contrast_level_ordering: tuple[str, ...] | None = field(default=None, validator=is_optional_tuple_of_str)

contrast_ref¶

contrast_ref: str | None = field(default=None, validator=is_optional_str)

contrast_type¶

contrast_type: str | None = field(default=None, validator=(validators.optional(is_choice_str(('pairwise', 'sequential', 'poly', 'treatment', 'sum', 'helmert', 'custom')))))

focal_at_quantile¶

focal_at_quantile: int | None = field(default=None, validator=is_optional_positive_int)

focal_at_range¶

focal_at_range: int | None = field(default=None, validator=is_optional_positive_int)

focal_at_values¶

focal_at_values: tuple[float | str, ...] | None = field(default=None)

focal_var¶

focal_var: str = field(validator=(validators.instance_of(str)))

has_conditions¶

has_conditions: bool

Return True if conditioning variables are specified.

has_contrast¶

has_contrast: bool

Return True if any contrast is specified (named function or bracket expr).

has_contrast_expr¶

has_contrast_expr: bool

Return True if a bracket contrast expression is specified.

has_rhs_contrasts¶

has_rhs_contrasts: bool

Return True if any RHS condition has a bracket contrast expression.

ResolvedConditions¶

Typed buckets for resolved conditioning specifications.

Attributes:

Attributes¶

at_overrides¶

at_overrides: dict[str, float]

grid_categoricals¶

grid_categoricals: dict[str, list[str]]

grid_numerics¶

grid_numerics: dict[str, list[float]]

has_grid¶

has_grid: bool

Return True if any condition requires grid expansion.

raw_at_overrides¶

raw_at_overrides: dict[str, float] = Factory(dict)

Pre-transform at-overrides with original (raw) variable names.

Populated by the dispatch layer after _resolve_at_overrides remaps keys to design-matrix names. The slopes path needs raw keys because it operates on a data grid with raw column names, not design-matrix names. Empty when no transform resolution was applied.

set_categoricals¶

set_categoricals: dict[str, str]

Functions¶

apply_bracket_contrasts¶

apply_bracket_contrasts(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply bracket contrast expression to an EMM MeeState.

Computes contrasts by building a weight matrix from the bracket expression and multiplying by the EMM estimates.

Parameters:

Returns:

apply_bracket_contrasts_grouped¶

apply_bracket_contrasts_grouped(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply bracket contrasts within each condition group of a crossed MeeState.

Similar to apply_contrasts_grouped() but for bracket contrast expressions. Applies the same contrast matrix independently within each group of n_focal consecutive rows.

Parameters:

Returns:

apply_contrasts¶

apply_contrasts(mee_state: MeeState, contrast_type: str, fit: FitState | None = None, *, degree: int | None = None, ref_idx: int | None = None, level_ordering: tuple[str, ...] | None = None) -> MeeState

Apply contrast matrix to marginal means/effects.

High-level function that takes a MeeState with EMM estimates and returns a new MeeState with contrast estimates.

Parameters:

Returns:

Examples:

>>> mee = compute_emm(bundle, fit, "treatment", "treatment")
>>> contrasts = apply_contrasts(mee, "pairwise")
>>> contrasts.estimate  # B-A, C-A, C-B differences

Note: Variance propagation for inference is deferred to .infer(). This function only computes point estimates.

apply_contrasts_grouped¶

apply_contrasts_grouped(mee_state: MeeState, contrast_type: str | None, *, degree: int | None = None, ref_idx: int | None = None, level_ordering: tuple[str, ...] | None = None) -> MeeState

Apply contrasts within each condition group of a crossed MeeState.

When EMMs have been computed over a crossed grid (focal levels x condition groups), this function applies the contrast matrix independently within each group of n_focal consecutive rows, then stacks the results.

The number of focal levels is inferred from the focal_var column of mee_state.grid.

Parameters:

Returns:

apply_rhs_bracket_contrast¶

apply_rhs_bracket_contrast(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply a bracket contrast on a RHS condition column.

After the main computation produces a crossed MeeState with a condition column (e.g., Dose with levels [High, Low]), this collapses that column by applying the bracket contrast.

The condition column must vary slowest in the grid (standard layout from crossed EMM/slope computation).

Parameters:

Returns:

build_all_pairwise_matrix¶

build_all_pairwise_matrix(n_levels: int) -> np.ndarray

Build all pairwise contrasts between EMM levels.

Creates C(n, 2) = n*(n-1)/2 contrasts for all unique pairs. Unlike build_pairwise_matrix(), this includes all pairs, not just comparisons to the reference level.

Parameters:

Returns:

Examples:

>>> C = build_all_pairwise_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.],
       [ 0., -1.,  1.]])

build_bracket_contrast_matrix¶

build_bracket_contrast_matrix(expr: ContrastExpr, levels: list[str]) -> tuple[np.ndarray, list[str]]

Build contrast matrix and labels from bracket contrast expression.

Handles auto-normalization (each operand side sums to +1 or -1) and wildcard expansion (* expands to all unmentioned levels).

Parameters:

Returns:

build_contrast_matrix¶

build_contrast_matrix(contrast_type: str | dict, levels: list, normalize: bool = False) -> np.ndarray

Build a contrast matrix based on contrast type.

Dispatcher function that builds the appropriate contrast matrix based on the contrast type specification.

Parameters:

Returns:

Examples:

>>> build_contrast_matrix("pairwise", ["A", "B", "C"])
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.],
       [ 0., -1.,  1.]])

build_helmert_matrix¶

build_helmert_matrix(n_levels: int) -> np.ndarray

Build Helmert contrasts (each level vs mean of previous levels).

Creates n_levels - 1 contrasts where level k is compared to the mean of levels 0, 1, ..., k-1.

Parameters:

Returns:

Examples:

>>> C = build_helmert_matrix(3)
>>> C
array([[-1. ,  1. ,  0. ],
       [-0.5, -0.5,  1. ]])

>>> C = build_helmert_matrix(4)
>>> C
array([[-1.        ,  1.        ,  0.        ,  0.        ],
       [-0.5       , -0.5       ,  1.        ,  0.        ],
       [-0.33333333, -0.33333333, -0.33333333,  1.        ]])

build_pairwise_matrix¶

build_pairwise_matrix(n_levels: int) -> np.ndarray

Build (n-1) linearly independent pairwise contrasts.

Creates “treatment-style” contrasts comparing each level to the first (reference) level. This produces n_levels - 1 contrasts that span the space of all pairwise differences.

Parameters:

Returns:

Examples:

>>> C = build_pairwise_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.]])

build_poly_matrix¶

build_poly_matrix(n_levels: int, degree: int | None = None) -> np.ndarray

Build orthogonal polynomial contrast matrix for EMMs.

Creates orthogonal polynomial contrasts for ordered factors. Tests for linear, quadratic, cubic (etc.) trends across the factor levels.

Parameters:

Returns:

Examples:

>>> build_poly_matrix(5)  # 4x5 matrix
>>> build_poly_matrix(5, degree=2)  # 2x5 matrix (linear + quadratic)

build_reference_grid¶

build_reference_grid(bundle: DataBundle, focal_vars: list[str], *, at: dict[str, Any] | None = None, covariate_means: dict[str, float] | None = None) -> pl.DataFrame

Construct reference grid for marginal effects evaluation.

Creates a grid of covariate values at which to evaluate predictions. This function transforms the parsed explore formula’s conditions into a concrete evaluation grid.

The grid construction follows emmeans conventions:

• Focal variables: vary across their unique levels/values from data

• Non-focal continuous: set to their overall mean from the data

• Non-focal categorical: create rows for all levels (equal weight averaging)

• Conditioned variables (@ spec): use the specified values

Parameters:

Returns:

Integration with explore(): Called from model._compute_marginal_effects() after parsing::

    # From parsed explore formula "treatment ~ age@50"
    parsed = ExploreFormulaSpec(
        focal_var="treatment",
        conditions=[Condition(var="age", at_values=(50,))]
    )

    # Convert conditions to at dict
    at_dict = {"age": 50.0}

    # Build grid
    grid = build_reference_grid(
        bundle=self._bundle,
        focal_vars=["treatment"],
        at=at_dict,
    )
    # grid: pl.DataFrame with columns ["treatment"] and rows ["A", "B", "C"]
    # age is NOT in the grid (it's held fixed at 50)

Examples:

Grid for categorical focal::

grid = build_reference_grid(bundle, ["treatment"])
# Returns grid with one row per treatment level

Grid with conditioning::

grid = build_reference_grid(bundle, ["treatment"], at={"age": 50})
# Returns grid at age=50

Multiple focal variables (interaction EMMs)::

grid = build_reference_grid(bundle, ["treatment", "sex"])
# Returns Cartesian product: treatment x sex levels

build_sequential_matrix¶

build_sequential_matrix(n_levels: int) -> np.ndarray

Build sequential (successive differences) contrasts.

Creates n_levels - 1 contrasts comparing each level to the previous one.

Parameters:

Returns:

Examples:

>>> C = build_sequential_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [ 0., -1.,  1.]])

>>> C = build_sequential_matrix(4)
>>> C
array([[-1.,  1.,  0.,  0.],
       [ 0., -1.,  1.,  0.],
       [ 0.,  0., -1.,  1.]])

build_sum_to_zero_matrix¶

build_sum_to_zero_matrix(n_levels: int) -> np.ndarray

Build sum-to-zero contrasts (deviation coding).

Creates contrasts comparing each level to the grand mean.

Parameters:

Returns:

Examples:

>>> C = build_sum_to_zero_matrix(3)
>>> C
array([[ 0.667, -0.333, -0.333],
       [-0.333,  0.667, -0.333]])

build_treatment_matrix¶

build_treatment_matrix(n_levels: int, ref_idx: int = 0) -> np.ndarray

Build treatment (Dunnett-style) contrasts against a reference level.

Creates n_levels - 1 contrasts, each comparing one non-reference level to the specified reference level.

Parameters:

Returns:

Examples:

>>> C = build_treatment_matrix(3, ref_idx=0)
>>> C
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.]])

>>> C = build_treatment_matrix(3, ref_idx=2)
>>> C
array([[ 1.,  0., -1.],
       [ 0.,  1., -1.]])

combine_resolved¶

combine_resolved(a: ResolvedConditions, b: ResolvedConditions) -> ResolvedConditions

Merge two ResolvedConditions, with b taking precedence on conflicts.

Parameters:

Returns:

compose_contrast_matrix¶

compose_contrast_matrix(C: np.ndarray, X_ref: np.ndarray) -> np.ndarray

Compose contrast matrix with prediction matrix.

L_emm @ beta = C @ (X_ref @ beta) = C @ EMMs L_emm @ beta = C @ (X_ref @ beta) = C @ EMMs

Parameters:

Returns:

compute_compound_bracket_contrasts¶

compute_compound_bracket_contrasts(bundle: object, fit: object, focal_var: str, contrast_expr: ContrastExpr, *, data: pl.DataFrame, spec: object | None = None, effect_scale: str = 'link', resolved: object | None = None) -> MeeState

Compute bracket contrasts for a compound focal variable.

Handles compound variables like Drug:Dose by building a crossed EMM grid over the component variables, creating compound level names, and then applying the bracket contrast.

For example, Drug:Dose[Active:High - Placebo:Low] builds EMMs over Drug × Dose, labels each cell as Active:High, Active:Low, etc., and applies the contrast Active:High - Placebo:Low.

Parameters:

Returns:

compute_conditional_emm¶

compute_conditional_emm(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, spec: object, varying_offsets: object, grouping_var: str, effect_scale: str = 'link', levels: list[str] | None = None, at_overrides: dict[str, float] | None = None, set_categoricals: dict[str, str] | None = None) -> MeeState

Compute per-group conditional EMMs incorporating intercept BLUPs.

For each group g and each focal level, the conditional EMM is: η_g = X_ref_row @ β + b_intercept_g (+ other BLUP contributions)

When effect_scale=“response”: estimates = g⁻¹(η_g).

Parameters:

Returns:

compute_conditional_slopes¶

compute_conditional_slopes(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, spec: object, varying_offsets: object, grouping_var: str, effect_scale: str = 'link') -> MeeState

Compute per-group conditional slopes incorporating BLUPs.

For each group, the conditional slope is: slope_g = β_x + b_x_g (random slopes model) slope_g = β_x (random intercepts only)

When effect_scale=“response” and the link is non-identity, the response-scale slope is: slope_response_g = slope_link_g × dμ/dη(η_g)

where η_g is the linear predictor at the reference point for group g.

Parameters:

Returns:

compute_contrasts¶

compute_contrasts(emm: np.ndarray, contrast_matrix: np.ndarray) -> np.ndarray

Apply contrast matrix to EMMs.

Transforms a vector of estimated marginal means into contrast estimates by matrix multiplication. This is the final step when the explore formula includes a contrast function like pairwise(treatment).

contrasts = C @ EMM contrasts = C @ EMM

Where:

• C is the contrast matrix, shape (n_contrasts, n_levels)

• EMM is the vector of marginal means, shape (n_levels,)

• contrasts is the result, shape (n_contrasts,)

• Each row sums to 0 (contrasts compare, not estimate absolute levels)

• Each row sums to 0 (contrasts compare, not estimate absolute levels)

• The result represents a comparison (e.g., B - A)

Parameters:

Returns:

Integration with explore(): Called after compute_emm when a contrast function is specified::

    # In model._compute_marginal_effects():
    if parsed.has_contrast:
        # First compute EMMs
        emms = compute_emm(...)

        # Build appropriate contrast matrix
        if parsed.contrast_type == "pairwise":
            C = build_all_pairwise_matrix(n_levels)
        elif parsed.contrast_type == "sequential":
            C = build_sequential_matrix(n_levels)
        elif parsed.contrast_type == "poly":
            C = build_poly_matrix(n_levels, degree=parsed.contrast_degree)

        # Apply contrasts
        contrast_estimates = compute_contrasts(emms, C)

        return build_mee_state(
            grid=contrast_grid,  # with contrast labels
            estimate=contrast_estimates,
            mee_type="contrasts",
            ...
        )

Note: For inference on contrasts, the variance of contrasts is: Var(C @ EMM) = C @ Var(EMM) @ C.T

This is handled by the inference methods (delta method for asymp, direct computation for bootstrap).

Examples:

Apply pairwise contrasts to 3-level EMMs::

emm = np.array([2.0, 3.5, 2.8])  # A, B, C means
C = build_all_pairwise_matrix(3)
# C = [[-1, 1, 0],   # B - A
#      [-1, 0, 1],   # C - A
#      [0, -1, 1]]   # C - B

contrasts = compute_contrasts(emm, C)
# contrasts: array([1.5, 0.8, -0.7])

compute_emm¶

compute_emm(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, levels: list[str] | None = None, at_overrides: dict[str, float] | None = None, set_categoricals: dict[str, str] | None = None, spec: object | None = None, how: str = 'mem', effect_scale: str = 'link') -> MeeState

Compute estimated marginal means for a categorical focal variable.

Supports two averaging methods via the how parameter:

• "mem" (default): Marginal Estimated Mean. Balanced reference grid (emmeans-style). Predictions at a grid where covariates are at their means.

• "ame": Average Marginal Effect / g-computation. 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ᵢ) · β).

Parameters:

Returns:

Examples:

Compute EMMs with default (reference grid) averaging::

mee = compute_emm(bundle, fit, "treatment", "treatment")

G-computation averaging for a logistic regression::

mee = compute_emm(bundle, fit, "treatment", "treatment",
                  spec=spec, how="ame", effect_scale="response")

Note: For how="ame" with effect_scale="response" 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 how="mem".

compute_joint_test¶

compute_joint_test(fit: FitState, bundle: DataBundle, spec: ModelSpec, terms: list[str] | None = None, *, errors: str = 'iid', data: pl.DataFrame | None = None) -> JointTestState

Compute joint hypothesis tests for model terms.

Tests each model term (continuous variables, factors, and interactions) using direct coefficient tests. This matches R’s emmeans::joint_tests() behavior.

For each term:

• Continuous variables: F-test on the single coefficient (df1=1)

• Categorical factors: Joint F-test on all indicator coefficients

• Interactions: Joint F-test on all interaction coefficients

The test type is determined by the model family:

• Gaussian (LM/LMER): F-test with df_resid

• Non-Gaussian (GLM/GLMER): Chi-square test (asymptotic)

Parameters:

Returns:

Examples:

Basic usage::

from marginal import compute_joint_test
state = compute_joint_test(fit, bundle, spec)
# state.terms: ("x", "group", "x:group")
# state.statistic: F or chi2 values per term
# state.p_value: p-values per term

Welch ANOVA::

state = compute_joint_test(fit, bundle, spec, errors="unequal_var", data=df)
# Per-term Satterthwaite df2 in state.df2

Test specific terms only::

state = compute_joint_test(fit, bundle, spec, terms=["group"])

Note: Intercept is never tested. Random effect grouping factors are excluded for mixed models.

compute_mee_inference¶

compute_mee_inference(mee: MeeState, vcov: np.ndarray, df_resid: float | np.ndarray | None, conf_level: float = 0.95, null: float = 0.0, alternative: str = 'two-sided') -> MeeState

Compute delta method inference for marginal effects.

SE = sqrt(diag(L @ vcov @ L’)) SE = sqrt(diag(L @ vcov @ L’))

where L is the design matrix (X_ref for EMMs, selector for slopes).

Applies multiplicity adjustment for contrast CIs to match R’s emmeans:

• Tukey HSD for pairwise contrasts.

• Multivariate-t (MVT) for sequential, treatment, sum, helmert contrasts.

• No adjustment for polynomial contrasts or non-contrast MEEs.

For response-scale EMMs (effect_scale="response"), CIs are computed on the link scale then back-transformed via the inverse link function, matching R’s emmeans behavior.

Parameters:

Returns:

Examples:

Compute inference for EMMs::

from marginal import compute_emm
from inference import compute_mee_inference

mee = compute_emm(bundle, fit, "treatment", "treatment")
mee_with_inf = compute_mee_inference(mee, fit.vcov, fit.df_resid)
# mee_with_inf.se: standard errors
# mee_with_inf.ci_lower, mee_with_inf.ci_upper: confidence bounds

Note: For EMMs, L_matrix is the reference design matrix X_ref. For slopes, L_matrix is a selector row that picks the coefficient. For contrasts, L_matrix is the contrast matrix applied to EMMs.

compute_mee_inference_fallback¶

compute_mee_inference_fallback(mee: 'MeeState', bundle: 'DataBundle', fit: 'FitState', data: 'pl.DataFrame', df_resid: float | None, conf_level: float = 0.95, null: float = 0.0, alternative: str = 'two-sided') -> 'MeeState'

Compute inference for MEE without L_matrix (fallback path).

Used when L_matrix is not available (legacy MEE). Computes simplified SEs via compute_mee_se, then builds test statistics, p-values, and confidence intervals.

Parameters:

Returns:

compute_mee_se¶

compute_mee_se(mee: 'MeeState', bundle: 'DataBundle', fit: 'FitState', data: 'pl.DataFrame') -> np.ndarray

Compute standard errors for MEE estimates (means or slopes).

When an L_matrix is available, uses the delta method: SE = sqrt(diag(L @ vcov @ L')). This correctly handles multi-row L_matrices (e.g. per-group conditional slopes).

Otherwise falls back to simplified type-based dispatch:

• "means": SE = sqrt(MSE / n_per_group) for each group level

• "slopes": SE = sqrt(vcov[idx, idx]) from coefficient vcov diagonal

Parameters:

Returns:

compute_slopes¶

compute_slopes(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, spec: object | None = None, effect_scale: str = 'link') -> MeeState

Compute marginal slope for a continuous focal variable.

For a linear model, this extracts the coefficient for the variable. For models with interactions, this would need to average marginal effects across the grid (future enhancement).

Marginal slopes answer: “By how much does the predicted outcome change for a one-unit increase in this variable?”

Parameters:

Returns:

Examples:

Compute marginal effect of age::

from marginal import compute_slopes
mee = compute_slopes(bundle, fit, "age", "age")
# For y ~ age: returns MeeState with estimate = [coef_age]

Note: This is the fast path for Gaussian identity-link models without interactions. For models with interactions or GLMs (non-identity link), the dispatcher routes to compute_slopes_finite_diff which handles both cases via centered finite differences.

compute_slopes_crossed¶

compute_slopes_crossed(bundle: 'DataBundle', fit: 'FitState', focal_var: str, *, resolved: object, data: pl.DataFrame, spec: 'ModelSpec | None' = None, formula_spec: 'FormulaSpec | None' = None, effect_scale: str = 'link', delta_frac: float = 0.001) -> MeeState

Compute crossed slopes over focal variable x condition grid.

Builds a Cartesian product of condition levels (grid categoricals, grid numerics, scalar pins) and computes a finite-difference AME for each combination. Analogous to _compute_emm_crossed but for continuous focal variables.

Parameters:

Returns:

compute_slopes_finite_diff¶

compute_slopes_finite_diff(bundle: DataBundle, fit: FitState, focal_var: str, explore_formula: str, *, spec: ModelSpec, formula_spec: FormulaSpec, data: pl.DataFrame, how: str = 'mem', effect_scale: str = 'link', delta_frac: float = 0.001) -> MeeState

Compute marginal slopes via centered finite differences.

Supports two averaging methods via the how parameter:

• "mem": Average over a balanced reference grid (Cartesian product of categorical levels, continuous covariates at means). Matches R’s emmeans::emtrends.

• "ame": Average over actual data rows. Gives a true Average Marginal Effect (AME) that preserves the observed covariate distribution.

For linear models (identity link), both approaches give identical results.

1. delta = delta_frac × range(focal_var)

2. delta = delta_frac × range(focal_var)

3. Build evaluation grid (balanced or observed data)

4. Perturb: grid_plus = grid[focal + delta/2], grid_minus = grid[focal − delta/2]

5. X_plus, X_minus = evaluate_newdata(formula_spec, grid_*)

6. L_diff = (X_plus − X_minus) / delta (link-scale functional)

7. If effect_scale="response" and how="ame": J_i = [f'(η_i+) X_i+ − f'(η_i−) X_i−] / delta (per-observation Jacobian with f’ at perturbed points)

8. L_avg = mean(J, axis=0, keepdims=True)

Parameters:

Returns:

dispatch_marginal_computation¶

dispatch_marginal_computation(parsed: ExploreFormulaSpec, bundle: DataBundle, fit: FitState, data: pl.DataFrame, *, spec: ModelSpec | None = None, formula_spec: object | None = None, varying_offsets: VaryingState | None = None, effect_scale: str = 'link', varying: str = 'exclude', how: str = 'auto', inverse_transforms: bool = True, by: str | None = None) -> MeeState

Route a parsed explore formula to the appropriate marginal computation.

Validates the focal variable, determines whether it is categorical or continuous, and dispatches to the correct computation function.

Supports effect_scale= and varying= for scale transforms and conditional (group-specific) effects in mixed models. When RHS conditions contain bracket contrasts (e.g., Drug[A - B] ~ Dose[High - Low]), applies them as a post-processing step after the main computation.

Parameters:

Returns:

get_contrast_labels¶

get_contrast_labels(levels: list[str], contrast_type: str = 'pairwise') -> list[str]

Generate human-readable labels for contrasts.

Parameters:

Returns:

Examples:

>>> get_contrast_labels(["A", "B", "C"], "pairwise")
['B - A', 'C - A']
>>> get_contrast_labels(["A", "B", "C"], "all_pairwise")
['B - A', 'C - A', 'C - B']

parse_explore_formula¶

parse_explore_formula(formula: str, model_terms: list[str] | None = None) -> ExploreFormulaSpec

Parse an explore formula string.

Parameters:

Returns:

Examples:

Simple term::

>>> parse_explore_formula("treatment")
ExploreFormulaSpec(focal_var='treatment', contrast_type=None, conditions=())

Contrast function::

>>> parse_explore_formula("pairwise(treatment)")
ExploreFormulaSpec(focal_var='treatment', contrast_type='pairwise', conditions=())

With conditioning::

>>> parse_explore_formula("treatment ~ age@50")
ExploreFormulaSpec(
    focal_var='treatment',
    contrast_type=None,
    conditions=(Condition(var='age', at_values=(50.0,)),)
)

resolve_conditions¶

resolve_conditions(conditions: tuple[Condition, ...], bundle: DataBundle, data: pl.DataFrame) -> ResolvedConditions

Classify each Condition into the appropriate typed bucket.

Uses bundle.factor_levels to distinguish categorical from continuous variables, and the data DataFrame to compute range / quantile grids for continuous variables.

Parameters:

Returns:

Modules¶

bracket_contrasts¶

Bracket contrast matrix builder and application.

Builds contrast weight matrices from bracket contrast expressions (e.g., Drug[Active - Placebo], Drug[* - Placebo], Drug[(A + B) - C]).

Auto-normalization: each side of - normalizes so weights sum to 1. Wildcard expansion: * expands to all levels not mentioned on the other side, producing one contrast row per expanded level.

Functions:

Classes¶

Functions¶

apply_bracket_contrasts¶

apply_bracket_contrasts(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply bracket contrast expression to an EMM MeeState.

Computes contrasts by building a weight matrix from the bracket expression and multiplying by the EMM estimates.

Parameters:

Returns:

apply_bracket_contrasts_grouped¶

apply_bracket_contrasts_grouped(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply bracket contrasts within each condition group of a crossed MeeState.

Similar to apply_contrasts_grouped() but for bracket contrast expressions. Applies the same contrast matrix independently within each group of n_focal consecutive rows.

Parameters:

Returns:

apply_rhs_bracket_contrast¶

apply_rhs_bracket_contrast(mee_state: MeeState, expr: ContrastExpr) -> MeeState

Apply a bracket contrast on a RHS condition column.

After the main computation produces a crossed MeeState with a condition column (e.g., Dose with levels [High, Low]), this collapses that column by applying the bracket contrast.

The condition column must vary slowest in the grid (standard layout from crossed EMM/slope computation).

Parameters:

Returns:

build_bracket_contrast_matrix¶

build_bracket_contrast_matrix(expr: ContrastExpr, levels: list[str]) -> tuple[np.ndarray, list[str]]

Build contrast matrix and labels from bracket contrast expression.

Handles auto-normalization (each operand side sums to +1 or -1) and wildcard expansion (* expands to all unmentioned levels).

Parameters:

Returns:

compute_compound_bracket_contrasts¶

compute_compound_bracket_contrasts(bundle: object, fit: object, focal_var: str, contrast_expr: ContrastExpr, *, data: pl.DataFrame, spec: object | None = None, effect_scale: str = 'link', resolved: object | None = None) -> MeeState

Compute bracket contrasts for a compound focal variable.

Handles compound variables like Drug:Dose by building a crossed EMM grid over the component variables, creating compound level names, and then applying the bracket contrast.

For example, Drug:Dose[Active:High - Placebo:Low] builds EMMs over Drug × Dose, labels each cell as Active:High, Active:Low, etc., and applies the contrast Active:High - Placebo:Low.

Parameters:

Returns:

dispatch_bracket_contrasts¶

dispatch_bracket_contrasts(bundle: DataBundle, fit: FitState, focal_var: str, contrast_expr: ContrastExpr, *, data: pl.DataFrame, spec: ModelSpec | None = None, effect_scale: str = 'link', how: str = 'mem', resolved: ResolvedConditions | None = None) -> MeeState

Compute bracket contrast expression for a categorical focal variable.

First computes EMMs, then applies the bracket contrast matrix. When resolved contains grid conditions, computes crossed EMMs and applies grouped bracket contrasts.

Parameters:

Returns:

compute¶

Marginal effects dispatch and routing.

Routes parsed explore formulas to the appropriate computation: EMMs for categorical focal variables, slopes for continuous, or contrasts. Supports effect_scale= (link/response scale), how= (mem/ame), and varying= (marginal/conditional).

Functions:

Classes¶

Functions¶

compute_emm_categorical¶

compute_emm_categorical(bundle: DataBundle, fit: FitState, focal_var: str, *, levels: list[str] | None = None, at_overrides: dict[str, float] | None = None, set_categoricals: dict[str, str] | None = None, spec: ModelSpec | None = None, effect_scale: str = 'link', how: str = 'mem') -> MeeState

Compute EMMs for a categorical focal variable.

Parameters:

Returns:

dispatch_marginal_computation¶

dispatch_marginal_computation(parsed: ExploreFormulaSpec, bundle: DataBundle, fit: FitState, data: pl.DataFrame, *, spec: ModelSpec | None = None, formula_spec: object | None = None, varying_offsets: VaryingState | None = None, effect_scale: str = 'link', varying: str = 'exclude', how: str = 'auto', inverse_transforms: bool = True, by: str | None = None) -> MeeState

Route a parsed explore formula to the appropriate marginal computation.

Validates the focal variable, determines whether it is categorical or continuous, and dispatches to the correct computation function.

Supports effect_scale= and varying= for scale transforms and conditional (group-specific) effects in mixed models. When RHS conditions contain bracket contrasts (e.g., Drug[A - B] ~ Dose[High - Low]), applies them as a post-processing step after the main computation.

Parameters:

Returns:

conditions¶

Condition resolution for explore formula RHS conditioning.

Resolves parsed Condition objects into typed buckets that downstream EMM / slope / contrast computation can consume:

• at_overrides: single numeric pins (Income@50)

• set_categoricals: single categorical pins (Ethnicity@Asian)

• grid_categoricals: multi-level categoricals to cross (bare Ethnicity)

• grid_numerics: multi-value numerics to cross (Income@:range(5))

Classes:

Functions:

Classes¶

ResolvedConditions¶

Typed buckets for resolved conditioning specifications.

Attributes:

Attributes¶

at_overrides¶

at_overrides: dict[str, float]

grid_categoricals¶

grid_categoricals: dict[str, list[str]]

grid_numerics¶

grid_numerics: dict[str, list[float]]

has_grid¶

has_grid: bool

Return True if any condition requires grid expansion.

raw_at_overrides¶

raw_at_overrides: dict[str, float] = Factory(dict)

Pre-transform at-overrides with original (raw) variable names.

Populated by the dispatch layer after _resolve_at_overrides remaps keys to design-matrix names. The slopes path needs raw keys because it operates on a data grid with raw column names, not design-matrix names. Empty when no transform resolution was applied.

set_categoricals¶

set_categoricals: dict[str, str]

Functions¶

combine_resolved¶

combine_resolved(a: ResolvedConditions, b: ResolvedConditions) -> ResolvedConditions

Merge two ResolvedConditions, with b taking precedence on conflicts.

Parameters:

Returns:

get_column_values¶

get_column_values(data: pl.DataFrame, bundle: DataBundle, var: str) -> np.ndarray

Extract a data column as a numpy array for range/quantile computation.

Looks in the original data DataFrame first, falling back to the design matrix column if the variable name matches an X_names entry.

Parameters:

Returns:

resolve_conditions¶

resolve_conditions(conditions: tuple[Condition, ...], bundle: DataBundle, data: pl.DataFrame) -> ResolvedConditions

Classify each Condition into the appropriate typed bucket.

Uses bundle.factor_levels to distinguish categorical from continuous variables, and the data DataFrame to compute range / quantile grids for continuous variables.

Parameters:

Returns:

contrasts¶

Contrast computation for marginal effects.

Applies contrast matrices to EMM estimates. Matrix builders live in marginal.matrices.

apply_contrasts: Apply contrasts to MeeState apply_contrasts: Apply contrasts to MeeState apply_contrasts_grouped: Apply contrasts within condition groups compute_contrasts: Apply contrast matrix to EMM vector

Functions:

Classes¶

Functions¶

apply_contrasts¶

apply_contrasts(mee_state: MeeState, contrast_type: str, fit: FitState | None = None, *, degree: int | None = None, ref_idx: int | None = None, level_ordering: tuple[str, ...] | None = None) -> MeeState

Apply contrast matrix to marginal means/effects.

High-level function that takes a MeeState with EMM estimates and returns a new MeeState with contrast estimates.

Parameters:

Returns:

Examples:

>>> mee = compute_emm(bundle, fit, "treatment", "treatment")
>>> contrasts = apply_contrasts(mee, "pairwise")
>>> contrasts.estimate  # B-A, C-A, C-B differences

Note: Variance propagation for inference is deferred to .infer(). This function only computes point estimates.

apply_contrasts_grouped¶

apply_contrasts_grouped(mee_state: MeeState, contrast_type: str | None, *, degree: int | None = None, ref_idx: int | None = None, level_ordering: tuple[str, ...] | None = None) -> MeeState

Apply contrasts within each condition group of a crossed MeeState.

When EMMs have been computed over a crossed grid (focal levels x condition groups), this function applies the contrast matrix independently within each group of n_focal consecutive rows, then stacks the results.

The number of focal levels is inferred from the focal_var column of mee_state.grid.

Parameters:

Returns:

compute_contrasts¶

compute_contrasts(emm: np.ndarray, contrast_matrix: np.ndarray) -> np.ndarray

Apply contrast matrix to EMMs.

Transforms a vector of estimated marginal means into contrast estimates by matrix multiplication. This is the final step when the explore formula includes a contrast function like pairwise(treatment).

contrasts = C @ EMM contrasts = C @ EMM

Where:

• C is the contrast matrix, shape (n_contrasts, n_levels)

• EMM is the vector of marginal means, shape (n_levels,)

• contrasts is the result, shape (n_contrasts,)

• Each row sums to 0 (contrasts compare, not estimate absolute levels)

• Each row sums to 0 (contrasts compare, not estimate absolute levels)

• The result represents a comparison (e.g., B - A)

Parameters:

Returns:

Integration with explore(): Called after compute_emm when a contrast function is specified::

    # In model._compute_marginal_effects():
    if parsed.has_contrast:
        # First compute EMMs
        emms = compute_emm(...)

        # Build appropriate contrast matrix
        if parsed.contrast_type == "pairwise":
            C = build_all_pairwise_matrix(n_levels)
        elif parsed.contrast_type == "sequential":
            C = build_sequential_matrix(n_levels)
        elif parsed.contrast_type == "poly":
            C = build_poly_matrix(n_levels, degree=parsed.contrast_degree)

        # Apply contrasts
        contrast_estimates = compute_contrasts(emms, C)

        return build_mee_state(
            grid=contrast_grid,  # with contrast labels
            estimate=contrast_estimates,
            mee_type="contrasts",
            ...
        )

Note: For inference on contrasts, the variance of contrasts is: Var(C @ EMM) = C @ Var(EMM) @ C.T

This is handled by the inference methods (delta method for asymp, direct computation for bootstrap).

Examples:

Apply pairwise contrasts to 3-level EMMs::

emm = np.array([2.0, 3.5, 2.8])  # A, B, C means
C = build_all_pairwise_matrix(3)
# C = [[-1, 1, 0],   # B - A
#      [-1, 0, 1],   # C - A
#      [0, -1, 1]]   # C - B

contrasts = compute_contrasts(emm, C)
# contrasts: array([1.5, 0.8, -0.7])

dispatch_contrasts¶

dispatch_contrasts(bundle: DataBundle, fit: FitState, focal_var: str, contrast_type: str | None, contrast_degree: int | None, data: pl.DataFrame, *, spec: ModelSpec | None = None, effect_scale: str = 'link', how: str = 'mem', resolved: ResolvedConditions | None = None, focal_at_values: tuple[float | str, ...] | None = None, contrast_ref: str | None = None, level_ordering: tuple[str, ...] | None = None) -> MeeState

Compute contrasts for a categorical focal variable.

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 over the requested subset of levels.

Parameters:

Returns:

emm¶

Estimated marginal means computation.

This module provides EMM computation for categorical focal variables.

compute_emm: Compute estimated marginal means at grid points compute_emm: Compute estimated marginal means at grid points

Functions:

Classes¶

Functions¶

compute_conditional_emm¶

compute_conditional_emm(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, spec: object, varying_offsets: object, grouping_var: str, effect_scale: str = 'link', levels: list[str] | None = None, at_overrides: dict[str, float] | None = None, set_categoricals: dict[str, str] | None = None) -> MeeState

Compute per-group conditional EMMs incorporating intercept BLUPs.

For each group g and each focal level, the conditional EMM is: η_g = X_ref_row @ β + b_intercept_g (+ other BLUP contributions)

When effect_scale=“response”: estimates = g⁻¹(η_g).

Parameters:

Returns:

compute_emm¶

compute_emm(bundle: 'DataBundle', fit: 'FitState', focal_var: str, explore_formula: str, *, levels: list[str] | None = None, at_overrides: dict[str, float] | None = None, set_categoricals: dict[str, str] | None = None, spec: object | None = None, how: str = 'mem', effect_scale: str = 'link') -> MeeState

Compute estimated marginal means for a categorical focal variable.

Supports two averaging methods via the how parameter:

• "mem" (default): Marginal Estimated Mean. Balanced reference grid (emmeans-style). Predictions at a grid where covariates are at their means.

• "ame": Average Marginal Effect / g-computation. 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ᵢ) · β).

Parameters:

Returns:

Examples:

Compute EMMs with default (reference grid) averaging::

mee = compute_emm(bundle, fit, "treatment", "treatment")

G-computation averaging for a logistic regression::

mee = compute_emm(bundle, fit, "treatment", "treatment",
                  spec=spec, how="ame", effect_scale="response")

Note: For how="ame" with effect_scale="response" 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 how="mem".

compute_emm_crossed¶

compute_emm_crossed(bundle: DataBundle, fit: FitState, focal_var: str, *, resolved: ResolvedConditions, levels: list[str] | None = None, spec: ModelSpec | None = None, effect_scale: str = 'link', how: str = 'mem') -> MeeState

Compute crossed EMMs over focal levels x condition grid.

Builds a Cartesian product of focal levels with all grid conditions, computing one EMM row per combination.

Supports two averaging methods via how:

• "mem": Reference grid at covariate means (emmeans-style).

• "ame": G-computation over observed data rows with counterfactual treatment assignment per combo.

Parameters:

Returns:

explore¶

Explore formula parser.

Tokenizes and parses explore formula strings into :class:ExploreFormulaSpec containers using :class:ExploreScanner and :class:ExploreParser.

Grammar::

explore_formula := lhs [ '~' rhs ]

lhs := focal_term | contrast_fn '(' focal_term { ',' arg } ')'
focal_term := varname [ '@' at_spec ]
at_spec := value | '[' value_list ']' | [':']'range(' n ')' | [':']'quantile(' n ')'
contrast_fn := 'pairwise' | 'sequential' | 'poly'
             | 'treatment' | 'dummy' | 'sum' | 'deviation' | 'helmert'
arg := IDENTIFIER '=' value | value

rhs := condition [ '+' condition ]*
condition := term [ '@' at_spec ]

Classes:

Functions:

Classes¶

Condition¶

A conditioning specification in explore formula.

Represents a variable to condition on, optionally with specific values, a method for generating values, or a bracket contrast expression.

Created by: parse_explore_formula(), ExploreParser Consumed by: resolve_all_conditions(), dispatch_marginal_computation() Augmented by: Never

Attributes:

Attributes¶

at_quantile¶

at_quantile: int | None = field(default=None, validator=is_optional_positive_int)

at_range¶

at_range: int | None = field(default=None, validator=is_optional_positive_int)

at_values¶

at_values: tuple | None = field(default=None)

contrast_expr¶

contrast_expr: ContrastExpr | None = field(default=None, validator=(validators.optional(validators.instance_of(ContrastExpr))))

var¶

var: str = field(validator=(validators.instance_of(str)))

ExploreFormulaError¶

ExploreFormulaError(message: str, formula: str, position: int | None = None) -> None

Bases: ValueError

Error in explore formula syntax.

Provides helpful error messages with position indicators for syntax errors.

Parameters:

Attributes:

Parameters:

Attributes¶

formula¶

formula = formula

position¶

position = position

ExploreFormulaSpec¶

Parsed explore formula.

Represents a parsed explore formula with focal variable, optional contrast type, and conditioning specifications.

Created by: parse_explore_formula() Consumed by: dispatch_marginal_computation(), plot_predict(), plot_explore() Augmented by: attrs.evolve() in resolve_focal_at_spec() materializes focal at-values

Attributes:

Attributes¶

conditions¶

conditions: tuple[Condition, ...] = field(factory=tuple, converter=tuple)

contrast_degree¶

contrast_degree: int | None = field(default=None, validator=is_optional_positive_int)

contrast_expr¶

contrast_expr: ContrastExpr | None = field(default=None)

contrast_level_ordering¶

contrast_level_ordering: tuple[str, ...] | None = field(default=None, validator=is_optional_tuple_of_str)

contrast_ref¶

contrast_ref: str | None = field(default=None, validator=is_optional_str)

contrast_type¶

contrast_type: str | None = field(default=None, validator=(validators.optional(is_choice_str(('pairwise', 'sequential', 'poly', 'treatment', 'sum', 'helmert', 'custom')))))

focal_at_quantile¶

focal_at_quantile: int | None = field(default=None, validator=is_optional_positive_int)

focal_at_range¶

focal_at_range: int | None = field(default=None, validator=is_optional_positive_int)

focal_at_values¶

focal_at_values: tuple[float | str, ...] | None = field(default=None)

focal_var¶

focal_var: str = field(validator=(validators.instance_of(str)))

has_conditions¶

has_conditions: bool

Return True if conditioning variables are specified.

has_contrast¶

has_contrast: bool

Return True if any contrast is specified (named function or bracket expr).

has_contrast_expr¶

has_contrast_expr: bool

Return True if a bracket contrast expression is specified.

has_rhs_contrasts¶

has_rhs_contrasts: bool

Return True if any RHS condition has a bracket contrast expression.

Functions¶

parse_explore_formula¶

parse_explore_formula(formula: str, model_terms: list[str] | None = None) -> ExploreFormulaSpec

Parse an explore formula string.

Parameters:

Returns:

Examples:

Simple term::

>>> parse_explore_formula("treatment")
ExploreFormulaSpec(focal_var='treatment', contrast_type=None, conditions=())

Contrast function::

>>> parse_explore_formula("pairwise(treatment)")
ExploreFormulaSpec(focal_var='treatment', contrast_type='pairwise', conditions=())

With conditioning::

>>> parse_explore_formula("treatment ~ age@50")
ExploreFormulaSpec(
    focal_var='treatment',
    contrast_type=None,
    conditions=(Condition(var='age', at_values=(50.0,)),)
)

parse_lhs¶

parse_lhs(lhs: str, formula: str) -> tuple[str, str | None, int | None, tuple[float | str, ...] | None, int | None, int | None]

Parse LHS of explore formula.

Backward-compatible wrapper that tokenizes and parses just the LHS portion of an explore formula.

Parameters:

Returns:

Note: This wrapper does not return contrast_ref. Use parse_explore_formula() for the full result including contrast_ref.

explore_parser¶

Explore formula recursive descent parser.

Parses token streams from :class:ExploreScanner into :class:ExploreFormulaSpec containers.

Grammar::

explore_formula  := lhs [ TILDE rhs ]

lhs              := contrast_fn_call | contrast_expr | focal_term
contrast_fn_call := fn_name LEFT_PAREN focal_term { COMMA arg } RIGHT_PAREN
fn_name          := 'pairwise' | 'sequential' | 'poly'
                  | 'treatment' | 'dummy'
                  | 'sum' | 'deviation'
                  | 'helmert'
arg              := kwarg | positional_arg
kwarg            := IDENTIFIER EQUAL value
positional_arg   := value

contrast_expr    := IDENTIFIER LEFT_BRACKET contrast_list RIGHT_BRACKET
contrast_list    := contrast_item { COMMA contrast_item }
contrast_item    := contrast_operand MINUS contrast_operand
contrast_operand := STAR
                  | LEFT_PAREN level_name { PLUS level_name } RIGHT_PAREN
                  | level_name
level_name       := IDENTIFIER | STRING | NUMBER

focal_term       := IDENTIFIER [ AT at_spec ]
at_spec          := value
                  | LEFT_BRACKET value_list RIGHT_BRACKET
                  | [COLON] range_or_quantile
range_or_quantile := IDENTIFIER LEFT_PAREN NUMBER RIGHT_PAREN
                     (where IDENTIFIER is 'range' or 'quantile')
value_list       := value { COMMA value }
value            := NUMBER | STRING | IDENTIFIER

rhs              := condition { PLUS condition }
condition        := IDENTIFIER LEFT_BRACKET contrast_list RIGHT_BRACKET
                  | IDENTIFIER [ AT at_spec ]

Classes:

Attributes¶

Classes¶

ExploreParser¶

ExploreParser(tokens: list[Token], formula: str) -> None

Recursive descent parser for explore formula syntax.

Consumes tokens produced by :class:ExploreScanner and builds an :class:ExploreFormulaSpec container.

Parameters:

Functions:

Functions¶

parse¶

parse() -> ExploreFormulaSpec

Parse explore formula tokens into ExploreFormulaSpec container.

Returns:

Functions¶

explore_scanner¶

Explore formula scanner/tokenizer.

Extends the formula scanner with @ token support for explore-specific at-spec syntax (e.g., Days@50, Income@range(5)).

Classes:

Classes¶

ExploreScanner¶

Bases: Scanner

Scanner for explore formulas.

Extends the base formula scanner to:

1. Recognize @ as an AT token.

2. Skip intercept insertion (explore formulas have no implicit intercept).

3. Allow multiple tildes check to be skipped (explore ~ separates LHS/RHS, not response/predictors).

Parameters:

Functions:

Attributes:

Attributes¶

code¶

code = code

current¶

current = 0

start¶

start = 0

tokens¶

tokens: list[Token] = []

Functions¶

add_token¶

add_token(kind: str, literal: object = None) -> None

advance¶

advance() -> str

at_end¶

at_end() -> bool

backquote¶

backquote() -> None

char¶

char() -> None

floatnum¶

floatnum() -> None

identifier¶

identifier() -> None

match¶

match(expected: str) -> bool

number¶

number() -> None

peek¶

peek() -> str

peek_next¶

peek_next() -> str

scan¶

scan(add_intercept: bool = False) -> list[Token]

Scan explore formula string.

Overrides base to skip intercept insertion and skip the multiple-tilde validation (explore formulas use ~ differently).

Parameters:

Returns:

scan_token¶

scan_token() -> None

Scan a single token, adding @ support.

Functions¶

factors¶

Factor level extraction and handling utilities.

Functions:

Classes¶

Functions¶

detect_factor_levels_from_data¶

detect_factor_levels_from_data(data: 'pl.DataFrame', var: str) -> list[str]

Infer factor levels from unique values in data column.

Extracts unique values and sorts them for consistent ordering. Values are converted to strings for uniform handling.

Parameters:

Returns:

Examples:

>>> import polars as pl
>>> df = pl.DataFrame({"group": ["B", "A", "B", "C", "A"]})
>>> detect_factor_levels_from_data(df, "group")
['A', 'B', 'C']

get_factor_levels¶

get_factor_levels(bundle: 'DataBundle', var: str, *, fallback_data: 'pl.DataFrame | None' = None, allow_inference: bool = True) -> list[str]

Extract factor levels for a variable from bundle or data.

Priority order:

1. bundle.factor_levels[var] if present

2. Infer from fallback_data[var].unique() if provided and allow_inference=True

3. Raise ValueError if not found and no fallback

Parameters:

Returns:

Examples:

Get levels from bundle metadata::

levels = get_factor_levels(bundle, "treatment")
# ["control", "drug_a", "drug_b"]

Get levels with data fallback::

levels = get_factor_levels(bundle, "group", fallback_data=df)
# Inferred from df["group"].unique()

grid¶

Reference grid construction for marginal effects.

Functions:

Classes¶

Functions¶

build_reference_grid¶

build_reference_grid(bundle: DataBundle, focal_vars: list[str], *, at: dict[str, Any] | None = None, covariate_means: dict[str, float] | None = None) -> pl.DataFrame

Construct reference grid for marginal effects evaluation.

Creates a grid of covariate values at which to evaluate predictions. This function transforms the parsed explore formula’s conditions into a concrete evaluation grid.

The grid construction follows emmeans conventions:

• Focal variables: vary across their unique levels/values from data

• Non-focal continuous: set to their overall mean from the data

• Non-focal categorical: create rows for all levels (equal weight averaging)

• Conditioned variables (@ spec): use the specified values

Parameters:

Returns:

Integration with explore(): Called from model._compute_marginal_effects() after parsing::

    # From parsed explore formula "treatment ~ age@50"
    parsed = ExploreFormulaSpec(
        focal_var="treatment",
        conditions=[Condition(var="age", at_values=(50,))]
    )

    # Convert conditions to at dict
    at_dict = {"age": 50.0}

    # Build grid
    grid = build_reference_grid(
        bundle=self._bundle,
        focal_vars=["treatment"],
        at=at_dict,
    )
    # grid: pl.DataFrame with columns ["treatment"] and rows ["A", "B", "C"]
    # age is NOT in the grid (it's held fixed at 50)

Examples:

Grid for categorical focal::

grid = build_reference_grid(bundle, ["treatment"])
# Returns grid with one row per treatment level

Grid with conditioning::

grid = build_reference_grid(bundle, ["treatment"], at={"age": 50})
# Returns grid at age=50

Multiple focal variables (interaction EMMs)::

grid = build_reference_grid(bundle, ["treatment", "sex"])
# Returns Cartesian product: treatment x sex levels

inference¶

Inference for marginal effects using the delta method.

This module provides inference computation for MeeState results. Uses the delta method: SE = sqrt(diag(L @ vcov @ L’))

compute_mee_inference: Add inference to MeeState via delta method compute_mee_inference: Add inference to MeeState via delta method compute_mee_se: Compute SEs for MEE (means or slopes) compute_mee_inference_fallback: Fallback inference without L_matrix

Functions:

Classes¶

Functions¶

compute_mee_inference¶

compute_mee_inference(mee: MeeState, vcov: np.ndarray, df_resid: float | np.ndarray | None, conf_level: float = 0.95, null: float = 0.0, alternative: str = 'two-sided') -> MeeState

Compute delta method inference for marginal effects.

SE = sqrt(diag(L @ vcov @ L’)) SE = sqrt(diag(L @ vcov @ L’))

where L is the design matrix (X_ref for EMMs, selector for slopes).

Applies multiplicity adjustment for contrast CIs to match R’s emmeans:

• Tukey HSD for pairwise contrasts.

• Multivariate-t (MVT) for sequential, treatment, sum, helmert contrasts.

• No adjustment for polynomial contrasts or non-contrast MEEs.

For response-scale EMMs (effect_scale="response"), CIs are computed on the link scale then back-transformed via the inverse link function, matching R’s emmeans behavior.

Parameters:

Returns:

Examples:

Compute inference for EMMs::

from marginal import compute_emm
from inference import compute_mee_inference

mee = compute_emm(bundle, fit, "treatment", "treatment")
mee_with_inf = compute_mee_inference(mee, fit.vcov, fit.df_resid)
# mee_with_inf.se: standard errors
# mee_with_inf.ci_lower, mee_with_inf.ci_upper: confidence bounds

Note: For EMMs, L_matrix is the reference design matrix X_ref. For slopes, L_matrix is a selector row that picks the coefficient. For contrasts, L_matrix is the contrast matrix applied to EMMs.

compute_mee_inference_fallback¶

compute_mee_inference_fallback(mee: 'MeeState', bundle: 'DataBundle', fit: 'FitState', data: 'pl.DataFrame', df_resid: float | None, conf_level: float = 0.95, null: float = 0.0, alternative: str = 'two-sided') -> 'MeeState'

Compute inference for MEE without L_matrix (fallback path).

Used when L_matrix is not available (legacy MEE). Computes simplified SEs via compute_mee_se, then builds test statistics, p-values, and confidence intervals.

Parameters:

Returns:

compute_mee_se¶

compute_mee_se(mee: 'MeeState', bundle: 'DataBundle', fit: 'FitState', data: 'pl.DataFrame') -> np.ndarray

Compute standard errors for MEE estimates (means or slopes).

When an L_matrix is available, uses the delta method: SE = sqrt(diag(L @ vcov @ L')). This correctly handles multi-row L_matrices (e.g. per-group conditional slopes).

Otherwise falls back to simplified type-based dispatch:

• "means": SE = sqrt(MSE / n_per_group) for each group level

• "slopes": SE = sqrt(vcov[idx, idx]) from coefficient vcov diagonal

Parameters:

Returns:

joint_tests¶

ANOVA-style joint hypothesis tests for model terms.

Functions:

Classes¶

Functions¶

compute_joint_test¶

compute_joint_test(fit: FitState, bundle: DataBundle, spec: ModelSpec, terms: list[str] | None = None, *, errors: str = 'iid', data: pl.DataFrame | None = None) -> JointTestState

Compute joint hypothesis tests for model terms.

Tests each model term (continuous variables, factors, and interactions) using direct coefficient tests. This matches R’s emmeans::joint_tests() behavior.

For each term:

• Continuous variables: F-test on the single coefficient (df1=1)

• Categorical factors: Joint F-test on all indicator coefficients

• Interactions: Joint F-test on all interaction coefficients

The test type is determined by the model family:

• Gaussian (LM/LMER): F-test with df_resid

• Non-Gaussian (GLM/GLMER): Chi-square test (asymptotic)

Parameters:

Returns:

Examples:

Basic usage::

from marginal import compute_joint_test
state = compute_joint_test(fit, bundle, spec)
# state.terms: ("x", "group", "x:group")
# state.statistic: F or chi2 values per term
# state.p_value: p-values per term

Welch ANOVA::

state = compute_joint_test(fit, bundle, spec, errors="unequal_var", data=df)
# Per-term Satterthwaite df2 in state.df2

Test specific terms only::

state = compute_joint_test(fit, bundle, spec, terms=["group"])

Note: Intercept is never tested. Random effect grouping factors are excluded for mixed models.

matrices¶

Contrast matrix builders for EMM comparisons.

Provides functions to build hypothesis contrast matrices used in marginal effects computation. These matrices transform EMM vectors into contrast estimates (e.g., pairwise differences, sequential differences, polynomial trends).

Distinct from design.coding which builds design matrix coding matrices for categorical variable encoding during model fitting.

Created by: absorbed from maths/inference/contrasts.py (zero consumers outside marginal/) + poly_contrasts from contrasts.py. Consumed by: marginal/contrasts.py (apply_contrasts, apply_contrasts_grouped).

Functions:

Functions¶

build_all_pairwise_matrix¶

build_all_pairwise_matrix(n_levels: int) -> np.ndarray

Build all pairwise contrasts between EMM levels.

Creates C(n, 2) = n*(n-1)/2 contrasts for all unique pairs. Unlike build_pairwise_matrix(), this includes all pairs, not just comparisons to the reference level.

Parameters:

Returns:

Examples:

>>> C = build_all_pairwise_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.],
       [ 0., -1.,  1.]])

build_contrast_matrix¶

build_contrast_matrix(contrast_type: str | dict, levels: list, normalize: bool = False) -> np.ndarray

Build a contrast matrix based on contrast type.

Dispatcher function that builds the appropriate contrast matrix based on the contrast type specification.

Parameters:

Returns:

Examples:

>>> build_contrast_matrix("pairwise", ["A", "B", "C"])
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.],
       [ 0., -1.,  1.]])

build_helmert_matrix¶

build_helmert_matrix(n_levels: int) -> np.ndarray

Build Helmert contrasts (each level vs mean of previous levels).

Creates n_levels - 1 contrasts where level k is compared to the mean of levels 0, 1, ..., k-1.

Parameters:

Returns:

Examples:

>>> C = build_helmert_matrix(3)
>>> C
array([[-1. ,  1. ,  0. ],
       [-0.5, -0.5,  1. ]])

>>> C = build_helmert_matrix(4)
>>> C
array([[-1.        ,  1.        ,  0.        ,  0.        ],
       [-0.5       , -0.5       ,  1.        ,  0.        ],
       [-0.33333333, -0.33333333, -0.33333333,  1.        ]])

build_pairwise_matrix¶

build_pairwise_matrix(n_levels: int) -> np.ndarray

Build (n-1) linearly independent pairwise contrasts.

Creates “treatment-style” contrasts comparing each level to the first (reference) level. This produces n_levels - 1 contrasts that span the space of all pairwise differences.

Parameters:

Returns:

Examples:

>>> C = build_pairwise_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [-1.,  0.,  1.]])

build_poly_matrix¶

build_poly_matrix(n_levels: int, degree: int | None = None) -> np.ndarray

Build orthogonal polynomial contrast matrix for EMMs.

Creates orthogonal polynomial contrasts for ordered factors. Tests for linear, quadratic, cubic (etc.) trends across the factor levels.

Parameters:

Returns:

Examples:

>>> build_poly_matrix(5)  # 4x5 matrix
>>> build_poly_matrix(5, degree=2)  # 2x5 matrix (linear + quadratic)

build_sequential_matrix¶

build_sequential_matrix(n_levels: int) -> np.ndarray

Build sequential (successive differences) contrasts.

Creates n_levels - 1 contrasts comparing each level to the previous one.

Parameters:

Returns:

Examples:

>>> C = build_sequential_matrix(3)
>>> C
array([[-1.,  1.,  0.],
       [ 0., -1.,  1.]])

>>> C = build_sequential_matrix(4)
>>> C
array([[-1.,  1.,  0.,  0.],
       [ 0., -1.,  1.,  0.],
       [ 0.,  0., -1.,  1.]])

build_sum_to_zero_matrix¶

build_sum_to_zero_matrix(n_levels: int) -> np.ndarray

Build sum-to-zero contrasts (deviation coding).

Creates contrasts comparing each level to the grand mean.

Parameters:

Returns: