GAMM — Generalized Additive Mixed Models

Generalized Additive Mixed Model routines.

class aurora.models.gamm.RandomEffect(grouping, variables=<factory>, include_intercept=True, covariance='unstructured')

Bases: object

Specification of a random effect term.

Random effects model within-group correlation in hierarchical or longitudinal data. Each random effect term specifies: - Which variables have random coefficients - What grouping structure to use - Whether to include a random intercept - The covariance structure among random effects

Parameters:

• grouping (str or int) – Grouping variable name or column index. All observations with the same value of this variable belong to the same group.

• variables (tuple of (str or int), optional) – Variables with random slopes. If empty, only includes random intercept (if include_intercept=True).

• include_intercept (bool, default=True) – Whether to include a random intercept (1 | group).

• covariance (str, default='unstructured') –

  Covariance structure among random effects. Options: - ‘unstructured’: Full covariance matrix (q(q+1)/2 parameters) - ‘diagonal’: Independent random effects (q parameters) - ‘identity’: Equal variances, no correlation (1 parameter) - ‘ar1’: Autoregressive AR(1) for temporal correlation - ‘compound_symmetry’ or ‘cs’: Exchangeable correlation (ICC model) - ‘exponential’: Spatial exponential decay (requires coordinates) - ‘matern’: Matérn spatial covariance (requires coordinates)

  For temporal/longitudinal data, use ‘ar1’. For clustered data with equal within-cluster correlation, use ‘compound_symmetry’.

Variables:

n_effects (int) – Total number of random effects (intercept + slopes)

Examples

>>> # Random intercept: (1 | subject)
>>> re1 = RandomEffect(grouping='subject')
>>> re1.n_effects
1

>>> # Random intercept + slope: (1 + time | subject)
>>> re2 = RandomEffect(grouping='subject', variables=('time',))
>>> re2.n_effects
2

>>> # Random slopes only: (time + age | subject)
>>> re3 = RandomEffect(
...     grouping='subject',
...     variables=('time', 'age'),
...     include_intercept=False
... )
>>> re3.n_effects
2

>>> # Diagonal covariance (independent effects)
>>> re4 = RandomEffect(
...     grouping='subject',
...     variables=('time',),
...     covariance='diagonal'
... )

Notes

The total number of random effects is:

q = (1 if include_intercept else 0) + len(variables)

Nested random effects like (1 | clinic/subject) should be specified as two separate RandomEffect instances with appropriate grouping.

Crossed random effects like (1 | subject) + (1 | clinic) are specified as multiple RandomEffect instances in a list.

grouping: str | int

variables: tuple[str | int, ...]

include_intercept: bool = True

covariance: Literal['unstructured', 'diagonal', 'identity', 'ar1', 'compound_symmetry', 'cs', 'exponential', 'matern'] = 'unstructured'

__post_init__()

Validate random effect specification.

property n_effects: int

Number of random effects in this term.

Returns:

Number of random effects (intercept + slopes)

Return type:

int

__repr__()

String representation.

Return type:

str

aurora.models.gamm.validate_random_effects(random_effects)

Validate a list of random effects specifications.

Parameters:

random_effects (list of RandomEffect) – Random effects to validate

Raises:

• TypeError – If random_effects is not a list or contains non-RandomEffect items

• ValueError – If random effects specifications are invalid

Return type:

None

Examples

>>> re1 = RandomEffect(grouping='subject')
>>> re2 = RandomEffect(grouping='clinic')
>>> validate_random_effects([re1, re2])  # OK

>>> validate_random_effects([re1, "invalid"])  # Raises TypeError
Traceback (most recent call last):
    ...
TypeError: All items in random_effects must be RandomEffect instances

Notes

This function checks: - All items are RandomEffect instances - Each RandomEffect is internally valid (via __post_init__)

aurora.models.gamm.get_group_indices(groups)

Get unique groups and indices for each group.

Parameters:

groups (ndarray, shape (n,)) – Group indicators (integers)

Returns:

• unique_groups (ndarray) – Sorted array of unique group values

• group_indices (dict) – Dictionary mapping group_id -> array of observation indices

Return type:

tuple[ndarray, dict[int, ndarray]]

Examples

>>> groups = np.array([1, 1, 2, 2, 2, 3])
>>> unique, indices = get_group_indices(groups)
>>> unique
array([1, 2, 3])
>>> indices[1]
array([0, 1])
>>> indices[2]
array([2, 3, 4])

aurora.models.gamm.count_random_effects(random_effects)

Count total number of random effect coefficients.

Parameters:

random_effects (list of RandomEffect) – Random effects specifications

Returns:

• total_effects (int) – Total number of random effect coefficients across all terms

• effects_per_term (dict) – Dictionary mapping grouping variable -> number of effects

Return type:

tuple[int, dict[str | int, int]]

Examples

>>> re1 = RandomEffect(grouping='subject', variables=('time',))
>>> re2 = RandomEffect(grouping='clinic')
>>> total, per_term = count_random_effects([re1, re2])
>>> total
3
>>> per_term['subject']
2
>>> per_term['clinic']
1

Notes

For each random effect term with q effects and G groups, we have q * G random coefficients in total.

class aurora.models.gamm.CovarianceStructure

Bases: ABC

Base class for random effect covariance structures.

A covariance structure defines how to parameterize the variance-covariance matrix Ψ for random effects. Different structures impose different constraints (e.g., diagonal, identity) to reduce the number of parameters.

n_parameters(n_effects)

Number of free parameters for given number of random effects

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Build covariance matrix from parameter vector

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract parameter vector from covariance matrix

Parameters:

psi (ndarray)

Return type:

ndarray

abstractmethod n_parameters(n_effects)

Number of free parameters.

Parameters:

n_effects (int) – Number of random effects (dimension of Ψ)

Returns:

Number of free parameters to specify Ψ

Return type:

int

abstractmethod construct_psi(params, n_effects)

Construct covariance matrix from parameters.

Parameters:

• params (ndarray) – Parameter vector (length = n_parameters(n_effects))

• n_effects (int) – Number of random effects

Returns:

psi – Variance-covariance matrix (symmetric, positive definite)

Return type:

ndarray, shape (n_effects, n_effects)

abstractmethod extract_params(psi)

Extract parameters from covariance matrix.

Parameters:

psi (ndarray, shape (q, q)) – Variance-covariance matrix

Returns:

params – Parameter vector

Return type:

ndarray

class aurora.models.gamm.UnstructuredCovariance

Bases: CovarianceStructure

Unstructured covariance (full matrix).

Allows all variances and covariances to vary freely. For q random effects, requires q(q+1)/2 parameters.

Parameterization uses Cholesky decomposition: Ψ = LL’ This ensures Ψ is positive definite.

Parameters:

None

Examples

>>> cov = UnstructuredCovariance()
>>> cov.n_parameters(2)  # 2 random effects
3

>>> # Parameters: [L11, L21, L22]
>>> params = np.array([1.0, 0.5, 0.8])
>>> psi = cov.construct_psi(params, n_effects=2)
>>> psi
array([[1.  , 0.5 ],
       [0.5 , 0.89]])

>>> # Extract back
>>> extracted = cov.extract_params(psi)
>>> np.allclose(extracted, params)
True

Notes

The Cholesky parameterization: - Ensures positive definiteness - Provides unconstrained optimization (except L_ii > 0) - Standard approach in lme4 and other mixed model software

n_parameters(n_effects)

Number of parameters = q(q+1)/2.

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Ψ = LL’ from Cholesky parameters.

Parameters are ordered as lower triangle of L: [L11, L21, L22, L31, L32, L33, …]

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract Cholesky parameters from Ψ.

Computes L such that Ψ = LL’, then extracts lower triangle.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.DiagonalCovariance

Bases: CovarianceStructure

Diagonal covariance (independent random effects).

Assumes random effects are independent (zero correlations). For q random effects, requires q parameters (variances only).

Parameters are log-transformed variances to ensure positivity: σ²_i = exp(params[i])

Examples

>>> cov = DiagonalCovariance()
>>> cov.n_parameters(3)
3

>>> # Parameters: log-variances
>>> params = np.array([0.0, 0.5, 1.0])  # log(σ²)
>>> psi = cov.construct_psi(params, n_effects=3)
>>> psi
array([[1.        , 0.        , 0.        ],
       [0.        , 1.64872127, 0.        ],
       [0.        , 0.        , 2.71828183]])

>>> np.diag(psi)  # Variances
array([1.        , 1.64872127, 2.71828183])

Notes

Log transformation: - Ensures variances are positive - Provides unconstrained optimization - Standard in many mixed model packages

n_parameters(n_effects)

Number of parameters = q (one variance per effect).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct diagonal Ψ from log-variances.

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract log-variances from diagonal Ψ.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.IdentityCovariance

Bases: CovarianceStructure

Identity covariance (equal variances, no correlation).

All random effects have the same variance and are uncorrelated. For q random effects, requires only 1 parameter: σ²

Ψ = σ² I_q

Parameter is log(σ²) to ensure positivity.

Examples

>>> cov = IdentityCovariance()
>>> cov.n_parameters(3)
1

>>> # Single parameter: log(σ²)
>>> params = np.array([0.5])  # log(σ²)
>>> psi = cov.construct_psi(params, n_effects=3)
>>> psi
array([[1.64872127, 0.        , 0.        ],
       [0.        , 1.64872127, 0.        ],
       [0.        , 0.        , 1.64872127]])

>>> # All variances are equal
>>> np.allclose(np.diag(psi), np.exp(0.5))
True

Notes

This is the most constrained structure, useful when: - Sample size is small - Random effects are believed to be exchangeable - Computational efficiency is important

n_parameters(n_effects)

Number of parameters = 1 (single variance).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct σ²I from log-variance.

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract log(σ²) from Ψ = σ²I.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.AR1Covariance

Bases: CovarianceStructure

Autoregressive order 1 (AR(1)) covariance structure.

Models temporal correlation where observations closer in time are more correlated: Cov(b_t, b_s) = σ² ρ^|t-s|

For q random effects (time points), requires 2 parameters: - params[0]: log(σ²) - log-variance - params[1]: arctanh(ρ) - transformed correlation (ensures ρ ∈ (-1, 1))

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS. Ch. 5.

[2]

Diggle et al. (2002). Analysis of Longitudinal Data, 2nd ed.

Examples

>>> cov = AR1Covariance()
>>> cov.n_parameters(4)
2

>>> # Parameters: [log(σ²), arctanh(ρ)]
>>> params = np.array([0.0, 0.5])  # σ²=1, ρ≈0.46
>>> psi = cov.construct_psi(params, n_effects=4)
>>> # Creates correlation matrix with AR(1) structure

Notes

The AR(1) structure is ideal for: - Equally-spaced longitudinal data - Time series with decaying correlations - Repeated measures where adjacent observations are most correlated

The inverse of AR(1) covariance matrix is tridiagonal, enabling efficient computation for large matrices.

n_parameters(n_effects)

Number of parameters = 2 (variance and correlation).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct AR(1) covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), arctanh(ρ)] - transformed variance and correlation

• n_effects (int) – Number of time points/random effects

Returns:

psi – AR(1) covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from AR(1) covariance matrix.

Estimates σ² from diagonal and ρ from first off-diagonal.

Parameters:

psi (ndarray)

Return type:

ndarray

inverse(params, n_effects)

Efficient tridiagonal inverse for AR(1) structure.

The inverse of an AR(1) matrix is tridiagonal, which allows for O(n) computation instead of O(n³).

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

class aurora.models.gamm.CompoundSymmetryCovariance

Bases: CovarianceStructure

Compound symmetry (exchangeable) covariance structure.

Models constant correlation between all pairs of observations: - Diagonal: Var(b_i) = σ² - Off-diagonal: Cov(b_i, b_j) = σ²ρ for i ≠ j

For q random effects, requires 2 parameters: - params[0]: log(σ²) - log-variance - params[1]: logit(ρ_scaled) - transformed correlation

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS.

[2]

Diggle et al. (2002). Analysis of Longitudinal Data, 2nd ed.

Examples

>>> cov = CompoundSymmetryCovariance()
>>> cov.n_parameters(5)
2

>>> params = np.array([0.0, 0.0])  # σ²=1, ρ=0.5
>>> psi = cov.construct_psi(params, n_effects=3)
>>> # Creates matrix with equal variances and equal correlations

Notes

Compound symmetry is equivalent to: - Random intercept model with iid residuals - Exchangeable correlation (ICC model)

Constraint: ρ > -1/(q-1) for positive definiteness, where q is the number of effects.

n_parameters(n_effects)

Number of parameters = 2 (variance and correlation).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct compound symmetry covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), logit((ρ + 1/(q-1))/(1 + 1/(q-1)))]

• n_effects (int) – Number of random effects

Returns:

psi – Compound symmetry covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from compound symmetry matrix.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.ExponentialSpatialCovariance(coordinates=None)

Bases: CovarianceStructure

Exponential spatial covariance structure.

Models spatial correlation that decays exponentially with distance: Cov(b_i, b_j) = σ² exp(-d_ij / φ)

where d_ij is the Euclidean distance between locations i and j, and φ is the range parameter.

For geostatistical data, requires coordinates at construction time. Parameters: - params[0]: log(σ²) - log-variance (sill) - params[1]: log(φ) - log-range parameter

References

[1]

Diggle & Ribeiro (2007). Model-based Geostatistics.

[2]

Cressie (1993). Statistics for Spatial Data.

[3]

Rue & Held (2005). Gaussian Markov Random Fields.

Examples

>>> coords = np.array([[0, 0], [1, 0], [0, 1], [1, 1]])
>>> cov = ExponentialSpatialCovariance(coordinates=coords)
>>> params = np.array([0.0, 0.5])  # σ²=1, φ≈1.65
>>> psi = cov.construct_psi(params, n_effects=4)

Notes

The exponential covariance is a member of the Matérn family with smoothness ν = 0.5. It produces relatively rough spatial surfaces.

The range parameter φ approximately equals the distance at which the correlation drops to ~37% (= 1/e).

Practical range (correlation ≈ 5%) is approximately 3φ.

Parameters:

coordinates (np.ndarray | None)

__init__(coordinates=None)

Initialize with spatial coordinates.

Parameters:

coordinates (ndarray (n, d), optional) – Spatial coordinates for n locations in d dimensions. If not provided, assumes 1D equally-spaced locations.

n_parameters(n_effects)

Number of parameters = 2 (variance and range).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct exponential spatial covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), log(φ)] - log-variance and log-range

• n_effects (int) – Number of spatial locations

Returns:

psi – Spatial covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from spatial covariance matrix.

Uses method of moments estimation based on variogram fitting.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.MaternCovariance(coordinates=None, nu=1.5)

Bases: CovarianceStructure

Matérn covariance structure for spatial data.

The Matérn family is the most commonly used in geostatistics due to its flexibility in modeling different degrees of smoothness.

Cov(d) = σ² × (2^(1-ν) / Γ(ν)) × (√(2ν) d/φ)^ν × K_ν(√(2ν) d/φ)

where K_ν is the modified Bessel function of the second kind.

Special cases: - ν = 0.5: Exponential covariance - ν = 1.5: Once differentiable - ν = 2.5: Twice differentiable - ν → ∞: Gaussian (squared exponential)

Parameters (for fixed ν): - params[0]: log(σ²) - log-variance (sill) - params[1]: log(φ) - log-range parameter

References

[1]

Matérn (1960). Spatial Variation.

[2]

Stein (1999). Interpolation of Spatial Data.

[3]

Rasmussen & Williams (2006). Gaussian Processes for ML.

Examples

>>> coords = np.array([[0, 0], [1, 0], [2, 0]])
>>> cov = MaternCovariance(coordinates=coords, nu=1.5)
>>> params = np.array([0.0, 0.0])  # σ²=1, φ=1
>>> psi = cov.construct_psi(params, n_effects=3)

Parameters:

• coordinates (np.ndarray | None)

• nu (float)

__init__(coordinates=None, nu=1.5)

Initialize Matérn covariance.

Parameters:

• coordinates (ndarray (n, d), optional) – Spatial coordinates. If None, assumes 1D equally-spaced.

• nu (float, default=1.5) – Smoothness parameter. Common choices: 0.5, 1.5, 2.5

n_parameters(n_effects)

Number of parameters = 2 (variance and range).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Matérn covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), log(φ)]

• n_effects (int) – Number of spatial locations

Returns:

psi – Matérn covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from Matérn covariance matrix.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.ToeplitzCovariance(band=2)

Bases: CovarianceStructure

Toeplitz (banded) covariance structure for temporal data.

Toeplitz matrices have constant values along diagonals, meaning the covariance between observations depends only on their lag |t-s|:

Cov(b_t, b_s) = r_{|t-s|}

This is a generalization of AR(1) that allows arbitrary values at each lag, rather than assuming exponential decay.

For q random effects (time points), requires (band+1) parameters where band is the number of non-zero off-diagonal bands: - params[0]: log(r_0) - log-variance - params[k]: arctanh(r_k / r_0) for k = 1, …, band (transformed correlations)

Parameters:

band (int, default=2) – Number of off-diagonal bands to include. Higher values allow more complex temporal patterns but require more parameters. - band=1: First-order correlations only (like AR(1) but with free parameter) - band=2: First and second-order correlations

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS. Ch. 5.

[2]

Pourahmadi (1999). “Joint mean-covariance models with applications to longitudinal data.” Biometrika, 86(3), 677-690.

Examples

>>> cov = ToeplitzCovariance(band=2)
>>> cov.n_parameters(5)
3

>>> params = np.array([0.0, 0.3, 0.1])  # log(σ²)=0, r_1/σ²≈0.29, r_2/σ²≈0.10
>>> psi = cov.construct_psi(params, n_effects=4)

Notes

Toeplitz covariance is appropriate for: - Equally-spaced longitudinal data - Time series where correlations don’t follow AR(p) exactly - Exploratory analysis before fitting structured models

The Toeplitz structure is a special case of a block structure and allows efficient computation using banded matrix algorithms.

__init__(band=2)

Initialize Toeplitz covariance.

Parameters:

band (int, default=2) – Number of off-diagonal bands (lags) to model.

n_parameters(n_effects)

Number of parameters = min(band, n_effects-1) + 1.

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Toeplitz covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), arctanh(ρ_1), arctanh(ρ_2), …] - transformed parameters

• n_effects (int) – Number of time points/random effects

Returns:

psi – Toeplitz covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from Toeplitz covariance matrix.

Estimates variance from diagonal and correlations from off-diagonals.

Parameters:

psi (ndarray)

Return type:

ndarray

aurora.models.gamm.get_covariance_structure(structure, **kwargs)

Get covariance structure instance by name.

Parameters:

• structure ({'unstructured', 'diagonal', 'identity', 'ar1',) – ‘compound_symmetry’, ‘cs’, ‘exponential’, ‘matern’} Covariance structure name

• **kwargs (dict) – Additional arguments for specific structures: - coordinates: ndarray for spatial structures - nu: smoothness for Matérn

Returns:

Covariance structure instance

Return type:

CovarianceStructure

Raises:

ValueError – If structure name is not recognized

Examples

>>> cov = get_covariance_structure('unstructured')
>>> isinstance(cov, UnstructuredCovariance)
True

>>> cov = get_covariance_structure('ar1')
>>> isinstance(cov, AR1Covariance)
True

>>> coords = np.array([[0, 0], [1, 0], [0, 1]])
>>> cov = get_covariance_structure('exponential', coordinates=coords)
>>> isinstance(cov, ExponentialSpatialCovariance)
True

>>> cov = get_covariance_structure('matern', coordinates=coords, nu=2.5)
>>> cov.nu
2.5

aurora.models.gamm.construct_Z_matrix(X, random_effects, groups_data)

Construct random effects design matrix Z.

For each random effect term, Z has a block for each group. Within each block, columns correspond to the random effects (intercept, slopes) for that group.

Parameters:

• X (ndarray, shape (n, p)) – Fixed effects design matrix (including intercept if needed)

• random_effects (list of RandomEffect) – Random effect specifications

• groups_data (dict) – Dictionary mapping grouping variable -> group indicators array. Keys should match RandomEffect.grouping values.

Returns:

• Z (ndarray, shape (n, q)) – Random effects design matrix where q = sum over terms of (n_groups_i * n_effects_i)

• Z_info (list of dict) – List of dictionaries (one per random effect term) containing: - ‘grouping’: grouping variable - ‘n_effects’: number of random effects per group - ‘n_groups’: number of groups - ‘groups’: array of unique group values - ‘start_col’: starting column in Z - ‘end_col’: ending column in Z (exclusive)

Return type:

tuple[ndarray, list[dict]]

Examples

>>> n = 6
>>> X = np.ones((n, 1))  # Intercept only
>>> groups = np.array([1, 1, 2, 2, 3, 3])
>>> groups_data = {'subject': groups}
>>>
>>> # Random intercept: (1 | subject)
>>> re = RandomEffect(grouping='subject')
>>> Z, info = construct_Z_matrix(X, [re], groups_data)
>>> Z.shape
(6, 3)  # 3 groups, 1 effect per group
>>> Z
array([[1., 0., 0.],
       [1., 0., 0.],
       [0., 1., 0.],
       [0., 1., 0.],
       [0., 0., 1.],
       [0., 0., 1.]])

>>> # Random intercept + slope: (1 + x1 | subject)
>>> X_with_x1 = np.column_stack([np.ones(n), np.arange(n)])
>>> re = RandomEffect(grouping='subject', variables=(1,))  # x1 is column 1
>>> Z, info = construct_Z_matrix(X_with_x1, [re], groups_data)
>>> Z.shape
(6, 6)  # 3 groups, 2 effects per group

Notes

Z is constructed as a block diagonal matrix where each block corresponds to one group in one random effect term:

For random intercept (1 | group):

Z_ij = 1 if observation i is in group j, else 0

For random intercept + slope (1 + x | group):

Z has two columns per group: [1, x_i] for each observation i in that group, zero elsewhere.

Multiple random effect terms are concatenated horizontally.

aurora.models.gamm.extract_random_effects(b, Z_info)

Extract random effect coefficients by group.

Parameters:

• b (ndarray, shape (q,)) – Random effect coefficients (BLUPs) from model fitting

• Z_info (list of dict) – Information about Z matrix structure (from construct_Z_matrix)

Returns:

random_effects – Dictionary mapping grouping_variable -> {group_id -> coefficients}. For each group, coefficients are [intercept, slope1, slope2, …] depending on the random effect specification.

Return type:

dict

Examples

>>> # For (1 | subject) with 3 groups
>>> b = np.array([0.5, -0.3, 0.2])
>>> Z_info = [{
...     'grouping': 'subject',
...     'n_effects': 1,
...     'n_groups': 3,
...     'groups': np.array([1, 2, 3]),
...     'start_col': 0,
...     'end_col': 3,
... }]
>>> random_effects = extract_random_effects(b, Z_info)
>>> random_effects['subject'][1]
array([0.5])
>>> random_effects['subject'][2]
array([-0.3])

Notes

This function is the inverse of construct_Z_matrix: it takes the vector of random effect coefficients and organizes them by group for interpretation and prediction.

class aurora.models.gamm.REMLResult(psi, sigma2, log_likelihood, theta, n_iterations, converged, V=None, P=None)

Bases: object

Result from REML variance component estimation.

Variables:

• psi (ndarray, shape (q, q)) – Estimated random effects covariance matrix Ψ.

• sigma2 (float) – Estimated residual variance σ².

• log_likelihood (float) – REML log-likelihood at optimum.

• theta (ndarray) – Optimized covariance parameters (internal parameterization).

• n_iterations (int) – Number of optimization iterations.

• converged (bool) – Whether optimization converged successfully.

• V (ndarray, shape (n, n)) – Marginal covariance matrix V = ZΨZ’ + σ²I.

• P (ndarray, shape (n, n)) – REML projection matrix P = V⁻¹ - V⁻¹X(X’V⁻¹X)⁻¹X’V⁻¹.

Parameters:

• psi (NDArray[np.floating])

• sigma2 (float)

• log_likelihood (float)

• theta (NDArray[np.floating])

• n_iterations (int)

• converged (bool)

• V (NDArray[np.floating] | None)

• P (NDArray[np.floating] | None)

psi: NDArray[np.floating]

sigma2: float

log_likelihood: float

theta: NDArray[np.floating]

n_iterations: int

converged: bool

V: NDArray[np.floating] | None = None

P: NDArray[np.floating] | None = None

aurora.models.gamm.estimate_variance_components(y, X, Z, Z_info, covariance='unstructured', method='L-BFGS-B', maxiter=1000, tol=1e-6, initial_psi=None, initial_sigma2=None, store_matrices=False)

Estimate variance components via REML.

Parameters:

• y (ndarray, shape (n,)) – Response vector.

• X (ndarray, shape (n, p)) – Fixed effects design matrix.

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• Z_info (list of dict) – Metadata about Z structure (from construct_Z_matrix). Each dict contains: ‘n_effects’, ‘n_groups’, ‘grouping’, ‘start_col’, ‘end_col’.

• covariance (str, default='unstructured') – Covariance structure: ‘unstructured’, ‘diagonal’, or ‘identity’. Applied to all random effect terms.

• method (str, default='L-BFGS-B') – Optimization method for scipy.optimize.minimize.

• maxiter (int, default=1000) – Maximum number of iterations.

• tol (float, default=1e-6) – Convergence tolerance.

• initial_psi (ndarray or list of ndarray or None, optional) – Initial value(s) for Ψ. If None, uses identity for each term. If list, one per random effect term. If single array, used for all terms.

• initial_sigma2 (float or None, optional) – Initial value for σ². If None, uses variance of residuals.

• store_matrices (bool, default=False) – Whether to store V and P matrices in result (memory intensive).

Returns:

result – REML estimation result. For multiple random effects, psi is block-diagonal containing all variance-covariance matrices.

Return type:

REMLResult

Examples

>>> # Single random intercept model
>>> n = 100
>>> y = np.random.randn(n)
>>> X = np.ones((n, 1))
>>> groups = np.repeat(np.arange(20), 5)
>>> Z = np.zeros((n, 20))
>>> Z[np.arange(n), groups] = 1
>>> Z_info = [{'n_effects': 1, 'n_groups': 20, 'grouping': 'subject',
...            'start_col': 0, 'end_col': 20}]
>>> result = estimate_variance_components(y, X, Z, Z_info)

>>> # Multiple random effects (crossed): (1 | school) + (1 | class)
>>> # Z_info = [
>>> #     {'n_effects': 1, 'n_groups': 10, 'grouping': 'school', 'start_col': 0, 'end_col': 10},
>>> #     {'n_effects': 1, 'n_groups': 30, 'grouping': 'class', 'start_col': 10, 'end_col': 40}
>>> # ]
>>> # result = estimate_variance_components(y, X, Z, Z_info)

aurora.models.gamm.estimate_fixed_effects(y, X, V)

Estimate fixed effects β via generalized least squares.

Parameters:

• y (ndarray, shape (n,)) – Response vector.

• X (ndarray, shape (n, p)) – Fixed effects design matrix.

• V (ndarray, shape (n, n)) – Marginal covariance matrix.

Returns:

• beta (ndarray, shape (p,)) – Fixed effects estimates.

• cov_beta (ndarray, shape (p, p)) – Covariance matrix of β: (X’V⁻¹X)⁻¹.

Return type:

tuple[ndarray, ndarray]

aurora.models.gamm.estimate_random_effects(y, X, Z, beta, psi, sigma2, Z_info=None)

Estimate random effects b (BLUPs) via conditional expectation.

Parameters:

• y (ndarray, shape (n,)) – Response vector.

• X (ndarray, shape (n, p)) – Fixed effects design matrix.

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• beta (ndarray, shape (p,)) – Fixed effects estimates.

• psi (ndarray) – Random effects covariance matrix. Can be: - shape (q_effects, q_effects): per-group covariance (single term) - shape (q, q): full block-diagonal for multiple terms

• sigma2 (float) – Residual variance.

• Z_info (list of dict or None, optional) – Information about Z structure. If provided and psi is per-group, will expand psi to block-diagonal.

Returns:

b – Random effects estimates (BLUPs).

Return type:

ndarray, shape (q,)

Notes

b = Ψ_full Z’V⁻¹(y - Xβ) where V = ZΨ_full Z’ + σ²I

For multiple random effects, psi should already be block-diagonal.

aurora.models.gamm.compute_V_matrix(Z, psi, sigma2, n_effects=None)

Compute marginal covariance matrix V = ZΨZ’ + σ²I.

Constructs the marginal covariance of the response vector under the linear mixed model y = Xβ + Zb + ε, where b ~ N(0, Ψ) and ε ~ N(0, σ²I).

Parameters:

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• psi (ndarray, shape (q_effects, q_effects)) – Random effects covariance matrix (per-group structure).

• sigma2 (float) – Residual variance σ².

• n_effects (int or None, optional) – Number of random effects per group. If provided and psi is smaller than the full dimension q, the per-group psi is expanded to a block-diagonal matrix with q // n_effects repeated blocks.

Returns:

V – Marginal covariance matrix V = Z Ψ_full Z’ + σ² I.

Return type:

ndarray, shape (n, n)

Notes

For a random intercept model with m groups, psi is (1, 1) but Z has m columns. This function automatically expands psi to diag(σ²_b, ..., σ²_b) when n_effects is provided.

See also

compute_P_matrix

REML projection matrix using V.

aurora.models.gamm.compute_P_matrix(V, X)

Compute REML projection matrix P = V⁻¹ - V⁻¹X(X’V⁻¹X)⁻¹X’V⁻¹.

The REML projection matrix removes the contribution of fixed effects from the inverse covariance, yielding the restricted (REML) likelihood that accounts for degrees of freedom lost in estimating β.

Parameters:

• V (ndarray, shape (n, n)) – Marginal covariance matrix V = ZΨZ’ + σ²I.

• X (ndarray, shape (n, p)) – Fixed effects design matrix.

Returns:

P – REML projection matrix. Satisfies PX = 0 and is symmetric positive semi-definite.

Return type:

ndarray, shape (n, n)

Notes

Computation uses Cholesky decomposition for numerical stability. Both V⁻¹ and (X’V⁻¹X)⁻¹ are computed via Cholesky solves rather than explicit inversion.

See also

compute_V_matrix

Constructs V from Z, Ψ, and σ².

reml_log_likelihood

Uses P to compute the REML log-likelihood.

aurora.models.gamm.reml_log_likelihood(y, X, V, P)

Compute REML log-likelihood.

Evaluates the Restricted Maximum Likelihood (REML) criterion, which accounts for the loss of degrees of freedom from estimating fixed effects.

Parameters:

• y (ndarray, shape (n,)) – Response vector.

• X (ndarray, shape (n, p)) – Fixed effects design matrix.

• V (ndarray, shape (n, n)) – Marginal covariance matrix V = ZΨZ’ + σ²I.

• P (ndarray, shape (n, n)) – REML projection matrix (from compute_P_matrix()).

Returns:

log_lik – REML log-likelihood value.

Return type:

float

Notes

The REML log-likelihood is:

l_REML = -0.5 [log|V| + log|X'V⁻¹X| + y'Py + (n-p)log(2π)]

Determinants are computed via Cholesky decomposition for numerical stability (log|A| = 2 Σ log diag(L) where A = LL').

References

Harville, D. A. (1974). “Bayesian inference for variance components using only error contrasts.” Biometrika, 61(2), 383-385.

See also

compute_P_matrix

Computes the REML projection matrix P.

reml_objective

Negative REML log-likelihood (for optimization).

class aurora.models.gamm.GAMMResult(coefficients, beta_parametric, beta_smooth, random_effects, variance_components, covariance_params, residual_variance, smoothing_parameters, edf_total, edf_parametric, edf_smooth, fitted_values, residuals, log_likelihood, aic, bic, converged, n_iterations, n_obs, n_groups, family, _X_parametric=None, _X_smooth=None, _Z=None, _Z_info=None, _y=None)

Bases: object

Result from GAMM fitting.

Variables:

• coefficients (ndarray) – Combined coefficients [β_parametric, β_smooth, b_random].

• beta_parametric (ndarray) – Parametric (fixed) effect coefficients.

• beta_smooth (dict[str, ndarray]) – Smooth term coefficients by term name.

• random_effects (dict[str, ndarray]) – Random effect coefficients by grouping variable.

• variance_components (list[ndarray]) – List of variance-covariance matrices Ψ, one per random effect term. For single term models, this is a list with one element.

• covariance_params (list[ndarray] | None) – List of raw covariance parameters for each random effect term. For structured covariance (AR1, CS, etc.), contains transformed parameters. For example, AR1 params are [log(σ²), arctanh(ρ)]. None for unstructured covariance.

• residual_variance (float) – Residual variance σ².

• smoothing_parameters (dict[str, float] | None) – Smoothing parameters λ by smooth term (if estimated).

• edf_total (float) – Total effective degrees of freedom.

• edf_parametric (float) – EDF for parametric terms.

• edf_smooth (dict[str, float]) – EDF by smooth term.

• fitted_values (ndarray) – Fitted values on the response scale (μ̂; for the Gaussian identity link this equals η̂ = Xβ + Zb).

• residuals (ndarray) – Response residuals y − μ̂.

• log_likelihood (float) – Log-likelihood (or REML log-likelihood). For non-Gaussian PQL fits this is the conditional log-likelihood at the BLUPs, not the marginal likelihood.

• aic (float) – Akaike Information Criterion.

• bic (float) – Bayesian Information Criterion.

• converged (bool) – Whether fitting converged.

• n_iterations (int) – Number of iterations.

• n_obs (int) – Number of observations.

• n_groups (int) – Number of groups.

• family (str) – Distribution family.

Parameters:

• coefficients (NDArray[np.floating])

• beta_parametric (NDArray[np.floating])

• beta_smooth (dict[str, NDArray[np.floating]])

• random_effects (dict[Any, dict[Any, NDArray[np.floating]]])

• variance_components (list[NDArray[np.floating]])

• covariance_params (list[NDArray[np.floating]] | None)

• residual_variance (float)

• smoothing_parameters (dict[str, float] | None)

• edf_total (float)

• edf_parametric (float)

• edf_smooth (dict[str, float])

• fitted_values (NDArray[np.floating])

• residuals (NDArray[np.floating])

• log_likelihood (float)

• aic (float)

• bic (float)

• converged (bool)

• n_iterations (int)

• n_obs (int)

• n_groups (int)

• family (str)

• _X_parametric (NDArray[np.floating] | None)

• _X_smooth (dict[str, NDArray[np.floating] | spmatrix] | None)

• _Z (NDArray[np.floating] | None)

• _Z_info (list[dict] | None)

• _y (NDArray[np.floating] | None)

coefficients: NDArray[np.floating]

beta_parametric: NDArray[np.floating]

beta_smooth: dict[str, NDArray[np.floating]]

random_effects: dict[Any, dict[Any, NDArray[np.floating]]]

variance_components: list[NDArray[np.floating]]

covariance_params: list[NDArray[np.floating]] | None

residual_variance: float

smoothing_parameters: dict[str, float] | None

edf_total: float

edf_parametric: float

edf_smooth: dict[str, float]

fitted_values: NDArray[np.floating]

residuals: NDArray[np.floating]

log_likelihood: float

aic: float

bic: float

converged: bool

n_iterations: int

n_obs: int

n_groups: int

family: str

predict(include_random=True)

Make predictions using the fitted model.

Parameters:

include_random (bool, default=True) – Whether to include random effects in predictions. If True, returns fitted values (fixed + random effects). If False, returns only fixed effects (population-level predictions).

Returns:

predictions – Model predictions.

Return type:

ndarray

Examples

>>> # Population-level predictions (fixed effects only)
>>> pred_fixed = result.predict(include_random=False)
>>>
>>> # Individual predictions (fixed + random effects)
>>> pred_full = result.predict(include_random=True)

summary()

Generate a formatted summary of the model fit.

Returns:

summary_str – Formatted summary string with fixed effects, random effects, and model fit statistics.

Return type:

str

Examples

>>> result = fit_gamm(formula='y ~ x + (1|subject)', data=df)
>>> print(result.summary())

aurora.models.gamm.fit_gamm_gaussian(X_parametric, X_smooth, Z, Z_info, y, S_smooth=None, lambda_smooth=None, covariance='unstructured', maxiter=100, tol=1e-6, backend='numpy', device=None)

Fit Gaussian GAMM (Linear Mixed Model with smooth terms).

Parameters:

• X_parametric (ndarray, shape (n, p_para)) – Parametric fixed effects design matrix.

• X_smooth (dict[str, ndarray], optional) – Smooth term basis matrices by term name.

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• Z_info (list of dict) – Metadata about Z structure from construct_Z_matrix.

• y (ndarray, shape (n,)) – Response vector.

• S_smooth (dict[str, ndarray], optional) – Penalty matrices by smooth term name.

• lambda_smooth (dict[str, float], optional) – Smoothing parameters by term name. If None, uses default λ=1.0.

• covariance (str, default='unstructured') – Covariance structure for random effects.

• maxiter (int, default=100) – Maximum iterations (for future iterative methods).

• tol (float, default=1e-6) – Convergence tolerance.

• backend (str, default='numpy') – Computational backend: ‘numpy’, ‘torch’, or ‘jax’.

• device (str, optional) – Device for computation (for torch backend): ‘cpu’, ‘cuda’, etc.

Returns:

result – Fitted GAMM result.

Return type:

GAMMResult

Notes

For Gaussian family, the model reduces to:

y = X_para β_para + Σ_k X_k β_k + Zb + ε b ~ N(0, Ψ), ε ~ N(0, σ²I)

This is solved in closed form via penalized least squares with REML estimation of variance components.

Algorithm: 1. Construct combined fixed effects matrix X = [X_para | X_smooth_1 | …] 2. Build combined penalty matrix S (block diagonal for smooth terms) 3. Estimate variance components Ψ, σ² via REML 4. Solve mixed model equations with estimated variances 5. Compute EDF, AIC, BIC

Examples

>>> # Random intercept with linear predictor
>>> n = 100
>>> X_para = np.ones((n, 1))
>>> Z = np.eye(n)[:, :10]  # 10 groups
>>> y = X_para @ [2.0] + Z @ np.random.randn(10) + np.random.randn(n)
>>> Z_info = [{'n_effects': 1, 'n_groups': 10}]
>>> result = fit_gamm_gaussian(X_para, None, Z, Z_info, y)

aurora.models.gamm.predict_gamm(result, X_parametric_new, X_smooth_new=None, Z_new=None, include_random=False)

Make predictions from fitted GAMM.

Parameters:

• result (GAMMResult) – Fitted GAMM result.

• X_parametric_new (ndarray, shape (n_new, p_para)) – Parametric design matrix for new data.

• X_smooth_new (dict[str, ndarray], optional) – Smooth term basis matrices for new data.

• Z_new (ndarray, shape (n_new, q), optional) – Random effects design matrix for new data.

• include_random (bool, default=False) – Whether to include random effects in predictions.

Returns:

predictions – Predicted values.

Return type:

ndarray, shape (n_new,)

Notes

• If include_random=False: η̂ = X_para β̂_para + Σ_k X_k β̂_k (population-level)

• If include_random=True: η̂ = X_para β̂_para + Σ_k X_k β̂_k + Z b̂ (conditional)

For new groups not in training data, random effects default to 0.

aurora.models.gamm.solve_mixed_model_equations(X, Z, y, S_smooth=None, psi_inv=None, lambda_smooth=0.0)

Solve penalized mixed model equations for Gaussian GAMM.

Solves the augmented system:

[X’X + λS X’Z ] [β] [X’y] [Z’X Z’Z + Ψ⁻¹] [b] = [Z’y]

Parameters:

• X (ndarray, shape (n, p)) – Fixed effects design matrix (parametric + smooth basis).

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• y (ndarray, shape (n,)) – Response vector.

• S_smooth (ndarray, shape (p, p), optional) – Smoothing penalty matrix for smooth terms.

• psi_inv (ndarray, shape (q, q), optional) – Inverse of random effects covariance Ψ⁻¹.

• lambda_smooth (float, default=0.0) – Smoothing parameter.

Returns:

• beta (ndarray, shape (p,)) – Fixed effect coefficients.

• b (ndarray, shape (q,)) – Random effect coefficients (BLUPs).

Return type:

tuple[ndarray, ndarray]

Notes

Uses Cholesky decomposition for numerical stability. For large systems, could exploit block structure or sparsity.

aurora.models.gamm.fit_gamm(formula=None, data=None, y=None, X=None, random_effects=None, groups_data=None, family='gaussian', covariance='unstructured', maxiter=100, tol=1e-6, use_sparse=False, backend='numpy', device=None)

Fit a Generalized Additive Mixed Model.

Can be called in two modes: 1. Formula mode: Pass formula and data 2. Matrix mode: Pass y, X, random_effects, and groups_data

Parameters:

• formula (str, optional) – R-style formula string (formula mode). Example: “y ~ x1 + (1 | subject)” or “y ~ x1 + (1 + time | subject)”

• data (DataFrame or dict, optional) – Data containing all variables referenced in formula (formula mode).

• y (array-like, shape (n,), optional) – Response variable (matrix mode).

• X (array-like, shape (n, p), optional) – Design matrix for parametric fixed effects (matrix mode). If None, uses intercept-only model.

• random_effects (list of RandomEffect, optional) – Random effect specifications (matrix mode).

• groups_data (dict or DataFrame, optional) – Grouping variables for random effects (matrix mode). Keys should match RandomEffect.grouping names.

• family (str, default='gaussian') – Distribution family. Gaussian uses exact REML estimation; non-Gaussian families (e.g. ‘poisson’, ‘binomial’) are fitted with Penalized Quasi-Likelihood (PQL).

• covariance (str, default='unstructured') –

  Covariance structure for random effects. Options: - ‘identity’: Independent random effects (default) - ‘diagonal’: Heterogeneous variances, no correlation - ‘unstructured’: Full covariance matrix - ‘ar1’: Autoregressive AR(1) for temporal correlation - ‘compound_symmetry’ or ‘cs’: Exchangeable correlation - ‘exponential’: Spatial decay (requires coordinates) - ‘matern’: Matérn spatial covariance (requires coordinates)

  For temporal/longitudinal data, use ‘ar1’. For clustered data with equal correlations, use ‘compound_symmetry’.

• maxiter (int, default=100) – Maximum iterations for optimization.

• tol (float, default=1e-6) – Convergence tolerance.

• use_sparse (bool, default=False) – Whether to use sparse matrix operations for smooth terms. When True, uses sparse CSR format for basis matrices and sparse linear solvers. Provides significant speedup (10-100×) and memory reduction (6-8×) for large problems with B-spline basis functions. Currently only supported when smooth terms use B-spline basis.

• backend (str, default='numpy') – Computational backend: ‘numpy’, ‘torch’, or ‘jax’.

• device (str, optional) – Device for computation (for torch backend): ‘cpu’, ‘cuda’, ‘cuda:0’, etc. If None, uses CUDA if available, else CPU.

Returns:

result – Fitted GAMM result with coefficients, variance components, and diagnostics.

Return type:

GAMMResult

Raises:

ValueError – If random_effects provided but groups_data is None. If group variables not found in groups_data.

Examples

Formula mode (recommended):

>>> # Random intercept model
>>> import pandas as pd
>>> import numpy as np
>>> from aurora.models.gamm import fit_gamm
>>>
>>> # Generate data
>>> n_groups, n_per_group = 10, 20
>>> n = n_groups * n_per_group
>>> data = pd.DataFrame({
...     'y': np.random.randn(n),
...     'x1': np.random.randn(n),
...     'subject': np.repeat(np.arange(n_groups), n_per_group)
... })
>>>
>>> # Fit model with formula
>>> result = fit_gamm(
...     formula="y ~ x1 + (1 | subject)",
...     data=data,
...     covariance='identity'
... )
>>>
>>> # Access results
>>> print(result.beta_parametric)  # Fixed effects
>>> print(result.variance_components)  # Random effect variance

>>> # Random intercept + slope
>>> result = fit_gamm(
...     formula="y ~ x1 + (1 + x1 | subject)",
...     data=data,
...     covariance='unstructured'
... )
>>> print(result.variance_components)  # 2x2 covariance matrix

Matrix mode (advanced):

>>> from aurora.models.gamm import RandomEffect
>>>
>>> # Manually construct matrices
>>> groups = data['subject'].values
>>> x = data['x1'].values
>>> y = data['y'].values
>>> X = np.column_stack([np.ones(n), x])
>>>
>>> # Specify random effects manually
>>> re = RandomEffect(grouping='subject')
>>> result = fit_gamm(
...     y=y,
...     X=X,
...     random_effects=[re],
...     groups_data={'subject': groups},
...     covariance='identity'
... )
>>> print(result.beta_parametric)

Notes

This is a high-level interface that automatically constructs design matrices for random effects and calls the appropriate fitting function based on the family.

For Gaussian family, uses exact REML estimation. For non-Gaussian families, uses Penalized Quasi-Likelihood (PQL; Breslow & Clayton 1993) with uncorrected estimates.

Non-Gaussian likelihood quantities. For non-Gaussian families, log_likelihood is the conditional log-likelihood Σᵢ ℓ(yᵢ; μ̂ᵢ) evaluated at the BLUPs, not the marginal likelihood; aic/bic are conditional information criteria derived from it. Likelihood-ratio tests on fixed effects based on these quantities are not valid (the anova module issues a warning).

Limitations. offset and prior weights are not currently supported by fit_gamm; for models requiring them, use fit_glm (no random effects) instead.

Formula mode supports: - R-style formula syntax with lme4-style random effects - Automatic design matrix construction from DataFrames - Random intercepts: (1 | group) - Random slopes: (1 + x | group) - Nested effects: (1 | a/b) - Crossed effects: (1 | a) + (1 | b)

aurora.models.gamm.fit_gamm_with_smooth(y, X_parametric, X_smooth, S_smooth, random_effects, groups_data, lambda_smooth=None, family='gaussian', covariance='unstructured', maxiter=100, tol=1e-6)

Fit GAMM with smooth terms (advanced interface).

Parameters:

• y (array-like, shape (n,)) – Response variable.

• X_parametric (array-like, shape (n, p)) – Design matrix for parametric fixed effects.

• X_smooth (dict of str -> ndarray) – Smooth term basis matrices by term name.

• S_smooth (dict of str -> ndarray) – Penalty matrices by term name.

• random_effects (list of RandomEffect) – Random effect specifications.

• groups_data (dict or DataFrame) – Grouping variables for random effects.

• lambda_smooth (dict of str -> float, optional) – Smoothing parameters by term name. If None, uses λ=1.0 for all.

• family (str, default='gaussian') – Distribution family.

• covariance (str, default='unstructured') – Covariance structure for random effects.

• maxiter (int, default=100) – Maximum iterations.

• tol (float, default=1e-6) – Convergence tolerance.

Returns:

result – Fitted GAMM result.

Return type:

GAMMResult

Examples

>>> # GAMM with smooth term + random intercept
>>> from aurora.smoothing import fit_cubic_spline
>>> from aurora.models.gamm import fit_gamm_with_smooth, RandomEffect
>>>
>>> # Data
>>> n = 100
>>> x_smooth = np.linspace(0, 2*np.pi, n)
>>> groups = np.repeat(np.arange(10), 10)
>>>
>>> # Build smooth basis
>>> import scipy.interpolate as interp
>>> k = 10
>>> knots = np.linspace(0, 2*np.pi, k-2)
>>> tck = interp.splrep(x_smooth, np.zeros(n), t=knots[1:-1], k=3)
>>> X_smooth_basis = interp.BSpline.design_matrix(x_smooth, tck[0], 3)
>>>
>>> # Parametric design
>>> X_para = np.ones((n, 1))
>>>
>>> # Random effects
>>> re = RandomEffect(grouping='subject')
>>>
>>> # Generate response
>>> y = 2.0 + np.sin(x_smooth) + np.random.randn(10)[groups] + np.random.randn(n)*0.3
>>>
>>> # Penalty matrix (second derivative)
>>> S = np.diag([0]*2 + [1]*(k-2))  # Penalize non-linear components
>>>
>>> result = fit_gamm_with_smooth(
...     y=y,
...     X_parametric=X_para,
...     X_smooth={'s(x)': X_smooth_basis},
...     S_smooth={'s(x)': S},
...     random_effects=[re],
...     groups_data={'subject': groups},
...     lambda_smooth={'s(x)': 0.1},
...     covariance='identity'
... )

Notes

This is the advanced interface for GAMMs with smooth terms. For basic usage without smooth terms, use fit_gamm() instead.

Smooth term construction (basis matrices and penalties) should be done using the smoothing module from Phase 3.

aurora.models.gamm.predict_from_gamm(result, X_new, groups_new=None, X_smooth_new=None, include_random=False)

Make predictions from fitted GAMM.

Parameters:

• result (GAMMResult) – Fitted GAMM result.

• X_new (array-like, shape (n_new, p)) – New parametric design matrix.

• groups_new (array-like, shape (n_new,), optional) – Group indicators for new observations. Required if include_random=True.

• X_smooth_new (dict of str -> ndarray, optional) – New smooth term basis matrices.

• include_random (bool, default=False) – Whether to include random effects in predictions.

Returns:

predictions – Predicted values.

Return type:

ndarray, shape (n_new,)

Examples

>>> # Population-level predictions (no random effects)
>>> X_new = np.column_stack([np.ones(20), np.random.randn(20)])
>>> pred_pop = predict_from_gamm(result, X_new)
>>>
>>> # Conditional predictions (with random effects)
>>> groups_new = np.array([0, 0, 1, 1, 2, 2, ...])  # Existing groups
>>> pred_cond = predict_from_gamm(
...     result, X_new,
...     groups_new=groups_new,
...     include_random=True
... )

Notes

• Population-level predictions are appropriate for new, unobserved groups

• Conditional predictions are appropriate for predicting within observed groups

• For new groups, random effects default to 0 (population mean)

class aurora.models.gamm.PQLResult(beta, b, psi, sigma2, fitted_values, linear_predictor, converged, n_iter_outer, n_iter_inner, deviance, log_likelihood)

Bases: object

Result of PQL estimation.

Variables:

• beta (ndarray, shape (p,)) – Estimated fixed effect coefficients

• b (ndarray, shape (q,)) – Estimated random effect coefficients (BLUPs)

• psi (ndarray, shape (n_effects, n_effects)) – Estimated variance-covariance matrix for random effects

• sigma2 (float) – Estimated residual variance (on working response scale)

• fitted_values (ndarray, shape (n,)) – Fitted values on response scale (μ)

• linear_predictor (ndarray, shape (n,)) – Linear predictor (η = Xβ + Zb)

• converged (bool) – Whether the algorithm converged

• n_iter_outer (int) – Number of outer iterations (variance component updates)

• n_iter_inner (int) – Total number of inner iterations (coefficient updates)

• deviance (float) – Final deviance

• log_likelihood (float) – Approximate log-likelihood

Parameters:

• beta (NDArray[np.floating])

• b (NDArray[np.floating])

• psi (NDArray[np.floating])

• sigma2 (float)

• fitted_values (NDArray[np.floating])

• linear_predictor (NDArray[np.floating])

• converged (bool)

• n_iter_outer (int)

• n_iter_inner (int)

• deviance (float)

• log_likelihood (float)

beta: NDArray[np.floating]

b: NDArray[np.floating]

psi: NDArray[np.floating]

sigma2: float

fitted_values: NDArray[np.floating]

linear_predictor: NDArray[np.floating]

converged: bool

n_iter_outer: int

n_iter_inner: int

deviance: float

log_likelihood: float

aurora.models.gamm.fit_pql(X, Z, y, family, S=None, lambda_=0.0, psi_init=None, maxiter_outer=20, maxiter_inner=10, tol_outer=1e-4, tol_inner=1e-6, update_psi=True)

Fit a GLMM using Penalized Quasi-Likelihood.

Parameters:

• X (ndarray, shape (n, p)) – Fixed effects design matrix

• Z (ndarray, shape (n, q)) – Random effects design matrix

• y (ndarray, shape (n,)) – Response variable

• family (Family or str) – Distribution family (‘gaussian’, ‘poisson’, ‘binomial’)

• S (ndarray, shape (p, p), optional) – Penalty matrix for smooth terms (default: zeros)

• lambda (float, default=0.0) – Smoothing parameter

• psi_init (ndarray, shape (n_effects, n_effects), optional) – Initial variance-covariance matrix (default: identity)

• maxiter_outer (int, default=20) – Maximum outer iterations (variance updates)

• maxiter_inner (int, default=10) – Maximum inner iterations per outer (coefficient updates)

• tol_outer (float, default=1e-4) – Convergence tolerance for variance components

• tol_inner (float, default=1e-6) – Convergence tolerance for coefficients

• update_psi (bool, default=True) – Whether to update variance components (False for fixed Ψ)

• lambda_ (float)

Returns:

result – Fitted PQL model

Return type:

PQLResult

Notes

For Gaussian family, this reduces to standard LMM which could be solved more efficiently with REML. Use fit_gamm_gaussian() instead for Gaussian data.

The algorithm iterates between: 1. Updating (β, b) given Ψ via weighted least squares 2. Updating Ψ given (β, b) via empirical covariance

Convergence is determined by relative change in Ψ falling below tol_outer. With update_psi=False (fixed Ψ), convergence is determined by the inner (β, b) fixed point instead.

For the Gamma family, the working weights assume dispersion φ = 1; the estimated dispersion is not propagated into the weighting (documented limitation).

For the Binomial family, fitted probabilities numerically 0 or 1 indicate complete or quasi-complete separation; a warning is issued and the result is reported as not converged.

Examples

>>> import numpy as np
>>> from aurora.models.gamm.pql import fit_pql
>>> # Generate Poisson data with random intercepts
>>> n_groups, n_per_group = 5, 20
>>> n = n_groups * n_per_group
>>> groups = np.repeat(np.arange(n_groups), n_per_group)
>>> X = np.column_stack([np.ones(n), np.random.randn(n)])
>>> Z = np.eye(n_groups)[groups]  # Random intercept design
>>> b_true = np.random.randn(n_groups) * 0.5
>>> eta = X @ [1.0, 0.3] + Z @ b_true
>>> y = np.random.poisson(np.exp(eta))
>>> result = fit_pql(X, Z, y, family='poisson')
>>> print(result.converged)
True

class aurora.models.gamm.LaplaceResult(beta, b, psi, sigma2, fitted_values, linear_predictor, hessian, converged, n_iter, log_likelihood)

Bases: object

Result of Laplace approximation.

Variables:

• beta (ndarray, shape (p,)) – Estimated fixed effect coefficients

• b (ndarray, shape (q,)) – Estimated random effect coefficients at conditional mode

• psi (ndarray, shape (n_effects, n_effects)) – Estimated variance-covariance matrix for random effects

• sigma2 (float) – Estimated residual variance (for Gaussian; 1.0 for others)

• fitted_values (ndarray, shape (n,)) – Fitted values on response scale (μ)

• linear_predictor (ndarray, shape (n,)) – Linear predictor (η = Xβ + Zb)

• hessian (ndarray, shape (q, q)) – Hessian of negative log joint density at b̂

• converged (bool) – Whether the algorithm converged

• n_iter (int) – Number of outer iterations

• log_likelihood (float) – Approximate marginal log-likelihood

Parameters:

• beta (NDArray[np.floating])

• b (NDArray[np.floating])

• psi (NDArray[np.floating])

• sigma2 (float)

• fitted_values (NDArray[np.floating])

• linear_predictor (NDArray[np.floating])

• hessian (NDArray[np.floating])

• converged (bool)

• n_iter (int)

• log_likelihood (float)

beta: NDArray[np.floating]

b: NDArray[np.floating]

psi: NDArray[np.floating]

sigma2: float

fitted_values: NDArray[np.floating]

linear_predictor: NDArray[np.floating]

hessian: NDArray[np.floating]

converged: bool

n_iter: int

log_likelihood: float

aurora.models.gamm.fit_laplace(X, Z, y, family, S=None, lambda_=0.0, psi_init=None, maxiter=50, tol=1e-6)

Fit a GLMM using Laplace approximation.

Parameters:

• X (ndarray, shape (n, p)) – Fixed effects design matrix

• Z (ndarray, shape (n, q)) – Random effects design matrix

• y (ndarray, shape (n,)) – Response variable

• family (Family or str) – Distribution family (‘gaussian’, ‘poisson’, ‘binomial’)

• S (ndarray, shape (p, p), optional) – Penalty matrix for smooth terms (default: zeros)

• lambda (float, default=0.0) – Smoothing parameter

• psi_init (ndarray, shape (n_effects, n_effects), optional) – Initial variance-covariance matrix (default: identity). Required for models with more than one random effect per group (e.g. random intercept + slope): its dimension determines how the columns of Z are grouped. Without it, random-intercept structure (one effect per group) is assumed.

• maxiter (int, default=50) – Maximum iterations

• tol (float, default=1e-6) – Convergence tolerance

• lambda_ (float)

Returns:

Fitted model results

Return type:

LaplaceResult

Notes

Laplace approximation is more accurate than PQL, especially for: - Binary data with small group sizes - Count data with small means - Models with large random effects

However, it is computationally more expensive due to: - Inner optimization to find conditional mode b̂ - Hessian computation at each iteration

Examples

>>> # Poisson GLMM with Laplace approximation
>>> import numpy as np
>>> from aurora.models.gamm import fit_laplace
>>>
>>> n_groups, n_per_group = 10, 20
>>> n = n_groups * n_per_group
>>> groups = np.repeat(np.arange(n_groups), n_per_group)
>>> x = np.random.randn(n)
>>> X = np.column_stack([np.ones(n), x])
>>>
>>> # Random effects design
>>> Z = np.zeros((n, n_groups))
>>> for i in range(n):
...     Z[i, groups[i]] = 1
>>>
>>> # Generate data
>>> b_true = np.random.randn(n_groups) * 0.5
>>> eta = 1.0 + 0.5 * x + b_true[groups]
>>> y = np.random.poisson(np.exp(eta))
>>>
>>> # Fit model
>>> result = fit_laplace(X, Z, y, family='poisson')
>>> print(f"Log-likelihood: {result.log_likelihood:.2f}")

aurora.models.gamm.plot_caterpillar(result, grouping=None, effect_index=0, confidence=0.95, sort=True, figsize=None, ax=None)

Create caterpillar plot for random effects.

Caterpillar plots show the estimated random effects (BLUPs) for each group with confidence intervals, sorted by magnitude. They are useful for identifying groups with unusually large or small effects.

Parameters:

• result (GAMMResult) – Fitted GAMM model result.

• grouping (str or int, optional) – Which grouping variable to plot. If None, uses the first random effect term. Required if model has multiple grouping variables.

• effect_index (int, default=0) – Which random effect to plot (0=intercept, 1=first slope, etc.).

• confidence (float, default=0.95) – Confidence level for intervals (e.g., 0.95 for 95% CI).

• sort (bool, default=True) – Whether to sort groups by effect magnitude.

• figsize (tuple, optional) – Figure size (width, height) in inches.

• ax (Axes, optional) – Matplotlib axes to plot on. If None, creates new figure.

Returns:

• fig (Figure) – Matplotlib figure.

• ax (Axes) – Matplotlib axes with plot.

Return type:

tuple[Figure, Axes]

Examples

>>> import numpy as np
>>> from aurora.models.gamm import fit_gamm
>>> from aurora.models.gamm.plotting import plot_caterpillar
>>>
>>> # Fit model
>>> data = pd.DataFrame({
...     'y': np.random.randn(100),
...     'x': np.random.randn(100),
...     'subject': np.repeat(np.arange(10), 10)
... })
>>> result = fit_gamm(formula="y ~ x + (1 | subject)", data=data)
>>>
>>> # Plot random intercepts
>>> fig, ax = plot_caterpillar(result)
>>> plt.show()

Notes

The confidence intervals are approximate and computed using the posterior variance of the random effects:

Var(b | y) ≈ (Z’WZ + Ψ⁻¹)⁻¹

Groups whose confidence intervals do not include zero have effects that are significantly different from the population mean.

aurora.models.gamm.plot_random_effects_qq(result, grouping=None, effect_index=0, figsize=(6, 6), ax=None)

Create Q-Q plot for random effects normality check.

Q-Q (quantile-quantile) plots compare the distribution of random effects to a theoretical normal distribution. Points falling along the diagonal line indicate normality.

Parameters:

• result (GAMMResult) – Fitted GAMM model result.

• grouping (str or int, optional) – Which grouping variable to plot. If None, uses the first random effect term.

• effect_index (int, default=0) – Which random effect to plot (0=intercept, 1=first slope, etc.).

• figsize (tuple, default=(6, 6)) – Figure size (width, height) in inches.

• ax (Axes, optional) – Matplotlib axes to plot on. If None, creates new figure.

Returns:

• fig (Figure) – Matplotlib figure.

• ax (Axes) – Matplotlib axes with plot.

Return type:

tuple[Figure, Axes]

Examples

>>> fig, ax = plot_random_effects_qq(result)
>>> plt.show()

Notes

The random effects are assumed to follow a normal distribution:

b ~ N(0, Ψ)

Departures from normality (heavy tails, skewness) can be identified visually from Q-Q plots. Consider: - S-shaped pattern: Heavy tails (outliers) - Systematic curve: Skewness - Points far from line: Individual outliers

aurora.models.gamm.plot_random_effects_density(result, grouping=None, effect_index=0, show_normal=True, figsize=(8, 5), ax=None)

Plot density of random effects with optional normal overlay.

Parameters:

• result (GAMMResult) – Fitted GAMM model result.

• grouping (str or int, optional) – Which grouping variable to plot.

• effect_index (int, default=0) – Which random effect to plot.

• show_normal (bool, default=True) – Whether to overlay theoretical normal distribution.

• figsize (tuple, default=(8, 5)) – Figure size.

• ax (Axes, optional) – Matplotlib axes to plot on.

Returns:

• fig (Figure) – Matplotlib figure.

• ax (Axes) – Matplotlib axes with plot.

Return type:

tuple[Figure, Axes]

Examples

>>> fig, ax = plot_random_effects_density(result)
>>> plt.show()

Notes

The histogram shows the empirical distribution of random effects, while the overlaid curve shows the theoretical normal distribution N(0, Ψ) implied by the model.

aurora.models.gamm.plot_diagnostics(result, plot_type='residuals', figsize=(8, 6), ax=None)

Create diagnostic plots for GAMM models.

Parameters:

• result (GAMMResult) – Fitted GAMM model result.

• plot_type ({'residuals', 'fitted', 'qq', 'scale-location'}) – Type of diagnostic plot: - ‘residuals’: Residuals vs fitted values - ‘fitted’: Fitted vs observed values - ‘qq’: Q-Q plot of residuals - ‘scale-location’: Scale-location plot (sqrt(|residuals|) vs fitted)

• figsize (tuple, default=(8, 6)) – Figure size.

• ax (Axes, optional) – Matplotlib axes to plot on.

Returns:

• fig (Figure) – Matplotlib figure.

• ax (Axes) – Matplotlib axes with plot.

Return type:

tuple[Figure, Axes]

Examples

>>> # Residuals vs fitted
>>> fig, ax = plot_diagnostics(result, plot_type='residuals')
>>> plt.show()
>>>
>>> # Q-Q plot of residuals
>>> fig, ax = plot_diagnostics(result, plot_type='qq')
>>> plt.show()

Notes

Residuals vs Fitted: - Should show no pattern (randomness around zero) - Fan shape indicates heteroscedasticity - Curvature indicates nonlinearity

Q-Q Plot: - Points should fall on diagonal line - Departures indicate non-normality of residuals

Scale-Location: - Check homoscedasticity (constant variance) - Should show horizontal line with random scatter

Fitted vs Observed: - Points should fall near y=x line - Shows overall model fit quality

aurora.models.gamm.plot_random_effects_summary(result, grouping=None, effect_index=0, figsize=(14, 10))

Create comprehensive summary of random effects diagnostics.

Produces a 2x2 grid with: 1. Caterpillar plot 2. Q-Q plot 3. Density plot 4. Residuals vs fitted

Parameters:

• result (GAMMResult) – Fitted GAMM model result.

• grouping (str or int, optional) – Which grouping variable to plot.

• effect_index (int, default=0) – Which random effect to plot.

• figsize (tuple, default=(14, 10)) – Figure size.

Returns:

fig – Matplotlib figure with 2x2 subplot grid.

Return type:

Figure

Examples

>>> fig = plot_random_effects_summary(result)
>>> plt.show()

Notes

This is a convenience function that creates a comprehensive diagnostic summary in a single figure.

aurora.models.gamm.plot_gamm_diagnostics(result, figsize=(12, 10))

Create a 2x2 panel of diagnostic plots for a fitted GAMM model.

The four panels are:

1. Residuals vs Fitted – checks for non-linearity and heteroscedasticity. A blue spline smoother is overlaid to reveal systematic patterns.

2. Normal Q-Q Plot – assesses normality of residuals.

3. Scale-Location – plots sqrt(|residuals|) against fitted values to detect heteroscedasticity. A red smoother is overlaid.

4. Histogram of Residuals – visualises the residual distribution with a theoretical normal density curve overlaid.

Parameters:

• result (GAMMResult) – Fitted GAMM result. Must expose fitted_values and residuals attributes.

• figsize (tuple of float, default=(12, 10)) – Figure size as (width, height) in inches.

Returns:

• fig (matplotlib.figure.Figure) – The created figure object.

• axes (numpy.ndarray of matplotlib.axes.Axes) – Array of shape (2, 2) containing the individual Axes.

Raises:

ImportError – If matplotlib or scipy is not installed.

See also

plot_random_effects

Caterpillar and correlation plots for random effects.

Examples

>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm(formula='y ~ x + (1|subject)', data=df)
>>> fig, axes = plot_diagnostics(result)
>>> plt.show()

aurora.models.gamm.plot_gamm_random_effects(result, group_name=None, figsize=(15, 5))

Visualize random effects distribution with caterpillar and correlation plots.

For random-intercept-only models, produces a single sorted caterpillar plot. For random intercept + slope models, produces a 1x3 panel:

1. Caterpillar plot (intercepts) – random intercepts sorted by value.

2. Caterpillar plot (slopes) – random slopes sorted by value.

3. Intercept-slope scatter – scatter plot of intercepts vs slopes with the estimated correlation annotated.

Parameters:

• result (GAMMResult) – Fitted GAMM result. Must expose random_effects and variance_components attributes.

• group_name (str, optional) – Name of the grouping variable to plot. If None, the first grouping variable in result.random_effects is used.

• figsize (tuple of float, default=(15, 5)) – Figure size as (width, height) in inches (only used when there are two random effect components).

Returns:

• fig (matplotlib.figure.Figure) – The created figure object.

• axes (matplotlib.axes.Axes or numpy.ndarray of matplotlib.axes.Axes) – A single Axes (1 component) or an array of 3 Axes (2 components).

Raises:

• NotImplementedError – If the grouping variable has more than 2 random effect components.

• ImportError – If matplotlib is not installed.

See also

interpret_variance_components

Textual summary of variance components.

plot_diagnostics

Residual diagnostic plots.

Examples

>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm(formula='y ~ x + (1 + x | subject)', data=df)
>>> fig, axes = plot_random_effects(result, 'subject')
>>> plt.show()

aurora.models.gamm.interpret_variance_components(variance_components, group_names=None)

Interpret variance-covariance components of random effects.

Produces a human-readable summary of the estimated variance-covariance matrices for each grouping variable, including standard deviations, correlations, and qualitative interpretation of effect sizes.

Parameters:

• variance_components (list of ndarray) – List of variance-covariance matrices, one per grouping variable. Each matrix has shape (q, q) where q is the number of random effects for that grouping factor.

• group_names (list of str, optional) – Names of grouping variables corresponding to the entries in variance_components. If None, generic names (Group_1, Group_2, …) are used.

Returns:

interpretation – Formatted multi-line string with variance, standard deviation, covariance, and correlation estimates together with qualitative interpretation (e.g., weak/strong correlation, compensatory effects).

Return type:

str

See also

plot_random_effects

Visualize random effects distribution.

Examples

>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm(formula='y ~ x + (1 + x | subject)', data=df)
>>> print(interpret_variance_components(
...     result.variance_components,
...     list(result.random_effects.keys())
... ))

aurora.models.gamm.compute_r2_conditional_marginal(result)

Compute marginal and conditional R-squared for GAMM.

Implements the method of Nakagawa & Schielzeth (2013) for partitioning variance explained in mixed models into fixed-effects-only (marginal) and fixed-plus-random (conditional) components.

Parameters:

result (GAMMResult) – Fitted GAMM result. Must expose attributes _X_parametric, beta_parametric, fitted_values, residuals, residual_variance, and variance_components.

Returns:

• r2_marginal (float) – Proportion of total variance explained by fixed effects alone, computed as var_fixed / var_total.

• r2_conditional (float) – Proportion of total variance explained by both fixed and random effects, computed as (var_fixed + var_random) / var_total.

Return type:

tuple[float, float]

Notes

Total variance is decomposed as:

var_total = var_fixed + var_random + var_residual

where var_fixed is the variance of the fixed-effects predictions, var_random is the mean trace of the variance-covariance components, and var_residual is the residual variance from the fitted model.

For non-Gaussian families, these R-squared values are approximate.

References

[1]

Nakagawa, S., & Schielzeth, H. (2013). A general and simple method for obtaining R2 from generalized linear mixed-effects models. Methods in Ecology and Evolution, 4(2), 133-142.

See also

interpret_variance_components

Textual summary of variance components.

Examples

>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm(formula='y ~ x + (1|subject)', data=df)
>>> r2_m, r2_c = compute_r2_conditional_marginal(result)
>>> print(f"R² marginal: {r2_m:.3f}, R² conditional: {r2_c:.3f}")

Fitting

aurora.models.gamm.interface.fit_gamm(formula=None, data=None, y=None, X=None, random_effects=None, groups_data=None, family='gaussian', covariance='unstructured', maxiter=100, tol=1e-6, use_sparse=False, backend='numpy', device=None)

Fit a Generalized Additive Mixed Model.

Can be called in two modes: 1. Formula mode: Pass formula and data 2. Matrix mode: Pass y, X, random_effects, and groups_data

Parameters:

• formula (str, optional) – R-style formula string (formula mode). Example: “y ~ x1 + (1 | subject)” or “y ~ x1 + (1 + time | subject)”

• data (DataFrame or dict, optional) – Data containing all variables referenced in formula (formula mode).

• y (array-like, shape (n,), optional) – Response variable (matrix mode).

• X (array-like, shape (n, p), optional) – Design matrix for parametric fixed effects (matrix mode). If None, uses intercept-only model.

• random_effects (list of RandomEffect, optional) – Random effect specifications (matrix mode).

• groups_data (dict or DataFrame, optional) – Grouping variables for random effects (matrix mode). Keys should match RandomEffect.grouping names.

• family (str, default='gaussian') – Distribution family. Gaussian uses exact REML estimation; non-Gaussian families (e.g. ‘poisson’, ‘binomial’) are fitted with Penalized Quasi-Likelihood (PQL).

• covariance (str, default='unstructured') –

  Covariance structure for random effects. Options: - ‘identity’: Independent random effects (default) - ‘diagonal’: Heterogeneous variances, no correlation - ‘unstructured’: Full covariance matrix - ‘ar1’: Autoregressive AR(1) for temporal correlation - ‘compound_symmetry’ or ‘cs’: Exchangeable correlation - ‘exponential’: Spatial decay (requires coordinates) - ‘matern’: Matérn spatial covariance (requires coordinates)

  For temporal/longitudinal data, use ‘ar1’. For clustered data with equal correlations, use ‘compound_symmetry’.

• maxiter (int, default=100) – Maximum iterations for optimization.

• tol (float, default=1e-6) – Convergence tolerance.

• use_sparse (bool, default=False) – Whether to use sparse matrix operations for smooth terms. When True, uses sparse CSR format for basis matrices and sparse linear solvers. Provides significant speedup (10-100×) and memory reduction (6-8×) for large problems with B-spline basis functions. Currently only supported when smooth terms use B-spline basis.

• backend (str, default='numpy') – Computational backend: ‘numpy’, ‘torch’, or ‘jax’.

• device (str, optional) – Device for computation (for torch backend): ‘cpu’, ‘cuda’, ‘cuda:0’, etc. If None, uses CUDA if available, else CPU.

Returns:

result – Fitted GAMM result with coefficients, variance components, and diagnostics.

Return type:

GAMMResult

Raises:

ValueError – If random_effects provided but groups_data is None. If group variables not found in groups_data.

Examples

Formula mode (recommended):

>>> # Random intercept model
>>> import pandas as pd
>>> import numpy as np
>>> from aurora.models.gamm import fit_gamm
>>>
>>> # Generate data
>>> n_groups, n_per_group = 10, 20
>>> n = n_groups * n_per_group
>>> data = pd.DataFrame({
...     'y': np.random.randn(n),
...     'x1': np.random.randn(n),
...     'subject': np.repeat(np.arange(n_groups), n_per_group)
... })
>>>
>>> # Fit model with formula
>>> result = fit_gamm(
...     formula="y ~ x1 + (1 | subject)",
...     data=data,
...     covariance='identity'
... )
>>>
>>> # Access results
>>> print(result.beta_parametric)  # Fixed effects
>>> print(result.variance_components)  # Random effect variance

>>> # Random intercept + slope
>>> result = fit_gamm(
...     formula="y ~ x1 + (1 + x1 | subject)",
...     data=data,
...     covariance='unstructured'
... )
>>> print(result.variance_components)  # 2x2 covariance matrix

Matrix mode (advanced):

>>> from aurora.models.gamm import RandomEffect
>>>
>>> # Manually construct matrices
>>> groups = data['subject'].values
>>> x = data['x1'].values
>>> y = data['y'].values
>>> X = np.column_stack([np.ones(n), x])
>>>
>>> # Specify random effects manually
>>> re = RandomEffect(grouping='subject')
>>> result = fit_gamm(
...     y=y,
...     X=X,
...     random_effects=[re],
...     groups_data={'subject': groups},
...     covariance='identity'
... )
>>> print(result.beta_parametric)

Notes

This is a high-level interface that automatically constructs design matrices for random effects and calls the appropriate fitting function based on the family.

For Gaussian family, uses exact REML estimation. For non-Gaussian families, uses Penalized Quasi-Likelihood (PQL; Breslow & Clayton 1993) with uncorrected estimates.

Non-Gaussian likelihood quantities. For non-Gaussian families, log_likelihood is the conditional log-likelihood Σᵢ ℓ(yᵢ; μ̂ᵢ) evaluated at the BLUPs, not the marginal likelihood; aic/bic are conditional information criteria derived from it. Likelihood-ratio tests on fixed effects based on these quantities are not valid (the anova module issues a warning).

Limitations. offset and prior weights are not currently supported by fit_gamm; for models requiring them, use fit_glm (no random effects) instead.

Formula mode supports: - R-style formula syntax with lme4-style random effects - Automatic design matrix construction from DataFrames - Random intercepts: (1 | group) - Random slopes: (1 + x | group) - Nested effects: (1 | a/b) - Crossed effects: (1 | a) + (1 | b)

aurora.models.gamm.interface.fit_gamm_with_smooth(y, X_parametric, X_smooth, S_smooth, random_effects, groups_data, lambda_smooth=None, family='gaussian', covariance='unstructured', maxiter=100, tol=1e-6)

Fit GAMM with smooth terms (advanced interface).

Parameters:

• y (array-like, shape (n,)) – Response variable.

• X_parametric (array-like, shape (n, p)) – Design matrix for parametric fixed effects.

• X_smooth (dict of str -> ndarray) – Smooth term basis matrices by term name.

• S_smooth (dict of str -> ndarray) – Penalty matrices by term name.

• random_effects (list of RandomEffect) – Random effect specifications.

• groups_data (dict or DataFrame) – Grouping variables for random effects.

• lambda_smooth (dict of str -> float, optional) – Smoothing parameters by term name. If None, uses λ=1.0 for all.

• family (str, default='gaussian') – Distribution family.

• covariance (str, default='unstructured') – Covariance structure for random effects.

• maxiter (int, default=100) – Maximum iterations.

• tol (float, default=1e-6) – Convergence tolerance.

Returns:

result – Fitted GAMM result.

Return type:

GAMMResult

Examples

>>> # GAMM with smooth term + random intercept
>>> from aurora.smoothing import fit_cubic_spline
>>> from aurora.models.gamm import fit_gamm_with_smooth, RandomEffect
>>>
>>> # Data
>>> n = 100
>>> x_smooth = np.linspace(0, 2*np.pi, n)
>>> groups = np.repeat(np.arange(10), 10)
>>>
>>> # Build smooth basis
>>> import scipy.interpolate as interp
>>> k = 10
>>> knots = np.linspace(0, 2*np.pi, k-2)
>>> tck = interp.splrep(x_smooth, np.zeros(n), t=knots[1:-1], k=3)
>>> X_smooth_basis = interp.BSpline.design_matrix(x_smooth, tck[0], 3)
>>>
>>> # Parametric design
>>> X_para = np.ones((n, 1))
>>>
>>> # Random effects
>>> re = RandomEffect(grouping='subject')
>>>
>>> # Generate response
>>> y = 2.0 + np.sin(x_smooth) + np.random.randn(10)[groups] + np.random.randn(n)*0.3
>>>
>>> # Penalty matrix (second derivative)
>>> S = np.diag([0]*2 + [1]*(k-2))  # Penalize non-linear components
>>>
>>> result = fit_gamm_with_smooth(
...     y=y,
...     X_parametric=X_para,
...     X_smooth={'s(x)': X_smooth_basis},
...     S_smooth={'s(x)': S},
...     random_effects=[re],
...     groups_data={'subject': groups},
...     lambda_smooth={'s(x)': 0.1},
...     covariance='identity'
... )

Notes

This is the advanced interface for GAMMs with smooth terms. For basic usage without smooth terms, use fit_gamm() instead.

Smooth term construction (basis matrices and penalties) should be done using the smoothing module from Phase 3.

aurora.models.gamm.pql.fit_pql(X, Z, y, family, S=None, lambda_=0.0, psi_init=None, maxiter_outer=20, maxiter_inner=10, tol_outer=1e-4, tol_inner=1e-6, update_psi=True)

Fit a GLMM using Penalized Quasi-Likelihood.

Parameters:

• X (ndarray, shape (n, p)) – Fixed effects design matrix

• Z (ndarray, shape (n, q)) – Random effects design matrix

• y (ndarray, shape (n,)) – Response variable

• family (Family or str) – Distribution family (‘gaussian’, ‘poisson’, ‘binomial’)

• S (ndarray, shape (p, p), optional) – Penalty matrix for smooth terms (default: zeros)

• lambda (float, default=0.0) – Smoothing parameter

• psi_init (ndarray, shape (n_effects, n_effects), optional) – Initial variance-covariance matrix (default: identity)

• maxiter_outer (int, default=20) – Maximum outer iterations (variance updates)

• maxiter_inner (int, default=10) – Maximum inner iterations per outer (coefficient updates)

• tol_outer (float, default=1e-4) – Convergence tolerance for variance components

• tol_inner (float, default=1e-6) – Convergence tolerance for coefficients

• update_psi (bool, default=True) – Whether to update variance components (False for fixed Ψ)

• lambda_ (float)

Returns:

result – Fitted PQL model

Return type:

PQLResult

Notes

For Gaussian family, this reduces to standard LMM which could be solved more efficiently with REML. Use fit_gamm_gaussian() instead for Gaussian data.

The algorithm iterates between: 1. Updating (β, b) given Ψ via weighted least squares 2. Updating Ψ given (β, b) via empirical covariance

Convergence is determined by relative change in Ψ falling below tol_outer. With update_psi=False (fixed Ψ), convergence is determined by the inner (β, b) fixed point instead.

For the Gamma family, the working weights assume dispersion φ = 1; the estimated dispersion is not propagated into the weighting (documented limitation).

For the Binomial family, fitted probabilities numerically 0 or 1 indicate complete or quasi-complete separation; a warning is issued and the result is reported as not converged.

Examples

>>> import numpy as np
>>> from aurora.models.gamm.pql import fit_pql
>>> # Generate Poisson data with random intercepts
>>> n_groups, n_per_group = 5, 20
>>> n = n_groups * n_per_group
>>> groups = np.repeat(np.arange(n_groups), n_per_group)
>>> X = np.column_stack([np.ones(n), np.random.randn(n)])
>>> Z = np.eye(n_groups)[groups]  # Random intercept design
>>> b_true = np.random.randn(n_groups) * 0.5
>>> eta = X @ [1.0, 0.3] + Z @ b_true
>>> y = np.random.poisson(np.exp(eta))
>>> result = fit_pql(X, Z, y, family='poisson')
>>> print(result.converged)
True

aurora.models.gamm.pql_smooth.fit_pql_with_smooth(X_parametric, X_smooth_dict, Z, Z_info, y, family, S_smooth_dict, lambda_smooth=None, Psi_init=None, maxiter_outer=20, maxiter_inner=10, tol_outer=1e-4, tol_inner=1e-6, verbose=False)

Fit non-Gaussian GAMM with smooth terms using PQL.

Parameters:

• X_parametric (ndarray, shape (n, p_para)) – Parametric fixed effects design matrix.

• X_smooth_dict (dict[str, ndarray]) – Smooth term basis matrices {term_name: X_smooth}. Each X_smooth has shape (n, K_j) for K_j basis functions.

• Z (ndarray, shape (n, q)) – Random effects design matrix.

• Z_info (list[dict]) – Random effects metadata, one dict per random effect term. Each dict contains ‘n_levels’ and ‘type’ keys.

• y (ndarray, shape (n,)) – Response vector.

• family (str) – Distribution family: ‘poisson’, ‘binomial’, ‘gamma’.

• S_smooth_dict (dict[str, ndarray]) – Penalty matrices for each smooth term {term_name: S}. Each S has shape (K_j, K_j).

• lambda_smooth (dict[str, float], optional) – Smoothing parameters for each smooth term. If None, uses GCV.

• Psi_init (list[ndarray], optional) – Initial variance-covariance matrices for random effects.

• maxiter_outer (int, default=20) – Maximum outer iterations for variance components.

• maxiter_inner (int, default=10) – Maximum inner iterations for (β, b) update.

• tol_outer (float, default=1e-4) – Convergence tolerance for variance components.

• tol_inner (float, default=1e-6) – Convergence tolerance for coefficients.

• verbose (bool, default=False) – Print iteration progress.

Returns:

result – Dictionary containing: - ‘beta_parametric’: Parametric coefficients - ‘beta_smooth’: Dictionary of smooth coefficients by term - ‘random_effects’: Random effect estimates - ‘variance_components’: List of Ψ matrices - ‘smoothing_parameters’: Dictionary of λ by smooth term - ‘fitted_values’: Fitted values on the response scale (μ = g⁻¹(η)) - ‘linear_predictor’: Fitted linear predictor (η) - ‘converged’: Whether algorithm converged - ‘n_iterations_outer’: Number of outer iterations - ‘edf_smooth’: EDF by smooth term

Return type:

dict

Examples

>>> import numpy as np
>>> from aurora.models.gamm.pql_smooth import fit_pql_with_smooth
>>> from aurora.smoothing.splines.bspline import BSplineBasis
>>> # Generate data
>>> n = 200
>>> x = np.linspace(0, 1, n)
>>> groups = np.repeat(np.arange(20), 10)
>>> # Create smooth basis
>>> basis = BSplineBasis.create_knots(x, n_basis=10, degree=3)
>>> B = BSplineBasis(basis).basis_matrix(x)
>>> X_smooth = {'s(x)': B}
>>> S = {'s(x)': BSplineBasis(basis).penalty_matrix(order=2)}
>>> # Random effects
>>> Z = np.eye(n)[:, groups]
>>> # Fit model
>>> result = fit_pql_with_smooth(
...     X_parametric=np.ones((n, 1)),
...     X_smooth_dict=X_smooth,
...     Z=Z,
...     Z_info=[{'n_levels': 20, 'type': 'intercept'}],
...     y=y,
...     family='poisson',
...     S_smooth_dict=S,
... )

Result

class aurora.models.gamm.fitting.GAMMResult(coefficients, beta_parametric, beta_smooth, random_effects, variance_components, covariance_params, residual_variance, smoothing_parameters, edf_total, edf_parametric, edf_smooth, fitted_values, residuals, log_likelihood, aic, bic, converged, n_iterations, n_obs, n_groups, family, _X_parametric=None, _X_smooth=None, _Z=None, _Z_info=None, _y=None)

Bases: object

Result from GAMM fitting.

Variables:

• coefficients (ndarray) – Combined coefficients [β_parametric, β_smooth, b_random].

• beta_parametric (ndarray) – Parametric (fixed) effect coefficients.

• beta_smooth (dict[str, ndarray]) – Smooth term coefficients by term name.

• random_effects (dict[str, ndarray]) – Random effect coefficients by grouping variable.

• variance_components (list[ndarray]) – List of variance-covariance matrices Ψ, one per random effect term. For single term models, this is a list with one element.

• covariance_params (list[ndarray] | None) – List of raw covariance parameters for each random effect term. For structured covariance (AR1, CS, etc.), contains transformed parameters. For example, AR1 params are [log(σ²), arctanh(ρ)]. None for unstructured covariance.

• residual_variance (float) – Residual variance σ².

• smoothing_parameters (dict[str, float] | None) – Smoothing parameters λ by smooth term (if estimated).

• edf_total (float) – Total effective degrees of freedom.

• edf_parametric (float) – EDF for parametric terms.

• edf_smooth (dict[str, float]) – EDF by smooth term.

• fitted_values (ndarray) – Fitted values on the response scale (μ̂; for the Gaussian identity link this equals η̂ = Xβ + Zb).

• residuals (ndarray) – Response residuals y − μ̂.

• log_likelihood (float) – Log-likelihood (or REML log-likelihood). For non-Gaussian PQL fits this is the conditional log-likelihood at the BLUPs, not the marginal likelihood.

• aic (float) – Akaike Information Criterion.

• bic (float) – Bayesian Information Criterion.

• converged (bool) – Whether fitting converged.

• n_iterations (int) – Number of iterations.

• n_obs (int) – Number of observations.

• n_groups (int) – Number of groups.

• family (str) – Distribution family.

Parameters:

• coefficients (NDArray[np.floating])

• beta_parametric (NDArray[np.floating])

• beta_smooth (dict[str, NDArray[np.floating]])

• random_effects (dict[Any, dict[Any, NDArray[np.floating]]])

• variance_components (list[NDArray[np.floating]])

• covariance_params (list[NDArray[np.floating]] | None)

• residual_variance (float)

• smoothing_parameters (dict[str, float] | None)

• edf_total (float)

• edf_parametric (float)

• edf_smooth (dict[str, float])

• fitted_values (NDArray[np.floating])

• residuals (NDArray[np.floating])

• log_likelihood (float)

• aic (float)

• bic (float)

• converged (bool)

• n_iterations (int)

• n_obs (int)

• n_groups (int)

• family (str)

• _X_parametric (NDArray[np.floating] | None)

• _X_smooth (dict[str, NDArray[np.floating] | spmatrix] | None)

• _Z (NDArray[np.floating] | None)

• _Z_info (list[dict] | None)

• _y (NDArray[np.floating] | None)

coefficients: NDArray[np.floating]

beta_parametric: NDArray[np.floating]

beta_smooth: dict[str, NDArray[np.floating]]

random_effects: dict[Any, dict[Any, NDArray[np.floating]]]

variance_components: list[NDArray[np.floating]]

covariance_params: list[NDArray[np.floating]] | None

residual_variance: float

smoothing_parameters: dict[str, float] | None

edf_total: float

edf_parametric: float

edf_smooth: dict[str, float]

fitted_values: NDArray[np.floating]

residuals: NDArray[np.floating]

log_likelihood: float

aic: float

bic: float

converged: bool

n_iterations: int

n_obs: int

n_groups: int

family: str

predict(include_random=True)

Make predictions using the fitted model.

Parameters:

include_random (bool, default=True) – Whether to include random effects in predictions. If True, returns fitted values (fixed + random effects). If False, returns only fixed effects (population-level predictions).

Returns:

predictions – Model predictions.

Return type:

ndarray

Examples

>>> # Population-level predictions (fixed effects only)
>>> pred_fixed = result.predict(include_random=False)
>>>
>>> # Individual predictions (fixed + random effects)
>>> pred_full = result.predict(include_random=True)

summary()

Generate a formatted summary of the model fit.

Returns:

summary_str – Formatted summary string with fixed effects, random effects, and model fit statistics.

Return type:

str

Examples

>>> result = fit_gamm(formula='y ~ x + (1|subject)', data=df)
>>> print(result.summary())

Random Effects

class aurora.models.gamm.random_effects.RandomEffect(grouping, variables=<factory>, include_intercept=True, covariance='unstructured')

Bases: object

Specification of a random effect term.

Random effects model within-group correlation in hierarchical or longitudinal data. Each random effect term specifies: - Which variables have random coefficients - What grouping structure to use - Whether to include a random intercept - The covariance structure among random effects

Parameters:

• grouping (str or int) – Grouping variable name or column index. All observations with the same value of this variable belong to the same group.

• variables (tuple of (str or int), optional) – Variables with random slopes. If empty, only includes random intercept (if include_intercept=True).

• include_intercept (bool, default=True) – Whether to include a random intercept (1 | group).

• covariance (str, default='unstructured') –

  Covariance structure among random effects. Options: - ‘unstructured’: Full covariance matrix (q(q+1)/2 parameters) - ‘diagonal’: Independent random effects (q parameters) - ‘identity’: Equal variances, no correlation (1 parameter) - ‘ar1’: Autoregressive AR(1) for temporal correlation - ‘compound_symmetry’ or ‘cs’: Exchangeable correlation (ICC model) - ‘exponential’: Spatial exponential decay (requires coordinates) - ‘matern’: Matérn spatial covariance (requires coordinates)

  For temporal/longitudinal data, use ‘ar1’. For clustered data with equal within-cluster correlation, use ‘compound_symmetry’.

Variables:

n_effects (int) – Total number of random effects (intercept + slopes)

Examples

>>> # Random intercept: (1 | subject)
>>> re1 = RandomEffect(grouping='subject')
>>> re1.n_effects
1

>>> # Random intercept + slope: (1 + time | subject)
>>> re2 = RandomEffect(grouping='subject', variables=('time',))
>>> re2.n_effects
2

>>> # Random slopes only: (time + age | subject)
>>> re3 = RandomEffect(
...     grouping='subject',
...     variables=('time', 'age'),
...     include_intercept=False
... )
>>> re3.n_effects
2

>>> # Diagonal covariance (independent effects)
>>> re4 = RandomEffect(
...     grouping='subject',
...     variables=('time',),
...     covariance='diagonal'
... )

Notes

The total number of random effects is:

q = (1 if include_intercept else 0) + len(variables)

Nested random effects like (1 | clinic/subject) should be specified as two separate RandomEffect instances with appropriate grouping.

Crossed random effects like (1 | subject) + (1 | clinic) are specified as multiple RandomEffect instances in a list.

grouping: str | int

variables: tuple[str | int, ...]

include_intercept: bool = True

covariance: Literal['unstructured', 'diagonal', 'identity', 'ar1', 'compound_symmetry', 'cs', 'exponential', 'matern'] = 'unstructured'

__post_init__()

Validate random effect specification.

property n_effects: int

Number of random effects in this term.

Returns:

Number of random effects (intercept + slopes)

Return type:

int

__repr__()

String representation.

Return type:

str

Covariance Structures

Covariance structures for random effects.

This module provides classes for different covariance parameterizations used in random effects modeling.

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS.

class aurora.models.gamm.covariance.CovarianceStructure

Bases: ABC

Base class for random effect covariance structures.

A covariance structure defines how to parameterize the variance-covariance matrix Ψ for random effects. Different structures impose different constraints (e.g., diagonal, identity) to reduce the number of parameters.

n_parameters(n_effects)

Number of free parameters for given number of random effects

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Build covariance matrix from parameter vector

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract parameter vector from covariance matrix

Parameters:

psi (ndarray)

Return type:

ndarray

abstractmethod n_parameters(n_effects)

Number of free parameters.

Parameters:

n_effects (int) – Number of random effects (dimension of Ψ)

Returns:

Number of free parameters to specify Ψ

Return type:

int

abstractmethod construct_psi(params, n_effects)

Construct covariance matrix from parameters.

Parameters:

• params (ndarray) – Parameter vector (length = n_parameters(n_effects))

• n_effects (int) – Number of random effects

Returns:

psi – Variance-covariance matrix (symmetric, positive definite)

Return type:

ndarray, shape (n_effects, n_effects)

abstractmethod extract_params(psi)

Extract parameters from covariance matrix.

Parameters:

psi (ndarray, shape (q, q)) – Variance-covariance matrix

Returns:

params – Parameter vector

Return type:

ndarray

class aurora.models.gamm.covariance.UnstructuredCovariance

Bases: CovarianceStructure

Unstructured covariance (full matrix).

Allows all variances and covariances to vary freely. For q random effects, requires q(q+1)/2 parameters.

Parameterization uses Cholesky decomposition: Ψ = LL’ This ensures Ψ is positive definite.

Parameters:

None

Examples

>>> cov = UnstructuredCovariance()
>>> cov.n_parameters(2)  # 2 random effects
3

>>> # Parameters: [L11, L21, L22]
>>> params = np.array([1.0, 0.5, 0.8])
>>> psi = cov.construct_psi(params, n_effects=2)
>>> psi
array([[1.  , 0.5 ],
       [0.5 , 0.89]])

>>> # Extract back
>>> extracted = cov.extract_params(psi)
>>> np.allclose(extracted, params)
True

Notes

The Cholesky parameterization: - Ensures positive definiteness - Provides unconstrained optimization (except L_ii > 0) - Standard approach in lme4 and other mixed model software

n_parameters(n_effects)

Number of parameters = q(q+1)/2.

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Ψ = LL’ from Cholesky parameters.

Parameters are ordered as lower triangle of L: [L11, L21, L22, L31, L32, L33, …]

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract Cholesky parameters from Ψ.

Computes L such that Ψ = LL’, then extracts lower triangle.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.DiagonalCovariance

Bases: CovarianceStructure

Diagonal covariance (independent random effects).

Assumes random effects are independent (zero correlations). For q random effects, requires q parameters (variances only).

Parameters are log-transformed variances to ensure positivity: σ²_i = exp(params[i])

Examples

>>> cov = DiagonalCovariance()
>>> cov.n_parameters(3)
3

>>> # Parameters: log-variances
>>> params = np.array([0.0, 0.5, 1.0])  # log(σ²)
>>> psi = cov.construct_psi(params, n_effects=3)
>>> psi
array([[1.        , 0.        , 0.        ],
       [0.        , 1.64872127, 0.        ],
       [0.        , 0.        , 2.71828183]])

>>> np.diag(psi)  # Variances
array([1.        , 1.64872127, 2.71828183])

Notes

Log transformation: - Ensures variances are positive - Provides unconstrained optimization - Standard in many mixed model packages

n_parameters(n_effects)

Number of parameters = q (one variance per effect).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct diagonal Ψ from log-variances.

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract log-variances from diagonal Ψ.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.IdentityCovariance

Bases: CovarianceStructure

Identity covariance (equal variances, no correlation).

All random effects have the same variance and are uncorrelated. For q random effects, requires only 1 parameter: σ²

Ψ = σ² I_q

Parameter is log(σ²) to ensure positivity.

Examples

>>> cov = IdentityCovariance()
>>> cov.n_parameters(3)
1

>>> # Single parameter: log(σ²)
>>> params = np.array([0.5])  # log(σ²)
>>> psi = cov.construct_psi(params, n_effects=3)
>>> psi
array([[1.64872127, 0.        , 0.        ],
       [0.        , 1.64872127, 0.        ],
       [0.        , 0.        , 1.64872127]])

>>> # All variances are equal
>>> np.allclose(np.diag(psi), np.exp(0.5))
True

Notes

This is the most constrained structure, useful when: - Sample size is small - Random effects are believed to be exchangeable - Computational efficiency is important

n_parameters(n_effects)

Number of parameters = 1 (single variance).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct σ²I from log-variance.

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

extract_params(psi)

Extract log(σ²) from Ψ = σ²I.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.AR1Covariance

Bases: CovarianceStructure

Autoregressive order 1 (AR(1)) covariance structure.

Models temporal correlation where observations closer in time are more correlated: Cov(b_t, b_s) = σ² ρ^|t-s|

For q random effects (time points), requires 2 parameters: - params[0]: log(σ²) - log-variance - params[1]: arctanh(ρ) - transformed correlation (ensures ρ ∈ (-1, 1))

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS. Ch. 5.

[2]

Diggle et al. (2002). Analysis of Longitudinal Data, 2nd ed.

Examples

>>> cov = AR1Covariance()
>>> cov.n_parameters(4)
2

>>> # Parameters: [log(σ²), arctanh(ρ)]
>>> params = np.array([0.0, 0.5])  # σ²=1, ρ≈0.46
>>> psi = cov.construct_psi(params, n_effects=4)
>>> # Creates correlation matrix with AR(1) structure

Notes

The AR(1) structure is ideal for: - Equally-spaced longitudinal data - Time series with decaying correlations - Repeated measures where adjacent observations are most correlated

The inverse of AR(1) covariance matrix is tridiagonal, enabling efficient computation for large matrices.

n_parameters(n_effects)

Number of parameters = 2 (variance and correlation).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct AR(1) covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), arctanh(ρ)] - transformed variance and correlation

• n_effects (int) – Number of time points/random effects

Returns:

psi – AR(1) covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from AR(1) covariance matrix.

Estimates σ² from diagonal and ρ from first off-diagonal.

Parameters:

psi (ndarray)

Return type:

ndarray

inverse(params, n_effects)

Efficient tridiagonal inverse for AR(1) structure.

The inverse of an AR(1) matrix is tridiagonal, which allows for O(n) computation instead of O(n³).

Parameters:

• params (ndarray)

• n_effects (int)

Return type:

ndarray

class aurora.models.gamm.covariance.CompoundSymmetryCovariance

Bases: CovarianceStructure

Compound symmetry (exchangeable) covariance structure.

Models constant correlation between all pairs of observations: - Diagonal: Var(b_i) = σ² - Off-diagonal: Cov(b_i, b_j) = σ²ρ for i ≠ j

For q random effects, requires 2 parameters: - params[0]: log(σ²) - log-variance - params[1]: logit(ρ_scaled) - transformed correlation

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS.

[2]

Diggle et al. (2002). Analysis of Longitudinal Data, 2nd ed.

Examples

>>> cov = CompoundSymmetryCovariance()
>>> cov.n_parameters(5)
2

>>> params = np.array([0.0, 0.0])  # σ²=1, ρ=0.5
>>> psi = cov.construct_psi(params, n_effects=3)
>>> # Creates matrix with equal variances and equal correlations

Notes

Compound symmetry is equivalent to: - Random intercept model with iid residuals - Exchangeable correlation (ICC model)

Constraint: ρ > -1/(q-1) for positive definiteness, where q is the number of effects.

n_parameters(n_effects)

Number of parameters = 2 (variance and correlation).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct compound symmetry covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), logit((ρ + 1/(q-1))/(1 + 1/(q-1)))]

• n_effects (int) – Number of random effects

Returns:

psi – Compound symmetry covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from compound symmetry matrix.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.ToeplitzCovariance(band=2)

Bases: CovarianceStructure

Toeplitz (banded) covariance structure for temporal data.

Toeplitz matrices have constant values along diagonals, meaning the covariance between observations depends only on their lag |t-s|:

Cov(b_t, b_s) = r_{|t-s|}

This is a generalization of AR(1) that allows arbitrary values at each lag, rather than assuming exponential decay.

For q random effects (time points), requires (band+1) parameters where band is the number of non-zero off-diagonal bands: - params[0]: log(r_0) - log-variance - params[k]: arctanh(r_k / r_0) for k = 1, …, band (transformed correlations)

Parameters:

band (int, default=2) – Number of off-diagonal bands to include. Higher values allow more complex temporal patterns but require more parameters. - band=1: First-order correlations only (like AR(1) but with free parameter) - band=2: First and second-order correlations

References

[1]

Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS. Ch. 5.

[2]

Pourahmadi (1999). “Joint mean-covariance models with applications to longitudinal data.” Biometrika, 86(3), 677-690.

Examples

>>> cov = ToeplitzCovariance(band=2)
>>> cov.n_parameters(5)
3

>>> params = np.array([0.0, 0.3, 0.1])  # log(σ²)=0, r_1/σ²≈0.29, r_2/σ²≈0.10
>>> psi = cov.construct_psi(params, n_effects=4)

Notes

Toeplitz covariance is appropriate for: - Equally-spaced longitudinal data - Time series where correlations don’t follow AR(p) exactly - Exploratory analysis before fitting structured models

The Toeplitz structure is a special case of a block structure and allows efficient computation using banded matrix algorithms.

__init__(band=2)

Initialize Toeplitz covariance.

Parameters:

band (int, default=2) – Number of off-diagonal bands (lags) to model.

n_parameters(n_effects)

Number of parameters = min(band, n_effects-1) + 1.

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Toeplitz covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), arctanh(ρ_1), arctanh(ρ_2), …] - transformed parameters

• n_effects (int) – Number of time points/random effects

Returns:

psi – Toeplitz covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from Toeplitz covariance matrix.

Estimates variance from diagonal and correlations from off-diagonals.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.ExponentialSpatialCovariance(coordinates=None)

Bases: CovarianceStructure

Exponential spatial covariance structure.

Models spatial correlation that decays exponentially with distance: Cov(b_i, b_j) = σ² exp(-d_ij / φ)

where d_ij is the Euclidean distance between locations i and j, and φ is the range parameter.

For geostatistical data, requires coordinates at construction time. Parameters: - params[0]: log(σ²) - log-variance (sill) - params[1]: log(φ) - log-range parameter

References

[1]

Diggle & Ribeiro (2007). Model-based Geostatistics.

[2]

Cressie (1993). Statistics for Spatial Data.

[3]

Rue & Held (2005). Gaussian Markov Random Fields.

Examples

>>> coords = np.array([[0, 0], [1, 0], [0, 1], [1, 1]])
>>> cov = ExponentialSpatialCovariance(coordinates=coords)
>>> params = np.array([0.0, 0.5])  # σ²=1, φ≈1.65
>>> psi = cov.construct_psi(params, n_effects=4)

Notes

The exponential covariance is a member of the Matérn family with smoothness ν = 0.5. It produces relatively rough spatial surfaces.

The range parameter φ approximately equals the distance at which the correlation drops to ~37% (= 1/e).

Practical range (correlation ≈ 5%) is approximately 3φ.

Parameters:

coordinates (np.ndarray | None)

__init__(coordinates=None)

Initialize with spatial coordinates.

Parameters:

coordinates (ndarray (n, d), optional) – Spatial coordinates for n locations in d dimensions. If not provided, assumes 1D equally-spaced locations.

n_parameters(n_effects)

Number of parameters = 2 (variance and range).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct exponential spatial covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), log(φ)] - log-variance and log-range

• n_effects (int) – Number of spatial locations

Returns:

psi – Spatial covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from spatial covariance matrix.

Uses method of moments estimation based on variogram fitting.

Parameters:

psi (ndarray)

Return type:

ndarray

class aurora.models.gamm.covariance.MaternCovariance(coordinates=None, nu=1.5)

Bases: CovarianceStructure

Matérn covariance structure for spatial data.

The Matérn family is the most commonly used in geostatistics due to its flexibility in modeling different degrees of smoothness.

Cov(d) = σ² × (2^(1-ν) / Γ(ν)) × (√(2ν) d/φ)^ν × K_ν(√(2ν) d/φ)

where K_ν is the modified Bessel function of the second kind.

Special cases: - ν = 0.5: Exponential covariance - ν = 1.5: Once differentiable - ν = 2.5: Twice differentiable - ν → ∞: Gaussian (squared exponential)

Parameters (for fixed ν): - params[0]: log(σ²) - log-variance (sill) - params[1]: log(φ) - log-range parameter

References

[1]

Matérn (1960). Spatial Variation.

[2]

Stein (1999). Interpolation of Spatial Data.

[3]

Rasmussen & Williams (2006). Gaussian Processes for ML.

Examples

>>> coords = np.array([[0, 0], [1, 0], [2, 0]])
>>> cov = MaternCovariance(coordinates=coords, nu=1.5)
>>> params = np.array([0.0, 0.0])  # σ²=1, φ=1
>>> psi = cov.construct_psi(params, n_effects=3)

Parameters:

• coordinates (np.ndarray | None)

• nu (float)

__init__(coordinates=None, nu=1.5)

Initialize Matérn covariance.

Parameters:

• coordinates (ndarray (n, d), optional) – Spatial coordinates. If None, assumes 1D equally-spaced.

• nu (float, default=1.5) – Smoothness parameter. Common choices: 0.5, 1.5, 2.5

n_parameters(n_effects)

Number of parameters = 2 (variance and range).

Parameters:

n_effects (int)

Return type:

int

construct_psi(params, n_effects)

Construct Matérn covariance matrix.

Parameters:

• params (ndarray) – [log(σ²), log(φ)]

• n_effects (int) – Number of spatial locations

Returns:

psi – Matérn covariance matrix

Return type:

ndarray (n_effects, n_effects)

extract_params(psi)

Extract parameters from Matérn covariance matrix.

Parameters:

psi (ndarray)

Return type:

ndarray

aurora.models.gamm.covariance.get_covariance_structure(structure, **kwargs)

Get covariance structure instance by name.

Parameters:

• structure ({'unstructured', 'diagonal', 'identity', 'ar1',) – ‘compound_symmetry’, ‘cs’, ‘exponential’, ‘matern’} Covariance structure name

• **kwargs (dict) – Additional arguments for specific structures: - coordinates: ndarray for spatial structures - nu: smoothness for Matérn

Returns:

Covariance structure instance

Return type:

CovarianceStructure

Raises:

ValueError – If structure name is not recognized

Examples

>>> cov = get_covariance_structure('unstructured')
>>> isinstance(cov, UnstructuredCovariance)
True

>>> cov = get_covariance_structure('ar1')
>>> isinstance(cov, AR1Covariance)
True

>>> coords = np.array([[0, 0], [1, 0], [0, 1]])
>>> cov = get_covariance_structure('exponential', coordinates=coords)
>>> isinstance(cov, ExponentialSpatialCovariance)
True

>>> cov = get_covariance_structure('matern', coordinates=coords, nu=2.5)
>>> cov.nu
2.5