fit

Model fitting, diagnostics, convergence, varying parameters, and prediction.

Call chain:

model.fit() -> fit_model() -> resolve_solver() -> fit_ols_qr / fit_glm_irls / fit_lmer_pls / fit_glmer_pirls

Attributes:

Classes:

Functions:

Modules:

Attributes¶

VALID_SOLVERS¶

VALID_SOLVERS: frozenset[str] = frozenset({'qr', 'irls', 'pls', 'pirls'})

Classes¶

FitResult¶

Immutable result of the fit lifecycle.

Attributes:

Attributes¶

augmented_data¶

augmented_data: pl.DataFrame | None

bundle¶

bundle: DataBundle

fit¶

fit: FitState

formula_spec¶

formula_spec: object

raw_data¶

raw_data: pl.DataFrame | None

varying_offsets¶

varying_offsets: VaryingState | None = None

varying_spread¶

varying_spread: VaryingSpreadState | None = None

Functions¶

augment_data_with_diagnostics¶

augment_data_with_diagnostics(*, raw_data: pl.DataFrame, fit: FitState, bundle: DataBundle) -> pl.DataFrame

Augment raw data with diagnostic columns after fit.

Adds fitted, resid, hat, std_resid, cooksd columns (names from AugmentedDataCols schema). Values are NaN for rows dropped due to missing data.

Parameters:

Returns:

build_mixed_post_fit_state¶

build_mixed_post_fit_state(fit: FitState, bundle: DataBundle, data: pl.DataFrame, *, stacklevel: int = 3) -> tuple[VaryingState | None, VaryingSpreadState | None]

Compute BLUPs, variance components, and emit convergence warnings.

Orchestrates the post-fit assembly for mixed-effects models: computes VaryingState (BLUPs) and VaryingSpreadState (variance components) from the fitted parameters, then checks for convergence issues.

Parameters:

Returns:

build_predict_grid¶

build_predict_grid(data: pl.DataFrame, focal_var: str, response_col: str, grouping_factors: tuple[str, ...], *, focal_values: list[float | str] | None = None, n_points: int | Literal['data'] = 50, varying_vars: list[str] | None = None, at: dict[str, Any] | None = None) -> pl.DataFrame

Build a Cartesian-product prediction grid.

Creates a grid where the focal variable is varied, condition variables are expanded, and all other predictors are held at reference values (mean for continuous, first sorted level for categorical). Grouping factors and the response column are excluded.

Parameters:

Returns:

check_convergence¶

check_convergence(fit: FitState, re_meta: REInfo) -> list[ConvergenceMessage]

Run convergence diagnostics on a fitted mixed model.

Extracts theta and sigma from FitState, computes theta lower bounds and per-factor RE info from REInfo, and delegates to diagnose_convergence() for the actual diagnostic checks.

Parameters:

Returns:

compute_diagnostics¶

compute_diagnostics(*, model_type: str, spec: ModelSpec, bundle: DataBundle, fit: FitState, coef_for_predict: np.ndarray, varying_spread: VaryingSpreadState | None, cv: CVState | None, has_intercept: bool = True) -> pl.DataFrame

Compute model-level diagnostics as a single-row DataFrame.

Builds goodness-of-fit diagnostics from fitted model state, with columns varying by model type (lm, glm, lmer, glmer).

Parameters:

Returns:

compute_metadata¶

compute_metadata(*, bundle: DataBundle) -> pl.DataFrame

Compute model metadata as a single-row DataFrame.

Returns sample/structural info about the model: observation counts, parameter count, and group counts (for mixed models).

Parameters:

Returns:

compute_optimizer_diagnostics¶

compute_optimizer_diagnostics(*, model_type: str, fit: FitState) -> pl.DataFrame

Compute optimizer convergence diagnostics as a single-row DataFrame.

Parameters:

Returns:

compute_predictions_from_formula¶

compute_predictions_from_formula(formula: str, data: pl.DataFrame, spec: object, bundle: object, fit: object, formula_spec: object, pred_type: str, varying: str, allow_new_levels: bool, n_points: int | Literal['data']) -> 'PredictionState'

Parse a predict formula, build the grid, compute predictions, and attach grid columns.

Combines parse_predict_formula, compute_predictions, and grid-column attachment into a single call for model.predict() formula mode.

Parameters:

Returns:

compute_r_squared¶

compute_r_squared(y: np.ndarray, residuals: np.ndarray, n: int, p: int, has_intercept: bool = True) -> tuple[float, float]

Compute R-squared and adjusted R-squared from raw arrays.

For models with an intercept, uses centered SS_tot = sum((y - mean(y))^2). For no-intercept models, uses uncentered SS_tot = sum(y^2), matching R’s summary.lm() behavior.

Parameters:

Returns:

compute_varying_spread_state¶

compute_varying_spread_state(theta: NDArray[np.floating], sigma: float, re_meta: REInfo) -> VaryingSpreadState

Compute VaryingSpreadState (variance components) from theta parameters.

Extracts residual variance (sigma²), random effect variances (tau²), correlations (rho), and intraclass correlation (ICC) from the fitted theta vector using the random effects structure.

Parameters:

Returns:

compute_varying_state¶

compute_varying_state(theta: NDArray[np.floating], u: NDArray[np.floating], re_meta: REInfo, data: pl.DataFrame | None = None) -> VaryingState

Compute VaryingState (BLUPs) from fitted random effects parameters.

Converts spherical random effects u to BLUPs b = Lambda @ u using the relative covariance factor Lambda built from theta. Constructs a grid of group/level combinations and maps BLUP values to named effects.

Parameters:

Returns:

execute_fit¶

execute_fit(spec: ModelSpec, bundle: DataBundle | None, data: pl.DataFrame, raw_data: pl.DataFrame | None, formula: str, custom_contrasts: dict | None, weights_col: str | None, offset_col: str | None, missing: str, is_mixed: bool, solver_override: str | None, fit_kwargs: dict) -> FitResult

Execute the full fit lifecycle: bundle rebuild → fit → post-fit state → diagnostics.

Parameters:

Returns:

fit_glm_irls¶

fit_glm_irls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 25, tol: float = 1e-08) -> FitState

Fit generalized linear model using Iteratively Reweighted Least Squares.

This adapter wraps the IRLS implementation in IRLS solves GLMs by iterating between computing working weights and solving a weighted least squares problem.

1. Initialize mu from y (or link function default)

2. Initialize mu from y (or link function default)

3. For each iteration: a. Compute working weights: W = 1 / (V(mu) * g’(mu)^2) b. Compute working response: z = eta + (y - mu) * g’(mu) c. Solve weighted least squares: beta = (X’WX)^{-1} X’Wz d. Update eta = X @ beta, mu = g^{-1}(eta)

4. Continue until convergence (change in deviance < tol)

• gaussian: Identity variance, identity link

• gaussian: Identity variance, identity link

• binomial: mu(1-mu) variance, logit/probit/cloglog link

• poisson: mu variance, log link

• gamma: mu^2 variance, inverse/log link

Parameters:

Returns:

See Also:

• glm: Underlying IRLS implementation

fit_glmer_pirls¶

fit_glmer_pirls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 25, max_outer_iter: int = 10000, tol: float = 1e-07, verbose: bool = False, nAGQ: int = 1, use_hessian: bool = False) -> FitState

Fit generalized linear mixed model using Penalized IRLS.

This adapter wraps the PIRLS implementation from PIRLS combines IRLS (for the GLM part) with PLS (for random effects), using Laplace approximation to integrate out the random effects.

Outer loop (BOBYQA optimization over theta): Outer loop (BOBYQA optimization over theta): For each theta: 1. Build Lambda from theta

Inner loop (PIRLS iterations):
    a. Compute working weights from current eta/mu
    b. Compute working response
    c. Solve weighted PLS for beta and u
    d. Update eta = X @ beta + Z @ Lambda @ u
    e. Update mu = g^{-1}(eta)
    f. Step-halving if deviance increased
    g. Check convergence

2. Return Laplace deviance

Select theta minimizing Laplace deviance

Parameters:

Returns:

See Also:

• glmer: Underlying PIRLS implementation

fit_lmer_pls¶

fit_lmer_pls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 10000, verbose: bool = False) -> FitState

Fit linear mixed-effects model using Penalized Least Squares.

This adapter wraps the PLS implementation from PLS is the algorithm from Bates et al. (2015) used in R’s lme4 package.

Outer loop (BOBYQA optimization over theta): Outer loop (BOBYQA optimization over theta): For each theta (relative covariance parameters): 1. Build Lambda (block-diagonal Cholesky factor from theta) 2. Form S_22 = Lambda’ Z’ Z Lambda + I 3. Sparse Cholesky factorization of S_22 4. Compute Schur complement for fixed effects 5. Solve for beta (fixed effects) and u (spherical RE) 6. Compute REML or ML deviance

Parameters:

Returns:

See Also:

• lmer: Underlying PLS implementation

fit_model¶

fit_model(spec: ModelSpec, bundle: DataBundle, *, solver: str | None = None, max_iter: int | None = None, max_outer_iter: int = 10000, tol: float | None = None, verbose: bool = False, nAGQ: int = 1, use_hessian: bool = False) -> FitState

Dispatch to appropriate fitter based on model specification.

This is the main entry point for fitting models. It examines the ModelSpec to determine the appropriate solver and delegates to the corresponding fitter function.

If the design matrix is rank-deficient (detected during bundle construction), the X matrix is reduced to estimable columns before fitting. After fitting, coefficients and vcov are expanded back to full size with NaN for dropped columns (matching R’s lm() behavior).

The solver selection follows the estimation method matrix:

Parameters:

Returns:

Examples:

>>> import numpy as np
>>> from containers import build_model_spec, DataBundle
>>> spec = build_model_spec(
...     formula="y ~ x",
...     response_var="y",
...     fixed_terms=["Intercept", "x"],
... )
>>> bundle = DataBundle(
...     X=np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
...     y=np.array([2.0, 4.0, 6.0]),
...     X_names=["Intercept", "x"],
...     y_name="y",
...     valid_mask=np.array([True, True, True]),
...     n_total=3,
... )
>>> state = fit_model(spec, bundle)
>>> state.converged
True
>>> state.coef  # [Intercept, x] = [0, 2]
array([0., 2.])

fit_ols_qr¶

fit_ols_qr(spec: ModelSpec, bundle: DataBundle) -> FitState

Fit ordinary or weighted least squares using QR decomposition.

Supports observation weights (WLS) and offset terms. When weights are present, solves the transformed system sqrt(W)*X, sqrt(W)*y via QR decomposition, which yields WLS coefficients and vcov directly. Offsets are subtracted from y before fitting and added back to fitted values.

1. Subtract offset from y (if present): y_adj = y - offset

2. Subtract offset from y (if present): y_adj = y - offset

3. Apply weights (if present): X_w = sqrt(w)*X, y_w = sqrt(w)*y_adj

4. QR decompose X_w with column pivoting for stability

5. Solve R * beta = Q.T @ y_w via back-substitution

6. Recompute original-scale: fitted = X @ beta + offset, resid = y - fitted

7. vcov = sigma_w^2 * (X’WX)^{-1}

8. Leverage from (possibly weighted) hat matrix

Matches R’s logLik.lm formula:: Matches R’s logLik.lm formula::

L = 0.5*sum(log(w)) - n/2 * (log(2*pi) + log(RSS_w/n) + 1)

The 0.5*sum(log(w)) term is the Jacobian from the weight transformation.

Parameters:

Returns:

Examples:

>>> import numpy as np
>>> from containers import build_model_spec, DataBundle
>>> spec = build_model_spec(
...     formula="y ~ x",
...     response_var="y",
...     fixed_terms=["Intercept", "x"],
... )
>>> bundle = DataBundle(
...     X=np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
...     y=np.array([2.0, 4.0, 6.0]),
...     X_names=["Intercept", "x"],
...     y_name="y",
...     valid_mask=np.array([True, True, True]),
...     n_total=3,
... )
>>> state = fit_ols_qr(spec, bundle)
>>> np.allclose(state.fitted + state.residuals, bundle.y)
True
>>> np.allclose(state.coef, [0.0, 2.0])  # Perfect fit: y = 2x
True

get_theta_lower_bounds¶

get_theta_lower_bounds(n_theta: int, re_structure: str, metadata: dict | None = None) -> list[float]

Get lower bounds for theta parameters.

Diagonal elements of Cholesky factor must be non-negative. Off-diagonal elements are unbounded.

Parameters:

Returns:

parse_fit_kwargs¶

parse_fit_kwargs(spec: ModelSpec, kwargs: dict[str, object], nAGQ: int | None) -> tuple[ModelSpec, str | None, dict[str, object]]

Validate and extract fitting parameters from **kwargs.

Pops solver, method, and nAGQ from kwargs, validates each, and assembles the remaining fit-specific keyword arguments into a dict suitable for fit_model().

Parameters:

Returns:

parse_predict_formula¶

parse_predict_formula(formula: str, data: pl.DataFrame, response_col: str, grouping_factors: tuple[str, ...], *, n_points: int | Literal['data'] = 50) -> tuple[pl.DataFrame, list[str]]

Parse an explore-style formula and build a prediction grid.

Translates the formula via :func:parse_explore_formula, rejects contrast formulas, and delegates to :func:build_predict_grid.

Parameters:

Returns:

per_factor_re_info¶

per_factor_re_info(re_meta: REInfo, group_names: list[str]) -> tuple[str | list[str], list[str] | dict[str, list[str]]]

Split global RE metadata into per-factor structures and names.

For crossed/nested/mixed models, the global re_structure is a single string (e.g. “crossed”) and random_names is a concatenated list across all factors. This function splits them into per-factor structures and per-factor name dicts suitable for BLUP decomposition and convergence diagnostics.

For single-factor models, returns the originals unchanged.

Parameters:

Returns:

resolve_condition_values¶

resolve_condition_values(cond: Condition, data: pl.DataFrame) -> list | None

Resolve a :class:Condition to concrete values or None.

Parameters:

Returns:

resolve_solver¶

resolve_solver(spec: ModelSpec) -> str

Select the appropriate solver for a model configuration.

Parameters:

Returns:

validate_fit_method¶

validate_fit_method(spec: ModelSpec, method_str: str) -> ModelSpec

Validate and apply a user-specified fitting method to a ModelSpec.

Checks that the method is compatible with the model’s family and random-effects structure, then returns an evolved spec with the new method.

Parameters:

Returns:

Modules¶

convergence¶

Convergence diagnostics for fitted mixed-effects models.

Encapsulates the repeated pattern of extracting theta, computing bounds, assembling per-factor RE info, and running diagnose_convergence().

Functions:

Classes¶

Functions¶

check_convergence¶

check_convergence(fit: FitState, re_meta: REInfo) -> list[ConvergenceMessage]

Run convergence diagnostics on a fitted mixed model.

Extracts theta and sigma from FitState, computes theta lower bounds and per-factor RE info from REInfo, and delegates to diagnose_convergence() for the actual diagnostic checks.

Parameters:

Returns:

diagnostics¶

Model-level diagnostics computation.

Pure functions that compute model diagnostics from containers. These were extracted from model/core.py to keep the model class as thin glue.

Attributes¶

Classes¶

Functions¶

augment_data_with_diagnostics¶

augment_data_with_diagnostics(*, raw_data: pl.DataFrame, fit: FitState, bundle: DataBundle) -> pl.DataFrame

Augment raw data with diagnostic columns after fit.

Adds fitted, resid, hat, std_resid, cooksd columns (names from AugmentedDataCols schema). Values are NaN for rows dropped due to missing data.

Parameters:

Returns:

compute_diagnostics¶

compute_diagnostics(*, model_type: str, spec: ModelSpec, bundle: DataBundle, fit: FitState, coef_for_predict: np.ndarray, varying_spread: VaryingSpreadState | None, cv: CVState | None, has_intercept: bool = True) -> pl.DataFrame

Compute model-level diagnostics as a single-row DataFrame.

Builds goodness-of-fit diagnostics from fitted model state, with columns varying by model type (lm, glm, lmer, glmer).

Parameters:

Returns:

compute_metadata¶

compute_metadata(*, bundle: DataBundle) -> pl.DataFrame

Compute model metadata as a single-row DataFrame.

Returns sample/structural info about the model: observation counts, parameter count, and group counts (for mixed models).

Parameters:

Returns:

compute_optimizer_diagnostics¶

compute_optimizer_diagnostics(*, model_type: str, fit: FitState) -> pl.DataFrame

Compute optimizer convergence diagnostics as a single-row DataFrame.

Parameters:

Returns:

compute_r_squared¶

compute_r_squared(y: np.ndarray, residuals: np.ndarray, n: int, p: int, has_intercept: bool = True) -> tuple[float, float]

Compute R-squared and adjusted R-squared from raw arrays.

For models with an intercept, uses centered SS_tot = sum((y - mean(y))^2). For no-intercept models, uses uncentered SS_tot = sum(y^2), matching R’s summary.lm() behavior.

Parameters:

Returns:

dispatch¶

Solver dispatch for model fitting.

Provides fit_model() which dispatches to the appropriate fitter based on model specification, and resolve_solver() which determines the solver type.

Handles rank-deficient design matrices by reducing X before fitting and expanding coefficients/vcov after, inserting NaN for dropped columns.

Functions:

Attributes:

Attributes¶

VALID_SOLVERS¶

VALID_SOLVERS: frozenset[str] = frozenset({'qr', 'irls', 'pls', 'pirls'})

Classes¶

Functions¶

fit_model¶

fit_model(spec: ModelSpec, bundle: DataBundle, *, solver: str | None = None, max_iter: int | None = None, max_outer_iter: int = 10000, tol: float | None = None, verbose: bool = False, nAGQ: int = 1, use_hessian: bool = False) -> FitState

Dispatch to appropriate fitter based on model specification.

This is the main entry point for fitting models. It examines the ModelSpec to determine the appropriate solver and delegates to the corresponding fitter function.

If the design matrix is rank-deficient (detected during bundle construction), the X matrix is reduced to estimable columns before fitting. After fitting, coefficients and vcov are expanded back to full size with NaN for dropped columns (matching R’s lm() behavior).

The solver selection follows the estimation method matrix:

Parameters:

Returns:

Examples:

>>> import numpy as np
>>> from containers import build_model_spec, DataBundle
>>> spec = build_model_spec(
...     formula="y ~ x",
...     response_var="y",
...     fixed_terms=["Intercept", "x"],
... )
>>> bundle = DataBundle(
...     X=np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
...     y=np.array([2.0, 4.0, 6.0]),
...     X_names=["Intercept", "x"],
...     y_name="y",
...     valid_mask=np.array([True, True, True]),
...     n_total=3,
... )
>>> state = fit_model(spec, bundle)
>>> state.converged
True
>>> state.coef  # [Intercept, x] = [0, 2]
array([0., 2.])

parse_fit_kwargs¶

parse_fit_kwargs(spec: ModelSpec, kwargs: dict[str, object], nAGQ: int | None) -> tuple[ModelSpec, str | None, dict[str, object]]

Validate and extract fitting parameters from **kwargs.

Pops solver, method, and nAGQ from kwargs, validates each, and assembles the remaining fit-specific keyword arguments into a dict suitable for fit_model().

Parameters:

Returns:

resolve_solver¶

resolve_solver(spec: ModelSpec) -> str

Select the appropriate solver for a model configuration.

Parameters:

Returns:

validate_fit_method¶

validate_fit_method(spec: ModelSpec, method_str: str) -> ModelSpec

Validate and apply a user-specified fitting method to a ModelSpec.

Checks that the method is compatible with the model’s family and random-effects structure, then returns an evolved spec with the new method.

Parameters:

Returns:

glm¶

GLM fitting via Iteratively Reweighted Least Squares (IRLS).

Functions:

Classes¶

Functions¶

fit_glm_irls¶

fit_glm_irls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 25, tol: float = 1e-08) -> FitState

Fit generalized linear model using Iteratively Reweighted Least Squares.

This adapter wraps the IRLS implementation in IRLS solves GLMs by iterating between computing working weights and solving a weighted least squares problem.

1. Initialize mu from y (or link function default)

2. Initialize mu from y (or link function default)

3. For each iteration: a. Compute working weights: W = 1 / (V(mu) * g’(mu)^2) b. Compute working response: z = eta + (y - mu) * g’(mu) c. Solve weighted least squares: beta = (X’WX)^{-1} X’Wz d. Update eta = X @ beta, mu = g^{-1}(eta)

4. Continue until convergence (change in deviance < tol)

• gaussian: Identity variance, identity link

• gaussian: Identity variance, identity link

• binomial: mu(1-mu) variance, logit/probit/cloglog link

• poisson: mu variance, log link

• gamma: mu^2 variance, inverse/log link

Parameters:

Returns:

See Also:

• glm: Underlying IRLS implementation

Modules¶

glmer¶

GLMM fitting via Penalized IRLS (PIRLS).

Functions:

Classes¶

Functions¶

fit_glmer_pirls¶

fit_glmer_pirls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 25, max_outer_iter: int = 10000, tol: float = 1e-07, verbose: bool = False, nAGQ: int = 1, use_hessian: bool = False) -> FitState

Fit generalized linear mixed model using Penalized IRLS.

This adapter wraps the PIRLS implementation from PIRLS combines IRLS (for the GLM part) with PLS (for random effects), using Laplace approximation to integrate out the random effects.

Outer loop (BOBYQA optimization over theta): Outer loop (BOBYQA optimization over theta): For each theta: 1. Build Lambda from theta

Inner loop (PIRLS iterations):
    a. Compute working weights from current eta/mu
    b. Compute working response
    c. Solve weighted PLS for beta and u
    d. Update eta = X @ beta + Z @ Lambda @ u
    e. Update mu = g^{-1}(eta)
    f. Step-halving if deviance increased
    g. Check convergence

2. Return Laplace deviance

Select theta minimizing Laplace deviance

Parameters:

Returns:

See Also:

• glmer: Underlying PIRLS implementation

grid¶

Prediction grid construction for formula-mode predictions.

Provides parse_predict_formula() which translates an explore-style formula into a prediction grid (Polars DataFrame), and build_predict_grid() which assembles the Cartesian-product grid from column specifications.

Shared by model.predict() (formula mode) and viz/predict.py (plot_predict).

Functions:

Classes¶

Functions¶

build_predict_grid¶

build_predict_grid(data: pl.DataFrame, focal_var: str, response_col: str, grouping_factors: tuple[str, ...], *, focal_values: list[float | str] | None = None, n_points: int | Literal['data'] = 50, varying_vars: list[str] | None = None, at: dict[str, Any] | None = None) -> pl.DataFrame

Build a Cartesian-product prediction grid.

Creates a grid where the focal variable is varied, condition variables are expanded, and all other predictors are held at reference values (mean for continuous, first sorted level for categorical). Grouping factors and the response column are excluded.

Parameters:

Returns:

compute_predictions_from_formula¶

compute_predictions_from_formula(formula: str, data: pl.DataFrame, spec: object, bundle: object, fit: object, formula_spec: object, pred_type: str, varying: str, allow_new_levels: bool, n_points: int | Literal['data']) -> 'PredictionState'

Parse a predict formula, build the grid, compute predictions, and attach grid columns.

Combines parse_predict_formula, compute_predictions, and grid-column attachment into a single call for model.predict() formula mode.

Parameters:

Returns:

parse_predict_formula¶

parse_predict_formula(formula: str, data: pl.DataFrame, response_col: str, grouping_factors: tuple[str, ...], *, n_points: int | Literal['data'] = 50) -> tuple[pl.DataFrame, list[str]]

Parse an explore-style formula and build a prediction grid.

Translates the formula via :func:parse_explore_formula, rejects contrast formulas, and delegates to :func:build_predict_grid.

Parameters:

Returns:

resolve_condition_values¶

resolve_condition_values(cond: Condition, data: pl.DataFrame) -> list | None

Resolve a :class:Condition to concrete values or None.

Parameters:

Returns:

lifecycle¶

Fit lifecycle orchestration.

Owns the multi-step fit sequence: bundle rebuild → fit → post-fit state → diagnostics augmentation. Called by model.fit() so the model class stays a thin facade.

Classes:

Functions:

Classes¶

FitResult¶

Immutable result of the fit lifecycle.

Attributes:

Attributes¶

augmented_data¶

augmented_data: pl.DataFrame | None

bundle¶

bundle: DataBundle

fit¶

fit: FitState

formula_spec¶

formula_spec: object

raw_data¶

raw_data: pl.DataFrame | None

varying_offsets¶

varying_offsets: VaryingState | None = None

varying_spread¶

varying_spread: VaryingSpreadState | None = None

Functions¶

execute_fit¶

execute_fit(spec: ModelSpec, bundle: DataBundle | None, data: pl.DataFrame, raw_data: pl.DataFrame | None, formula: str, custom_contrasts: dict | None, weights_col: str | None, offset_col: str | None, missing: str, is_mixed: bool, solver_override: str | None, fit_kwargs: dict) -> FitResult

Execute the full fit lifecycle: bundle rebuild → fit → post-fit state → diagnostics.

Parameters:

Returns:

lmer¶

LMM fitting via Penalized Least Squares (PLS).

Functions:

Classes¶

Functions¶

fit_lmer_pls¶

fit_lmer_pls(spec: ModelSpec, bundle: DataBundle, *, max_iter: int = 10000, verbose: bool = False) -> FitState

Fit linear mixed-effects model using Penalized Least Squares.

This adapter wraps the PLS implementation from PLS is the algorithm from Bates et al. (2015) used in R’s lme4 package.

Outer loop (BOBYQA optimization over theta): Outer loop (BOBYQA optimization over theta): For each theta (relative covariance parameters): 1. Build Lambda (block-diagonal Cholesky factor from theta) 2. Form S_22 = Lambda’ Z’ Z Lambda + I 3. Sparse Cholesky factorization of S_22 4. Compute Schur complement for fixed effects 5. Solve for beta (fixed effects) and u (spherical RE) 6. Compute REML or ML deviance

Parameters:

Returns:

See Also:

• lmer: Underlying PLS implementation

ols¶

OLS fitting via QR decomposition.

Functions:

Classes¶

Functions¶

fit_ols_qr¶

fit_ols_qr(spec: ModelSpec, bundle: DataBundle) -> FitState

Fit ordinary or weighted least squares using QR decomposition.

Supports observation weights (WLS) and offset terms. When weights are present, solves the transformed system sqrt(W)*X, sqrt(W)*y via QR decomposition, which yields WLS coefficients and vcov directly. Offsets are subtracted from y before fitting and added back to fitted values.

1. Subtract offset from y (if present): y_adj = y - offset

2. Subtract offset from y (if present): y_adj = y - offset

3. Apply weights (if present): X_w = sqrt(w)*X, y_w = sqrt(w)*y_adj

4. QR decompose X_w with column pivoting for stability

5. Solve R * beta = Q.T @ y_w via back-substitution

6. Recompute original-scale: fitted = X @ beta + offset, resid = y - fitted

7. vcov = sigma_w^2 * (X’WX)^{-1}

8. Leverage from (possibly weighted) hat matrix

Matches R’s logLik.lm formula:: Matches R’s logLik.lm formula::

L = 0.5*sum(log(w)) - n/2 * (log(2*pi) + log(RSS_w/n) + 1)

The 0.5*sum(log(w)) term is the Jacobian from the weight transformation.

Parameters:

Returns:

Examples:

>>> import numpy as np
>>> from containers import build_model_spec, DataBundle
>>> spec = build_model_spec(
...     formula="y ~ x",
...     response_var="y",
...     fixed_terms=["Intercept", "x"],
... )
>>> bundle = DataBundle(
...     X=np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
...     y=np.array([2.0, 4.0, 6.0]),
...     X_names=["Intercept", "x"],
...     y_name="y",
...     valid_mask=np.array([True, True, True]),
...     n_total=3,
... )
>>> state = fit_ols_qr(spec, bundle)
>>> np.allclose(state.fitted + state.residuals, bundle.y)
True
>>> np.allclose(state.coef, [0.0, 2.0])  # Perfect fit: y = 2x
True

predict¶

Prediction operations on containers.

Pure functions for computing predictions on new data, including random effects contribution for mixed models. Extracted from model/core.py.

Classes¶

Functions¶

build_X_for_newdata¶

build_X_for_newdata(formula_spec: FormulaSpec | None, X_names: tuple[str, ...], newdata: pl.DataFrame) -> NDArray[np.float64]

Build design matrix X for new data.

Uses the stored FormulaSpec to properly handle factors, transformations (log, poly, center), and interactions. This ensures new data is encoded consistently with the training data.

When no FormulaSpec is available (e.g. simulation-only workflows), falls back to manual column stacking.

Parameters:

Returns:

build_re_covariates¶

build_re_covariates(newdata: pl.DataFrame, factor_names: list[str], valid_indices: NDArray[np.intp], formula_spec: FormulaSpec | None, X_names: tuple[str, ...]) -> NDArray[np.float64]

Build random effects covariate matrix for valid newdata rows.

For each random effect term (e.g. Intercept, slope variable), extracts the appropriate covariate values from newdata for the valid rows.

Parameters:

Returns:

compute_predictions¶

compute_predictions(spec: ModelSpec, bundle: DataBundle, fit: FitState, formula_spec: FormulaSpec | None, training_data: pl.DataFrame | None, newdata: pl.DataFrame | None, pred_type: Literal['response', 'link'], *, varying: Literal['exclude', 'include'] = 'exclude', allow_new_levels: bool = False) -> PredictionState

Compute predictions for given data.

For training data (newdata=None), returns fitted values directly. For new data, builds the design matrix, computes the linear predictor, optionally adds random effects, and applies the inverse link function.

Parameters:

Returns:

compute_re_contribution¶

compute_re_contribution(re_meta: REInfo, theta: NDArray[np.floating], u: NDArray[np.floating], training_data: pl.DataFrame | None, newdata: pl.DataFrame, valid_mask: NDArray[np.bool_], allow_new_levels: bool, formula_spec: FormulaSpec | None, X_names: tuple[str, ...]) -> NDArray[np.float64]

Compute random effects contribution for new data predictions.

For each valid observation in newdata, maps group labels to trained group indices and computes the BLUP contribution as the dot product of the random effects design row and the group’s estimated BLUPs.

Parameters:

Returns:

resolve_coef_for_predict¶

resolve_coef_for_predict(coef: NDArray[np.floating], rank_info: RankInfo | None) -> NDArray[np.floating]

Coefficients safe for matrix multiplication (NaN -> 0).

When rank-deficient columns produce NaN coefficients, those NaN values must be zeroed out for X @ coef to work correctly. The dropped columns contribute nothing to predictions.

Parameters:

Returns:

validate_newdata_groups¶

validate_newdata_groups(re_meta: REInfo, training_data: pl.DataFrame | None, newdata: pl.DataFrame, allow_new_levels: bool) -> None

Validate grouping columns and levels in new data for mixed models.

Ensures that all grouping variables exist in newdata and that group levels are a subset of those seen during training (unless allow_new_levels is True).

Called from compute_predictions when varying="include" for mixed models, ensuring group structure is valid before attempting to compute BLUP contributions.

Parameters:

varying¶

Varying parameter extraction for mixed-effects models.

Extracts BLUP (Best Linear Unbiased Predictor) computation and variance component decomposition from the model class into pure functions on containers. These operations convert fitted spherical random effects into interpretable group-level parameters.

per_factor_re_info: Split global RE metadata into per-factor structures. per_factor_re_info: Split global RE metadata into per-factor structures. compute_varying_state: Compute BLUPs from theta and u via Lambda matrix. compute_varying_spread_state: Extract variance components (tau², rho, ICC).

Functions:

Attributes¶

Classes¶

Functions¶

build_mixed_post_fit_state¶

build_mixed_post_fit_state(fit: FitState, bundle: DataBundle, data: pl.DataFrame, *, stacklevel: int = 3) -> tuple[VaryingState | None, VaryingSpreadState | None]

Compute BLUPs, variance components, and emit convergence warnings.

Orchestrates the post-fit assembly for mixed-effects models: computes VaryingState (BLUPs) and VaryingSpreadState (variance components) from the fitted parameters, then checks for convergence issues.

Parameters:

Returns:

compute_varying_spread_state¶

compute_varying_spread_state(theta: NDArray[np.floating], sigma: float, re_meta: REInfo) -> VaryingSpreadState

Compute VaryingSpreadState (variance components) from theta parameters.

Extracts residual variance (sigma²), random effect variances (tau²), correlations (rho), and intraclass correlation (ICC) from the fitted theta vector using the random effects structure.

Parameters:

Returns:

compute_varying_state¶

compute_varying_state(theta: NDArray[np.floating], u: NDArray[np.floating], re_meta: REInfo, data: pl.DataFrame | None = None) -> VaryingState

Compute VaryingState (BLUPs) from fitted random effects parameters.

Converts spherical random effects u to BLUPs b = Lambda @ u using the relative covariance factor Lambda built from theta. Constructs a grid of group/level combinations and maps BLUP values to named effects.

Parameters:

Returns:

per_factor_re_info¶

per_factor_re_info(re_meta: REInfo, group_names: list[str]) -> tuple[str | list[str], list[str] | dict[str, list[str]]]

Split global RE metadata into per-factor structures and names.

For crossed/nested/mixed models, the global re_structure is a single string (e.g. “crossed”) and random_names is a concatenated list across all factors. This function splits them into per-factor structures and per-factor name dicts suitable for BLUP decomposition and convergence diagnostics.

For single-factor models, returns the originals unchanged.

Parameters:

Returns: