spacr.mixed_gpu¶

Fit profiled REML mixed models with PyTorch on CPU or CUDA.

The backend fits the same nested random-intercept model as statsmodels MixedLM:

y = X b + sum_k Z_k u_k + e, u_k ~ N(0, sigma^2 theta_k I),
e ~ N(0, sigma^2 I)

The group intercept supplies one Z_k and each variance component supplies another, nested within the outer group as in MixedLM.vc_formula. For example, vc_formula={'rowID': '0 + C(rowID)'} with groups=gene represents (1 | gene:rowID), not (1 | rowID).

The implementation evaluates the profiled REML deviance from Bates et al. (2015, equation 41):

d(theta) = log|Lambda’ Z’Z Lambda + I| + log|X’ W^-1 X|

(n - p) [1 + log(2 pi r^2 / (n - p))]

Here Lambda = diag(sqrt(theta)), W = I + Z Lambda^2 Z', and r^2 is the minimized penalized residual sum of squares. Device-side cross-products Z'Z, Z'X, Z'y, X'X, X'y, and y'y are computed once. Each optimizer evaluation then requires a q x q Cholesky factorization and triangular solves. Autograd differentiates the deviance and L-BFGS optimizes log(theta), which enforces positive variance ratios.

Reference measurements explain where acceleration is expected. Statsmodels MixedLM took 54 times as long as OLS on a 40-gene screen and 67 times as long on an 80-gene screen. At q = 1212, the Cholesky step took 204 ms on the CPU and 7.69 ms on an RTX 3090. For a design with 1,830 observations, 387 genes, 710 guides, 388 fixed effects, and 1,097 random levels, end-to-end times were 11.3 s for statsmodels, 0.80 s for this backend on CPU, and 0.47 s on the RTX 3090. Use benchmark_against_statsmodels() to measure the supported backends on another design and device.

Selecting a CUDA device never falls back silently to CPU. If PyTorch or CUDA is unavailable, MixedBackendUnavailable explains what is missing.

Exceptions¶

MixedBackendUnavailable

Indicate that PyTorch or the requested compute device is unavailable.

Classes¶

TorchMixedResults

Store a PyTorch REML fit with statsmodels-compatible result fields.

Functions¶

`available_memory`(→ int)	Return the bytes currently available on a CPU or CUDA device.
`benchmark_against_statsmodels`(y, X, groups[, vc, device])	Fit the same mixed model with statsmodels and PyTorch and compare them.
`cuda_available`(→ bool)	Return whether PyTorch can use a CUDA device at call time.
`describe_device`(→ str)	Describe the installed PyTorch and CUDA device for logs or tooltips.
`design_bytes`(→ int)	Return the byte size of a dense `n x q` random-effects design.
`fit_mixed_reml_torch`(y, X, groups[, vc, device, ...])	Fit `y ~ X + (1 \| groups) + variance components` by profiled REML.
`mixedlm_torch`(formula, data, groups[, vc_formula, device])	Fit a formula mixed model with the PyTorch REML backend.
`resolve_device`([device])	Resolve a device string without silently substituting the CPU.
`torch_available`(→ bool)	Return whether PyTorch is importable without importing it.

Module Contents¶

exception spacr.mixed_gpu.MixedBackendUnavailable[source]¶

Bases: RuntimeError

Indicate that PyTorch or the requested compute device is unavailable.

Initialize self. See help(type(self)) for accurate signature.

class spacr.mixed_gpu.TorchMixedResults[source]¶

Store a PyTorch REML fit with statsmodels-compatible result fields.

fe_params, pvalues, random_effects, converged, residuals, and fitted values use the meanings expected by spacr.ml.fit_mixed_model(). Variance entries in params and theta are relative to residual variance; cov_re and vcomp are absolute variances on the response scale. device, fit_seconds, n_deviance_evals, and gradient_norm record backend diagnostics.

Parameters:

fe_params – Estimated fixed-effect coefficients indexed by fixed-design column.
bse_fe – Standard errors of the fixed-effect coefficients, indexed like fe_params.
params – Combined statsmodels-shaped vector of fixed-effect coefficients followed by variance ratios.
bse – Standard errors aligned with params; variance-ratio entries are nan because this backend does not estimate their standard errors.
tvalues – Fixed-effect z statistics aligned with params; variance-ratio entries are nan.
pvalues – Two-sided standard-normal p-values aligned with params; variance-ratio entries are nan.
scale – Estimated residual variance on the response scale.
cov_re – One-by-one covariance matrix for the outer-group random intercept on the response scale.
vcomp – Absolute variances of the nested variance components, in the input component order.
random_effects – Conditional random-effect estimates keyed by outer-group level, with statsmodels-compatible component labels.
resid – Conditional residuals, y - fittedvalues, in input-observation order.
fittedvalues – Conditional fitted values including fixed and random effects, in input-observation order.
converged – Whether the final log-variance gradient norm is below the backend’s reported-convergence threshold of 1e-3.
llf – Restricted log-likelihood at the reported solution.
n_obs – Number of observations included in the fit.
k_fe – Number of fixed-effect coefficients.
backend – Backend identifier recorded in logs and run metadata.
device – PyTorch device used for the fit.
fit_seconds – Wall-clock seconds spent in L-BFGS optimization.
n_deviance_evals – Number of profiled-deviance evaluations performed during optimization and finalization.
theta – Variance ratios relative to residual variance, ordered as the outer-group intercept followed by nested components.
gradient_norm – Euclidean norm of the final gradient with respect to log variance ratios.

summary_line() → str[source]¶: Return a one-line fit summary for the run log.

property df_resid: int[source]¶: Return residual degrees of freedom as observations minus effects.

spacr.mixed_gpu.available_memory(device: str = 'cpu') → int[source]¶: Return the bytes currently available on a CPU or CUDA device.

spacr.mixed_gpu.benchmark_against_statsmodels(y, X, groups, vc=None, *, device: str = GPU_DEVICE)[source]¶

Fit the same mixed model with statsmodels and PyTorch and compare them.

Parameters:

y (array-like) – Response values with shape (n_observations,).
X (array-like) – Fixed-effects design with shape (n_observations, n_fixed_effects).
groups (array-like) – Outer grouping labels with shape (n_observations,).
vc (mapping of str to array-like, optional) – Nested variance-component labels.
device (str, optional) – Device used by the PyTorch fit.

Returns:

dict – Timings, speedup, maximum fixed-effect and variance-component disagreement, residual-scale disagreement, resolved device, and PyTorch deviance-evaluation count.

spacr.mixed_gpu.cuda_available() → bool[source]¶

Return whether PyTorch can use a CUDA device at call time.

This function imports PyTorch and is intended for fit-time validation, not lightweight settings-panel construction.

spacr.mixed_gpu.describe_device() → str[source]¶: Describe the installed PyTorch and CUDA device for logs or tooltips.

spacr.mixed_gpu.design_bytes(n: int, q: int, *, itemsize: int = 8) → int[source]¶

Return the byte size of a dense n x q random-effects design.

Parameters:

n – number of observation rows in the design.
q – number of random-effect columns in the design.
itemsize – bytes used by each matrix element.

Exact rather than estimated: the shape is known before the matrix exists.

spacr.mixed_gpu.fit_mixed_reml_torch(y, X, groups, vc=None, *, device: str = GPU_DEVICE, max_iter: int = 400, verbose: bool = False)[source]¶

Fit y ~ X + (1 | groups) + variance components by profiled REML.

Parameters:

y (array-like) – Response values with shape (n_observations,).
X (array-like) – Fixed-effects design with shape (n_observations, n_fixed_effects). A DataFrame preserves its column names in TorchMixedResults.fe_params.
groups (array-like) – Outer grouping labels with shape (n_observations,). A random intercept is always included, equivalent to re_formula='1'.
vc (mapping of str to array-like, optional) – Additional grouping labels. Each entry defines one variance component nested within groups, matching statsmodels vc_formula semantics.
device (str, optional) – PyTorch device. Requesting CUDA raises instead of falling back to CPU when no CUDA device is available.
max_iter (int, optional) – Maximum L-BFGS iterations. Typical screen-sized fits use 20–40; the default of 400 is a safety limit.
verbose (bool, optional) – Print the deviance at each optimizer evaluation.

Returns:

TorchMixedResults – Fixed effects, variance components, conditional fitted values and residuals, BLUPs, convergence diagnostics, and timing information.

Raises:

MixedBackendUnavailable – If PyTorch or the requested device is unavailable.
ValueError – If input dimensions disagree or the fixed-effects design is not identifiable.
MemoryError – If the dense random-effects design exceeds the configured share of available device memory.

Notes

Validation against statsmodels MixedLM(...).fit() produced the following differences:

quantity	nested fixture (1,620 rows)	screen-shaped fixture (1,830 rows, p=388, q=1,097)
fixed effects	1.20e-7 absolute	2.39e-4 absolute (2.0e-4 of one SE)
variance components	1.29e-3 relative	2.02e-4 relative
residual scale	4.12e-6 relative	1.94e-5 relative
standard errors	3.88e-4 relative	1.04e-3 median; 1.80e-2 maximum across 388
guide BLUPs	7.64e-5 absolute	not evaluated

Both backends maximize the same REML criterion. On the screen-shaped fixture, this fit ended at gradient norm 1.3e-11 and log-likelihood -804.143968098; the default statsmodels fit returned -804.143968682. Setting statsmodels gtol=1e-12 reduced fixed-effect disagreement to 1.52e-8 absolute and fixture variance-component disagreement to 1.02e-4 relative.

Reference end-to-end times for that fixture were 11.3 s for statsmodels, 0.80 s for this backend on CPU, and 0.47 s on an RTX 3090. Under concurrent CPU load, statsmodels took 21.3 s and the CUDA backend 0.57 s.

spacr.mixed_gpu.mixedlm_torch(formula, data, groups, vc_formula=None, *, device: str = GPU_DEVICE, **kwargs)[source]¶

Fit a formula mixed model with the PyTorch REML backend.

The call shape parallels statsmodels.formula.api.mixedlm() so spacr.ml.fit_mixed_model() can select either backend.

Parameters:

formula (str) – Patsy fixed-effects formula, such as 'response ~ predictor'.
data (pandas.DataFrame) – Frame in which the formula and grouping columns are evaluated.
groups (str or array-like) – Outer grouping column or one label per row.
vc_formula (mapping of str to str, optional) – Variance components in the supported form '0 + C(column)'.
device (str, optional) – PyTorch compute device.
**kwargs – Additional arguments passed to fit_mixed_reml_torch().

Returns:

TorchMixedResults – Fitted mixed-model results.

Raises:

ValueError – If a variance-component formula is unsupported or names an absent column.
MixedBackendUnavailable – If PyTorch or the requested device is unavailable.

spacr.mixed_gpu.resolve_device(device: str = GPU_DEVICE)[source]¶

Resolve a device string without silently substituting the CPU.

Parameters:: device (str, optional) – 'cuda', 'cpu', or another PyTorch device string.
Returns:: torch.device – Validated compute device.
Raises:: MixedBackendUnavailable – If PyTorch is not installed or a CUDA device was requested but is not available.

spacr.mixed_gpu.torch_available() → bool[source]¶: Return whether PyTorch is importable without importing it.

Nested helpers¶

fit_mixed_reml_torch._solve(log_theta)¶

The profiled REML deviance at theta = exp(log_theta).

Returns the deviance plus everything the caller needs afterwards, so the final quantities come from the same factorisation that scored the final theta rather than from a re-solve that could differ.

spacr/mixed_gpu.py:551

fit_mixed_reml_torch.closure()¶

Evaluate and differentiate the captured LBFGS objective.

Returns:: differentiable profiled REML deviance at the current captured log-theta. Existing optimizer gradients are cleared, the deviance is backpropagated, and verbose mode reports its numeric value with the exponentiated variance parameters.

spacr/mixed_gpu.py:585