torch-openreml
A PyTorch-based library for AI-REML estimation of linear mixed models.
Overview¶
torch-openreml fits linear mixed-effects models using the Average Information REML (AI-REML) algorithm on a PyTorch backend. It supports flexible specification of covariance structures through a modular system of matrices and operators, along with automatic or manual gradients and optional parameter transformations for constrained estimation.
Unlike traditional mixed-model software, it does not provide a formula interface. Instead, users define the fixed- and random-effects design matrices and covariance structures directly in code. The library is focused purely on the computational and optimization backend rather than model specification syntax.
Features¶
Torch-based backend:
Built on PyTorch, supporting execution on CPU, GPU, and other available accelerators.
AI-REML estimation engine
Variance component estimation using the Average Information REML (AI-REML) quasi-Newton optimization framework.
Extensible covariance structure
Composable covariance structures and operators from built-in and user-defined components.
Hybrid differentiation support
Support automatic differentiation and manually specified gradients.
Composable parameter transformations
Configurable, chainable transformation pipelines for flexible parameterization.
Installation¶
pip install torch-openreml
Dependencies: torch, pandas, tqdm (Python 3.12).
Getting Started¶
Dataset¶
To illustrate a quick start with the library, we begin by fitting a mixed-effects model using
the john_alpha dataset. This dataset contains field trial data from a resolvable alpha lattice design
conducted at Craibstone near Aberdeen.
It consists of 72 observations and 7 variables. In this example, we use yield (dry matter yield) as the
response variable, and rep (replicate identifier), block (incomplete block within replicate),
and gen (genotype or variety identifier) as covariates.
plot rep block gen yield row col
0 1 R1 B1 G11 4.1172 1 1
1 2 R1 B1 G04 4.4461 2 1
2 3 R1 B1 G05 5.8757 3 1
3 4 R1 B1 G22 4.5784 4 1
4 5 R1 B2 G21 4.6540 5 1
.. ... .. ... ... ... ... ...
67 68 R3 B5 G24 3.5655 68 1
68 69 R3 B6 G03 2.8873 69 1
69 70 R3 B6 G05 4.1972 70 1
70 71 R3 B6 G20 3.7349 71 1
71 72 R3 B6 G07 3.6096 72 1
[72 rows x 7 columns]
Model Specification¶
The model includes an intercept, a single categorical fixed effect (rep), a random intercept for gen, and a random interaction effect between rep and block.
The model is specified as:
with marginal covariance structure:
and distributional assumptions:
For the present model, which includes two random intercept components and their interaction, the covariance contribution from the random effects is expressed as:
where \(\mathbf{G}_{rep} = \mathbf{I}\) is fixed as the identity matrix for identifiability.
Import modules¶
We begin by importing the required modules.
import torch
import pandas as pd
from torch_openreml import MarginalREML, blup
from torch_openreml.covariance import DummyMatrix, IdentityMatrix, ScalarMatrix, Sum, CovariancePropagation, KroneckerProduct, Augment, BlockDiagonal
from torch_openreml.example_data import john_alpha
Covariance Builder¶
Next, we construct \(\mathbf{y}\), \(\mathbf{X}\), and the components required to define \(\mathbf{V}\). Both \(\mathbf{y}\) and \(\mathbf{X}\) are represented as torch tensors. The DummyMatrix class serves as a matrix builder: it constructs the dummy matrix upon evaluation and accepts either pandas.Series or lists of strings as input. The argument drop_first=True removes the first column of the dummy matrix to avoid redundancy, as an intercept term is already included.
The classes ScalarMatrix and IdentityMatrix are also matrix builders, parameterized by the required matrix dimension. The dimension may be given as an integer, or directly as the grouping factor itself, in which case the number of distinct values defines the dimension.
We then assemble the covariance structure using composable operators. The Augment operator binds its operands side by side, forming the joint design matrix \(\mathbf{Z} = [\mathbf{Z}_{gen}\ \mathbf{Z}_{rep:block}]\), and BlockDiagonal places \(\mathbf{G}_{gen}\) and the Kronecker product \(\mathbf{G}_{rep} \otimes \mathbf{G}_{block}\) on the diagonal of the joint \(\mathbf{G}\). The CovariancePropagation operator then represents the transformation \(\mathbf{Z}\mathbf{G}\mathbf{Z}^\top\). The KroneckerProduct operator computes the direct (Kronecker) product of two matrices, and Sum adds the residual covariance \(\mathbf{R}\).
Altogether, the covariance structure can be written as:
with \(\mathbf{G}_{rep} = \mathbf{I}\), \(\mathbf{G}_{gen} = \sigma^2_{gen}\mathbf{I}\), \(\mathbf{G}_{block} = \sigma^2_{block}\mathbf{I}\), and \(\mathbf{R} = \sigma^2_{\varepsilon}\mathbf{I}\).
The model parameters are defined as:
The logarithmic parameterization ensures that the variance components remain positive during optimization.
rep = john_alpha["rep"]
block = john_alpha["block"]
gen = john_alpha["gen"]
n = len(john_alpha)
y = torch.tensor(john_alpha["yield"].values)
X = Augment(torch.ones(n, 1), DummyMatrix(rep, drop_first=True))()
V = Sum(
random = CovariancePropagation(
Z = Augment(
Z_gen = DummyMatrix(gen),
Z_rep_block = DummyMatrix(rep, block),
),
G = BlockDiagonal(
G_gen = ScalarMatrix(gen),
G_rep_block = KroneckerProduct(
G_rep = IdentityMatrix(rep),
G_block = ScalarMatrix(block),
),
),
),
residual = ScalarMatrix(n),
)
print(V)
Sum(
operands={
'random': CovariancePropagation(
operands={
'Z': Augment(
operands={
'Z_gen': DummyMatrix(shape=(72, 24)),
'Z_rep_block': DummyMatrix(shape=(72, 18))
}
),
'G': BlockDiagonal(
operands={
'G_gen': ScalarMatrix(shape=(24, 24), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}}),
'G_rep_block': KroneckerProduct(
operands={
'G_rep': IdentityMatrix(shape=(3, 3)),
'G_block': ScalarMatrix(shape=(6, 6), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}})
}
)
}
)
}
),
'residual': ScalarMatrix(shape=(72, 72), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}})
}
)
MarginalREML Optimizer¶
Once the covariance structure has been defined, it is passed to MarginalREML to initialize the estimation procedure. The optimize method is then called with \(\mathbf{y}\), \(\mathbf{X}\), and an initial value for \(\boldsymbol{\theta}\) (set to zeros in this example). The verbose argument controls the level of diagnostic output.
Because the optimization is performed on the transformed parameter scale, the estimated parameters can be mapped back to variance components using V.build_params. The resulting values correspond to the variance components associated with the parameter names stored in V.param_names.
reml = MarginalREML(V)
theta_hat, beta_hat, n_iter = reml.optimize(y, X, torch.zeros(3), verbose=1)
print(theta_hat, V.build_params(theta_hat))
print(V.free_param_names)
print(beta_hat)
[ 1] ∥∇∥: 41.8197, ∥Δ∥: 8.9438, η: 1.00, ∥Δᶜ∥: 8.9438, log 𝓛: -34.3129
[ 2] ∥∇∥: 315227.4375, ∥Δ∥: 0.8838, η: 1.00, ∥Δᶜ∥: 0.8838, log 𝓛: -201185.3906 (-201151.0777)
[ 3] ∥∇∥: 119230.2656, ∥Δ∥: 0.8519, η: 1.00, ∥Δᶜ∥: 0.8519, log 𝓛: -71844.5469 (+129340.8438)
[ 4] ∥∇∥: 44273.9336, ∥Δ∥: 0.8143, η: 1.00, ∥Δᶜ∥: 0.8143, log 𝓛: -26065.6270 (+45778.9199)
[ 5] ∥∇∥: 16331.3926, ∥Δ∥: 0.7587, η: 1.00, ∥Δᶜ∥: 0.7587, log 𝓛: -9457.4141 (+16608.2129)
[ 6] ∥∇∥: 6001.7085, ∥Δ∥: 0.7142, η: 1.00, ∥Δᶜ∥: 0.7142, log 𝓛: -3393.0168 (+6064.3972)
[ 7] ∥∇∥: 2194.9878, ∥Δ∥: 0.7031, η: 1.00, ∥Δᶜ∥: 0.7031, log 𝓛: -1181.3339 (+2211.6830)
[ 8] ∥∇∥: 793.8170, ∥Δ∥: 0.6952, η: 1.00, ∥Δᶜ∥: 0.6952, log 𝓛: -382.5103 (+798.8235)
[ 9] ∥∇∥: 278.6989, ∥Δ∥: 0.6715, η: 1.00, ∥Δᶜ∥: 0.6715, log 𝓛: -102.3058 (+280.2045)
[ 10] ∥∇∥: 90.4160, ∥Δ∥: 0.5957, η: 1.00, ∥Δᶜ∥: 0.5957, log 𝓛: -11.4295 (+90.8763)
[ 11] ∥∇∥: 23.9232, ∥Δ∥: 0.4006, η: 1.00, ∥Δᶜ∥: 0.4006, log 𝓛: 12.6788 (+24.1083)
[ 12] ∥∇∥: 3.8678, ∥Δ∥: 0.1392, η: 1.00, ∥Δᶜ∥: 0.1392, log 𝓛: 16.5908 (+3.9120)
[ 13] ∥∇∥: 0.2448, ∥Δ∥: 0.0157, η: 1.00, ∥Δᶜ∥: 0.0157, log 𝓛: 16.8077 (+0.2169)
[ 14] ∥∇∥: 0.0104, ∥Δ∥: 0.0007, η: 1.00, ∥Δᶜ∥: 0.0007, log 𝓛: 16.8099 (+0.0021)
[ 15] ∥∇∥: 0.0005, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: 16.8097 (-0.0002)
[ 16] ∥∇∥: 0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: 16.8101 (+0.0004)
[ 17] ∥∇∥: 0.0001, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: 16.8098 (-0.0002)
[ 18] ∥∇∥: 0.0001, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: 16.8101 (+0.0002)
[ 19] ∥∇∥: 0.0001, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: 16.8101 (+0.0000)
[∇: score, Δ: 𝐉⁻¹∇, η: learning rate, Δᶜ: clip(𝛉 + ηΔ, lb, ub) - 𝛉, 𝓛: restricted likelihood]
✓ Converged at iteration 19, total time: 0.02s
tensor([-0.9728, -1.3281, -1.2529]) tensor([0.1429, 0.0702, 0.0816])
['random/G/G_gen/sigma^2', 'random/G/G_rep_block/G_block/sigma^2', 'residual/sigma^2']
tensor([ 4.5183, 0.2978, -0.4140])
BLUP¶
Once the variance components have been estimated, the random effects can be predicted using the best linear unbiased predictor (BLUP):
Here, \(\mathbf{Z}\) and \(\mathbf{G}\) are the joint design and covariance matrices defined above, represented by the Augment and BlockDiagonal operands nested within the CovariancePropagation operator. These matrices can be retrieved directly from the covariance structure rather than reconstructed manually. Calling V.tree at theta_hat evaluates the structure and returns the evaluated nodes, including the root "/" for \(\mathbf{V}\) and "random/Z" and "random/G" for the random-effects design and covariance matrices.
tree, _ = V.tree(theta_hat)
The blup function takes \(\mathbf{y}\), \(\mathbf{X}\), \(\mathbf{Z}\), \(\mathbf{G}\), and \(\mathbf{V}\) as inputs and returns a prediction for each random-effect level. In this example, the predictions correspond to the 24 genotypes followed by the 18 replicate-by-block combinations:
b_hat = blup(y, X, tree["random/Z"], tree["random/G"], tree["/"])
b_hat
tensor([ 0.5012, 0.0050, -0.7846, 0.0061, 0.4750, 0.0446, -0.3089, 0.0622,
-0.8099, -0.0894, -0.1964, 0.2258, 0.2317, 0.2434, 0.4247, 0.2010,
0.0781, -0.1102, 0.2896, -0.3390, 0.2561, 0.0241, -0.1770, -0.2531,
0.1231, -0.1412, -0.1504, -0.1068, 0.0737, 0.2015, -0.5326, -0.3012,
0.2432, 0.1349, 0.2753, 0.1804, 0.0506, -0.0478, 0.1511, 0.0538,
-0.0080, -0.1996])
BLUPs for an individual random-effect component can be obtained in the same way by extracting its corresponding design and covariance matrices. For example, the random intercepts for gen are defined by "random/Z/Z_gen" and "random/G/G_gen":
gen_blup = blup(y, X, tree["random/Z/Z_gen"], tree["random/G/G_gen"], tree["/"])
pd.Series(gen_blup, index=sorted(set(gen)))
G01 0.501184
G02 0.004963
G03 -0.784562
G04 0.006126
G05 0.474950
G06 0.044641
G07 -0.308947
G08 0.062230
G09 -0.809931
G10 -0.089373
G11 -0.196434
G12 0.225759
G13 0.231665
G14 0.243400
G15 0.424700
G16 0.200965
G17 0.078078
G18 -0.110180
G19 0.289576
G20 -0.338969
G21 0.256132
G22 0.024089
G23 -0.176997
G24 -0.253057
dtype: float32
Documentation¶
Section |
Description |
|---|---|
Vignettes |
|
Model formulation, REML and ML theory, score and AI matrix derivations. |
|
Full documentation for |
Citing¶
@software{torch_openreml,
author = {Weihao Li},
title = {torch-openreml},
year = {2026},
url = {https://github.com/anu-aagi/torch-openreml/}
}