"""
Solver dispatch for mixed models.
Public API:
lmm() — fit a linear mixed model (REML or ML)
glmm() — fit a generalized linear mixed model (Laplace approximation)
"""
from __future__ import annotations
import warnings
import numpy as np
from numpy.typing import ArrayLike
from scipy import stats
from pystatistics.core.result import Result
from pystatistics.core.compute.timing import Timer
from pystatistics.core.exceptions import ValidationError
from pystatistics.mixed._common import (
LMMParams, GLMMParams, VarCompSummary,
)
from pystatistics.mixed._random_effects import (
parse_random_effects, build_z_matrix, build_lambda,
theta_lower_bounds, theta_start, is_singular_fit, theta_to_factor,
)
from pystatistics.mixed._pls import solve_pls
from pystatistics.mixed._pls_structured import (
build_structured_context, solve_structured,
)
from pystatistics.mixed._glmm_pirls import (
mode_joint, mode_given_beta,
profiled_deviance as glmm_profiled_deviance,
laplace_deviance as glmm_laplace_deviance,
)
from pystatistics.mixed._optimizer import optimize_theta
from pystatistics.mixed._glmm_optimizer import robust_minimize
from pystatistics.mixed._satterthwaite import satterthwaite_df
from pystatistics.mixed.design import MixedDesign
from pystatistics.mixed.solution import LMMSolution, GLMMSolution
[docs]
def lmm(
y: ArrayLike,
X: ArrayLike,
groups: dict[str, ArrayLike],
*,
random_effects: dict[str, list[str]] | None = None,
random_data: dict[str, ArrayLike] | None = None,
reml: bool = True,
tol: float = 1e-8,
max_iter: int = 200,
compute_satterthwaite: bool = True,
conf_level: float = 0.95,
) -> LMMSolution:
"""Fit a linear mixed model.
Estimates fixed effects β, random effects variance components,
and conditional modes (BLUPs) of random effects using the profiled
REML/ML deviance approach from Bates et al. (2015).
Args:
y: Response vector (n,).
X: Fixed effects design matrix (n, p). Should include an
intercept column if desired.
groups: Dict mapping grouping factor names to group label arrays.
Example: {'subject': subject_ids}.
random_effects: Optional dict mapping group names to lists of
random effect terms. Default: random intercept per group.
Example: {'subject': ['1', 'time']} for (1 + time | subject).
random_data: Optional dict mapping variable names to data arrays
for random slope variables.
Example: {'time': time_array}.
reml: If True (default), use REML estimation. If False, use ML.
Use ML (reml=False) for likelihood ratio tests between models
with different fixed effects.
tol: Convergence tolerance for the optimizer. Default 1e-8.
max_iter: Maximum optimizer iterations. Default 200.
compute_satterthwaite: If True (default), compute Satterthwaite
denominator df for fixed effects. Set to False for speed
if p-values are not needed.
Returns:
LMMSolution with fixed effects, random effects, variance components,
model fit statistics, and R-style summary().
Examples:
# Random intercept model
>>> result = lmm(y, X, groups={'subject': subject_ids})
# Random intercept + slope
>>> result = lmm(y, X, groups={'subject': subject_ids},
... random_effects={'subject': ['1', 'time']},
... random_data={'time': time_array})
# Crossed random effects
>>> result = lmm(y, X, groups={'subject': subj, 'item': item})
"""
if conf_level <= 0 or conf_level >= 1:
raise ValidationError(f"conf_level must be in (0, 1), got {conf_level}")
timer = Timer()
timer.start()
# Validate inputs
design = MixedDesign.validate(
np.asarray(y, dtype=np.float64),
np.asarray(X, dtype=np.float64),
groups,
random_effects,
random_data,
)
with timer.section('setup'):
# Parse random effects (structured LMM path: no dense Z — work from the
# compact value_cols + group_ids; build the reusable structured context).
specs = parse_random_effects(
design.groups, design.random_effects, design.random_data, design.n,
build_dense=False,
)
ctx = build_structured_context(design.X, design.y, specs, reml)
# Starting values and bounds
theta0 = theta_start(specs)
lb = theta_lower_bounds(specs)
bounds = [(lb[i], None) for i in range(len(theta0))]
# Optimize θ — use multiple starting points for models with random slopes
# (the profiled deviance can have local minima when q > 1)
with timer.section('optimization'):
has_slopes = any(s.n_terms > 1 for s in specs)
if has_slopes:
# Generate candidate starting values: the default [1,0,...,1]
# plus variants with smaller diagonal values for slope terms
starts = [theta0]
for scale in (0.2, 0.5):
alt = theta0.copy()
idx = 0
for spec in specs:
q = spec.n_terms
for row in range(q):
for col in range(row + 1):
if row == col and row > 0:
alt[idx] = scale
idx += 1
starts.append(alt)
else:
starts = [theta0]
opt_result = optimize_theta(ctx, starts, bounds, lb, max_iter, tol)
converged = opt_result.success
theta_hat = opt_result.x
n_iter = opt_result.nit
if not converged:
warnings.warn(
f"LMM optimizer did not converge after {n_iter} iterations. "
f"Message: {opt_result.message}",
RuntimeWarning,
stacklevel=2,
)
# Boundary (singular) fit diagnostic (mirrors lme4's isSingular).
is_singular = is_singular_fit(theta_hat, specs)
if is_singular:
warnings.warn(
"boundary (singular) fit: a random-effects variance is near zero "
"or a correlation is near ±1. The fit is on the boundary of the "
"feasible region; consider simplifying the random-effects "
"structure. (See LMMSolution.is_singular.)",
RuntimeWarning,
stacklevel=2,
)
# Final PLS solve at optimal θ (structured path)
with timer.section('final_solve'):
pls = solve_structured(theta_hat, ctx)
# Compute variance components
with timer.section('variance_components'):
var_comps = _extract_var_components(theta_hat, pls.sigma_sq, specs)
n_groups_dict = {s.group_name: s.n_groups for s in specs}
# Extract BLUPs
with timer.section('blups'):
random_effs = _extract_blups(pls.b, specs)
# Compute Satterthwaite df and p-values
with timer.section('satterthwaite'):
se = _compute_se(pls)
if compute_satterthwaite:
df_satt = satterthwaite_df(theta_hat, ctx)
else:
# Use residual df as fallback
df_satt = np.full(design.p, float(design.n - design.p))
t_vals = pls.beta / se
p_vals = 2.0 * stats.t.sf(np.abs(t_vals), df_satt)
# Log-likelihood, AIC, BIC
with timer.section('model_fit'):
ll, aic, bic = _compute_fit_stats(
pls, theta_hat, design.n, design.p, specs, reml
)
# Coefficient names
coef_names = _make_coef_names(design.p)
timer.stop()
# Assemble params
params = LMMParams(
coefficients=pls.beta,
coefficient_names=tuple(coef_names),
se=se,
df_satterthwaite=df_satt,
t_values=t_vals,
p_values=p_vals,
var_components=tuple(var_comps),
residual_variance=pls.sigma_sq,
residual_std=np.sqrt(pls.sigma_sq),
log_likelihood=ll,
reml=reml,
aic=aic,
bic=bic,
n_obs=design.n,
n_groups=n_groups_dict,
converged=converged,
n_iter=n_iter,
is_singular=is_singular,
random_effects=random_effs,
fitted_values=pls.fitted,
residuals=pls.residuals,
theta=theta_hat,
)
warn_list = []
if not converged:
warn_list.append(f"Optimizer did not converge: {opt_result.message}")
if is_singular:
warn_list.append("boundary (singular) fit")
result = Result(
params=params,
info={
'method': 'REML' if reml else 'ML',
'optimizer': 'L-BFGS-B',
'converged': converged,
'n_iter': n_iter,
'deviance': opt_result.fun,
},
timing=timer.result(),
backend_name='cpu_lmm',
warnings=tuple(warn_list),
)
return LMMSolution(_result=result, _conf_level=conf_level)
[docs]
def glmm(
y: ArrayLike,
X: ArrayLike,
groups: dict[str, ArrayLike],
*,
family: 'str | Family' = 'binomial',
random_effects: dict[str, list[str]] | None = None,
random_data: dict[str, ArrayLike] | None = None,
offset: ArrayLike | None = None,
weights: ArrayLike | None = None,
correlated: dict[str, bool] | bool | None = None,
tol: float = 1e-8,
max_iter: int = 200,
conf_level: float = 0.95,
) -> GLMMSolution:
"""Fit a generalized linear mixed model.
Uses the Laplace approximation to the marginal likelihood (equivalent to
``lme4::glmer(..., nAGQ=1)``, R's default): the random-effects modes are
profiled by Penalized IRLS (PIRLS) in an inner loop, while BOTH the variance
parameters θ and the fixed effects β are optimized in the outer loop
(L-BFGS-B) over the Laplace deviance. Optimizing β in the outer loop — rather
than solving it jointly inside PIRLS (the cruder ``nAGQ=0`` scheme) — is what
makes the fixed effects match ``glmer``'s Laplace fit rather than a biased
approximation.
Args:
y: Response vector (n,).
X: Fixed effects design matrix (n, p).
groups: Dict mapping grouping factor names to group label arrays.
family: GLM family specification. String ('binomial', 'poisson')
or a Family instance from pystatistics.regression.families.
random_effects: Optional random effects specification.
random_data: Optional data for random slope variables.
offset: Optional per-observation offset ``(n,)`` added to the linear
predictor, ``η = Xβ + Zb + offset`` (e.g. ``log(exposure)`` for a
Poisson rate model). Not estimated.
weights: Optional per-observation prior weights ``(n,)`` (IRLS prior
weights). For a proportion response these are the binomial trial
counts; a two-column ``y = [successes, failures]`` response supplies
them automatically (R's ``cbind(k, n-k)`` aggregated binomial).
correlated: Whether a grouping factor's random-effect terms share a full
covariance (the default, ``True``) or are uncorrelated with a
**diagonal** covariance (R's ``(… || g)``). Pass ``False`` to make all
factors diagonal, or a ``{group_name: bool}`` dict per factor. Only
affects factors with two or more terms.
tol: Convergence tolerance.
max_iter: Maximum optimizer iterations.
Returns:
GLMMSolution with fixed effects, random effects, and model fit.
"""
from pystatistics.regression.families import resolve_family, Family
if conf_level <= 0 or conf_level >= 1:
raise ValidationError(f"conf_level must be in (0, 1), got {conf_level}")
timer = Timer()
timer.start()
# Resolve family
if not isinstance(family, Family):
family_obj = resolve_family(family)
else:
family_obj = family
# glmm's Laplace likelihood fixes the dispersion at 1 (correct for binomial
# and Poisson). A free-dispersion family (Gaussian σ², Gamma shape) would get
# a silently-wrong log-likelihood / deviance / AIC / BIC and a biased variance
# component — fail loud instead (A6), pointing Gaussian users to lmm().
if not family_obj.dispersion_is_fixed:
raise ValidationError(
f"glmm() supports only fixed-dispersion families (binomial, poisson); "
f"family {family_obj.name!r} has a free dispersion/scale parameter "
f"that glmm's Laplace approximation does not estimate (it fixes "
f"dispersion=1), which would produce an incorrect log-likelihood, "
f"deviance, AIC/BIC and variance components. For a Gaussian mixed "
f"model use lmm().")
# Response / prior-weights / offset processing.
#
# Aggregated binomial: a 2-column response ``[successes, failures]`` (R's
# ``cbind(k, n-k)``) is converted to a proportion response with the trial
# counts as prior weights — exactly how glmer handles it.
y_arr = np.asarray(y, dtype=np.float64)
prior_w = None if weights is None else np.asarray(weights, dtype=np.float64)
if y_arr.ndim == 2 and y_arr.shape[1] == 2:
if family_obj.name != 'binomial':
raise ValidationError(
"A two-column [successes, failures] response is only valid for "
f"family='binomial', not {family_obj.name!r}.")
trials = y_arr[:, 0] + y_arr[:, 1]
if np.any(trials <= 0):
raise ValidationError(
"Aggregated binomial: every row needs a positive trial count "
"(successes + failures > 0).")
y_arr = y_arr[:, 0] / trials
prior_w = trials if prior_w is None else prior_w * trials
elif y_arr.ndim != 1:
raise ValidationError(
f"y must be 1-D (or a 2-column binomial response), got shape "
f"{y_arr.shape}.")
offset_arr = None if offset is None else np.asarray(offset, dtype=np.float64)
# Validate inputs
design = MixedDesign.validate(
y_arr,
np.asarray(X, dtype=np.float64),
groups,
random_effects,
random_data,
)
if prior_w is not None and prior_w.shape != (design.n,):
raise ValidationError(
f"weights must have length {design.n}, got {prior_w.shape}.")
if offset_arr is not None and offset_arr.shape != (design.n,):
raise ValidationError(
f"offset must have length {design.n}, got {offset_arr.shape}.")
with timer.section('setup'):
specs = parse_random_effects(
design.groups, design.random_effects, design.random_data, design.n,
build_dense=False, correlated=correlated,
)
# Structure-exploiting context (batched per-group for a single grouping
# factor; sparse for crossed/nested) — the GLMM PIRLS solves through it
# so the cost scales with the block/sparsity structure, not O(#groups³).
ctx = build_structured_context(
design.X, design.y, specs, reml=False,
offset=offset_arr, prior_weights=prior_w,
)
theta0 = theta_start(specs)
lb = theta_lower_bounds(specs)
bounds = [(lb[i], None) for i in range(len(theta0))]
p = design.X.shape[1]
n_theta = len(theta0)
with timer.section('optimization'):
# Stage 1 (nAGQ=0 warm start): optimize θ only, with β solved jointly
# inside PIRLS. This gives a good (θ₀, β₀) to seed the Laplace stage —
# exactly how lme4::glmer bootstraps its nAGQ=1 optimization. A
# derivative-free fallback rescues the case where L-BFGS-B overshoots an
# interior variance optimum to the θ=0 boundary (a silent collapse).
theta_start_2, _dev1, _conv1 = robust_minimize(
lambda th: glmm_profiled_deviance(ctx, th, family_obj),
theta0, bounds, max_iter=max_iter, tol=tol)
beta_start = mode_joint(ctx, theta_start_2, family_obj).beta
# Stage 2 (Laplace / nAGQ=1): optimize BOTH θ and β over the Laplace
# deviance, profiling only the random-effects modes u via PIRLS, with the
# same robust fallback on the joint objective.
params0 = np.concatenate([theta_start_2, beta_start])
param_bounds = bounds + [(None, None)] * p # β is unbounded
best_x, final_dev, converged = robust_minimize(
lambda pr: glmm_laplace_deviance(ctx, pr, family_obj, n_theta),
params0, param_bounds, max_iter=max_iter, tol=tol)
theta_hat = best_x[:n_theta]
beta_hat = best_x[n_theta:]
if not converged:
warnings.warn(
"GLMM optimizer did not converge (L-BFGS-B and the derivative-free "
"fallback both failed to satisfy the tolerance). Treat the estimates "
"with caution.",
RuntimeWarning,
stacklevel=2,
)
# Boundary (singular) fit diagnostic (mirrors lme4's isSingular). The test
# is on θ only, so it is identical to the LMM path.
is_singular = is_singular_fit(theta_hat, specs)
if is_singular:
warnings.warn(
"boundary (singular) fit: a random-effects variance is near zero "
"or a correlation is near ±1. The fit is on the boundary of the "
"feasible region; consider simplifying the random-effects "
"structure. (See GLMMSolution.is_singular.)",
RuntimeWarning,
stacklevel=2,
)
# Final inner solve at (θ̂, β̂): the conditional modes û, the fitted mean,
# the Laplace log-det and the p×p Schur factor RX the fixed-effect
# covariance needs — all from the structured factor at the mode.
with timer.section('final_solve'):
mode = mode_given_beta(ctx, theta_hat, beta_hat, family_obj)
n_iter = int(mode.n_iter) # inner PIRLS iterations at the optimum
# Variance components (for GLMM, σ² = 1 by convention)
with timer.section('variance_components'):
var_comps = _extract_var_components(theta_hat, 1.0, specs)
n_groups_dict = {s.group_name: s.n_groups for s in specs}
# BLUPs (conditional modes b = Λû)
with timer.section('blups'):
random_effs = _extract_blups(mode.b, specs)
# Fixed effect SEs and Wald z-statistics. The Laplace fixed-effect
# covariance is the inverse of the p×p Schur complement of the penalized
# system at the mode (the RX factor lme4 stores), carried on the mode result.
with timer.section('inference'):
se = _compute_se_glmm(mode, p)
z_vals = beta_hat / se
p_vals = 2.0 * stats.norm.sf(np.abs(z_vals))
# Model fit
with timer.section('model_fit'):
wt = np.ones(design.n, dtype=np.float64) if prior_w is None else prior_w
deviance = family_obj.deviance(design.y, mode.mu, wt)
n_params = p + len(theta_hat)
# Laplace-approximated marginal log-likelihood:
# ll = conditional_loglik - 0.5 * ||u||^2 - 0.5 * log|L_theta|^2
# where conditional_loglik = sum(f(y_i | mu_i)) is the full
# conditional log-likelihood including normalizing constants.
# For GLMM, dispersion = 1.
cond_ll = family_obj.log_likelihood(design.y, mode.mu, wt, 1.0)
# log|L_θ|² = log|Λ'Z'WZΛ + I|, carried on the structured mode result.
ll = cond_ll - 0.5 * mode.penalty - 0.5 * mode.logdet_M
aic = -2.0 * ll + 2.0 * n_params
bic = -2.0 * ll + np.log(design.n) * n_params
coef_names = _make_coef_names(design.p)
timer.stop()
params = GLMMParams(
coefficients=beta_hat,
coefficient_names=tuple(coef_names),
se=se,
z_values=z_vals,
p_values=p_vals,
var_components=tuple(var_comps),
log_likelihood=ll,
deviance=deviance,
aic=aic,
bic=bic,
n_obs=design.n,
n_groups=n_groups_dict,
family_name=family_obj.name,
link_name=family_obj.link.name,
converged=converged,
n_iter=n_iter,
random_effects=random_effs,
fitted_values=mode.mu,
linear_predictor=mode.eta,
residuals=design.y - mode.mu,
theta=theta_hat,
is_singular=is_singular,
)
warn_list = []
if not converged:
warn_list.append("Optimizer did not converge (L-BFGS-B + fallback)")
if not mode.converged:
warn_list.append(f"PIRLS did not converge after {mode.n_iter} iterations")
if is_singular:
warn_list.append("boundary (singular) fit")
result = Result(
params=params,
info={
'method': 'Laplace',
'family': family_obj.name,
'link': family_obj.link.name,
'optimizer': 'L-BFGS-B+NelderMead',
'converged': converged,
'pirls_converged': mode.converged,
'n_iter': n_iter,
'pirls_iter': mode.n_iter,
'deviance': float(final_dev),
},
timing=timer.result(),
backend_name='cpu_glmm',
warnings=tuple(warn_list),
)
return GLMMSolution(_result=result, _conf_level=conf_level)
# =====================================================================
# Helpers
# =====================================================================
def _extract_var_components(
theta: np.ndarray,
sigma_sq: float,
specs: list,
) -> list[VarCompSummary]:
"""Extract variance component summaries from θ and σ².
The actual covariance of random effects is σ² × Λ Λ'.
For each grouping factor, compute the covariance matrix and
extract variance, std dev, and correlations.
"""
var_comps = []
theta_start = 0
for spec in specs:
q = spec.n_terms
n_theta = spec.theta_size
# Reconstruct the q × q Cholesky factor (lower-triangular, or diagonal
# when the factor is uncorrelated — off-diagonals are then exactly 0).
theta_k = theta[theta_start:theta_start + n_theta]
theta_start += n_theta
T = theta_to_factor(theta_k, spec)
# Covariance matrix: σ² × T T'
cov_matrix = sigma_sq * (T @ T.T)
# Term names
term_names = []
for term in spec.terms:
if term == '1':
term_names.append('(Intercept)')
else:
term_names.append(term)
# Extract variance, std dev, correlations
for i in range(q):
var_i = cov_matrix[i, i]
sd_i = np.sqrt(max(var_i, 0.0))
# Correlation with first term (only for 2nd+ terms)
if i > 0 and cov_matrix[0, 0] > 0 and var_i > 0:
corr = cov_matrix[i, 0] / (np.sqrt(cov_matrix[0, 0]) * sd_i)
# NUMERICAL GUARD: floating-point arithmetic can produce |r| > 1
corr = np.clip(corr, -1.0, 1.0)
else:
corr = None
var_comps.append(VarCompSummary(
group=spec.group_name,
name=term_names[i],
variance=float(var_i),
std_dev=float(sd_i),
corr=float(corr) if corr is not None else None,
))
return var_comps
def _extract_blups(b: np.ndarray, specs: list) -> dict[str, np.ndarray]:
"""Extract BLUPs per grouping factor from the flat b vector.
b is structured as [b_group1, b_group2, ...] where each b_groupk
has J_k * q_k elements laid out as [term0_group0, term0_group1, ...,
term1_group0, ...].
Returns dict: group_name → (J_k, q_k) array.
"""
result = {}
block_start = 0
for spec in specs:
block_size = spec.n_groups * spec.n_terms
b_block = b[block_start:block_start + block_size]
# Reshape: columns are terms, rows are groups
# b_block layout: [term0_g0, term0_g1, ..., term1_g0, term1_g1, ...]
b_matrix = np.zeros((spec.n_groups, spec.n_terms), dtype=np.float64)
for t in range(spec.n_terms):
start = t * spec.n_groups
b_matrix[:, t] = b_block[start:start + spec.n_groups]
result[spec.group_name] = b_matrix
block_start += block_size
return result
def _compute_se(pls) -> np.ndarray:
"""Compute standard errors of fixed effects from the p×p Schur factor.
SE = sqrt(diag(Var(β̂))) where Var(β̂) = σ² × (X'V*⁻¹X)⁻¹ and
V* = ZΛΛ'Z' + I. By the Woodbury identity, X'V*⁻¹X = RX·RXᵀ where RX
is the lower-Cholesky of the Schur complement already returned by the
PLS solve, so
Var(β̂) = σ² × (RX·RXᵀ)⁻¹ = σ² × RX⁻ᵀ RX⁻¹ = σ² × (RX⁻¹)ᵀ(RX⁻¹).
This is "form A" — it uses only the p×p factor, never the n×n V* solve
the previous implementation formed (an O(n³) hot spot, ~83% of lmm()
runtime at scale). It is machine-identical (~1e-14) to that dense path
and thousands of times faster.
Both the LMM and GLMM paths use this SAME ``RX⁻¹ᵀ·RX⁻¹`` form — ``RX`` is
the lower Cholesky of the Schur complement in each case (both come from
:func:`solve_pls`), so Var(β̂) = (RX·RXᵀ)⁻¹ = RX⁻ᵀ RX⁻¹.
"""
try:
RX_inv = np.linalg.inv(pls.RX)
vcov = pls.sigma_sq * (RX_inv.T @ RX_inv)
except np.linalg.LinAlgError:
# Singular RX (collinear fixed effects): fall back to a
# pseudo-inverse of RX·RXᵀ = X'V*⁻¹X.
RtR = pls.RX @ pls.RX.T
vcov = pls.sigma_sq * np.linalg.pinv(RtR)
se = np.sqrt(np.maximum(np.diag(vcov), 0.0))
return se
def _compute_se_glmm(pls, p: int) -> np.ndarray:
"""Compute fixed-effect standard errors for GLMM (σ² = 1 by convention).
The Laplace fixed-effect covariance is the inverse of the p×p Schur
complement of the penalized system at the mode. ``pls.RX`` is the LOWER
Cholesky factor of that Schur complement S (so S = RX·RXᵀ), hence
Var(β̂) = S⁻¹ = (RX·RXᵀ)⁻¹ = RX⁻ᵀ RX⁻¹ = (RX⁻¹)ᵀ (RX⁻¹).
This is the SAME convention as the LMM :func:`_compute_se` (``pls`` comes
from the same :func:`solve_pls`). The previous ``RX⁻¹·RX⁻¹ᵀ`` form was
wrong: it inverts RXᵀRX instead of RX·RXᵀ, which agrees only when the
fixed-effect design is (weighted-)orthogonal and is badly wrong for
correlated predictors (validated against lme4::glmer).
"""
try:
RX_inv = np.linalg.inv(pls.RX)
vcov = RX_inv.T @ RX_inv
except np.linalg.LinAlgError:
# Singular RX (collinear fixed effects): pseudo-inverse of RX·RXᵀ.
RtR = pls.RX @ pls.RX.T
vcov = np.linalg.pinv(RtR)
se = np.sqrt(np.maximum(np.diag(vcov), 0.0))
return se
def _compute_fit_stats(pls, theta, n, p, specs, reml):
"""Compute log-likelihood, AIC, BIC for LMM."""
sigma_sq = pls.sigma_sq
pwrss = pls.pwrss
# Number of variance parameters
n_theta = len(theta)
# Total parameters for AIC: fixed effects + variance components + σ²
n_params = p + n_theta + 1
# log|L|² = log|Λ'Z'ZΛ + I| comes directly from the structured solve.
log_det_L = pls.logdet_M
if reml:
df = n - p
# REML log-likelihood
# NUMERICAL GUARD: prevents log(0) in log-likelihood computation
log_det_RX = 2.0 * np.sum(np.log(np.maximum(np.abs(np.diag(pls.RX)), 1e-20)))
ll = -0.5 * (
log_det_L
+ log_det_RX
+ df * (1.0 + np.log(2.0 * np.pi * pwrss / df))
)
# REML AIC/BIC: R's lme4 counts all parameters (fixed + variance)
# npar = p (fixed effects) + n_theta (RE params) + 1 (sigma)
aic = -2.0 * ll + 2.0 * n_params
bic = -2.0 * ll + np.log(n) * n_params
else:
# ML log-likelihood
ll = -0.5 * (
log_det_L
+ n * (1.0 + np.log(2.0 * np.pi * pwrss / n))
)
aic = -2.0 * ll + 2.0 * n_params
bic = -2.0 * ll + np.log(n) * n_params
return float(ll), float(aic), float(bic)
def _make_coef_names(p: int) -> list[str]:
"""Generate default coefficient names."""
if p == 1:
return ['(Intercept)']
names = ['(Intercept)']
for i in range(1, p):
names.append(f'X{i}')
return names