For R Users

torch-openreml can be used in R through the reticulate package. We recommend installing torch-openreml inside a dedicated conda environment.

Installation

1. Install Conda

You can skip this step if conda is already installed on your system.

if (is.null(reticulate:::find_conda()[[1]])) {
    reticulate::install_miniconda()
}

2. Create a Conda Environment

if (reticulate::condaenv_exists("torch-openreml")) {
    reticulate::conda_remove("torch-openreml")
}

reticulate::conda_create("torch-openreml",
                         python_version = "3.12.12")

3. Install torch-openreml

reticulate::conda_install("torch-openreml",
                          pip = TRUE,
                          packages = c("git+https://github.com/anu-aagi/torch-openreml.git"))

Usage

The following example mirrors the Getting Started example. In R, it is often more convenient to construct design matrices directly using model.matrix(). As in Python, a matrix dimension may be given as the grouping factor itself, whose distinct values define the dimension.

library(reticulate)
use_condaenv("torch-openreml")

openreml <- import("torch_openreml", convert = FALSE)
torch <- import("torch", convert = FALSE)

BlockDiagonal <- openreml$covariance$BlockDiagonal
ScalarMatrix <- openreml$covariance$ScalarMatrix
IdentityMatrix <- openreml$covariance$IdentityMatrix
KroneckerProduct <- openreml$covariance$KroneckerProduct
CovariancePropagation <- openreml$covariance$CovariancePropagation
Sum <- openreml$covariance$Sum
MarginalREML <- openreml$MarginalREML

data <- agridat::john.alpha
n <- nrow(data)

y <- torch$tensor(data$yield, dtype = torch$float32)
X <- model.matrix(~ rep, data = data) |>
    torch$tensor(dtype = torch$float32)

Z <- model.matrix(~ 0 + gen + block:rep, data = data) |>
    torch$tensor(dtype = torch$float32)

V <- Sum(
    random = CovariancePropagation(
        Z = Z,
        G = BlockDiagonal(
            G_gen = ScalarMatrix(data$gen),
            G_rep_block = KroneckerProduct(
                G_rep = IdentityMatrix(data$rep),
                G_block = ScalarMatrix(data$block)
            )
        )
    ),
    residual = ScalarMatrix(n)
)

fit_openreml <- MarginalREML(V)
result <- fit_openreml$optimize(y, X, torch$zeros(3L), verbose = 1L)

print(py_to_r(fit_openreml$get_theta()$numpy()))
print(py_to_r(V$build_params(fit_openreml$get_theta())$numpy()))
print(py_to_r(fit_openreml$get_beta()$numpy()))

Random effects are predicted with the same best linear unbiased predictor (BLUP) as in the Python example. The joint design and covariance matrices are read back from the fitted covariance structure with V.tree, which evaluates the structure at the estimated parameters and returns each node under its path: "random/Z" and "random/G" hold \(\mathbf{Z}\) and \(\mathbf{G}\), and the root "/" holds \(\mathbf{V}\). Note that reticulate indexes Python objects from zero, so [[0]] selects the results from the (results, free_params_by_path) pair that tree() returns.

tree <- V$tree(fit_openreml$get_theta())[[0]]

b_hat <- openreml$blup(y, X, tree[["random/Z"]], tree[["random/G"]], tree[["/"]])
print(py_to_r(b_hat$numpy()))

The BLUPs of a single random-effect component follow the same pattern. Here the design matrix is a plain tensor built by model.matrix(), so the columns of the gen term are built on their own, while its covariance is the "random/G/G_gen" node of the tree.

Z_gen <- model.matrix(~ 0 + gen, data = data) |>
    torch$tensor(dtype = torch$float32)
gen_blup <- openreml$blup(y, X, Z_gen, tree[["random/G/G_gen"]], tree[["/"]])
print(setNames(py_to_r(gen_blup$numpy()), levels(data$gen)))

[  1] ∥∇∥:      41.8197, ∥Δ∥: 8.9438, η: 1.00, ∥Δᶜ∥: 8.9438, log 𝓛: -34.3129
[  2] ∥∇∥:  315227.5000, ∥Δ∥: 0.8838, η: 1.00, ∥Δᶜ∥: 0.8838, log 𝓛: -201185.4531 (-201151.1402)
[  3] ∥∇∥:  119231.8828, ∥Δ∥: 0.8520, η: 1.00, ∥Δᶜ∥: 0.8520, log 𝓛: -71845.0547 (+129340.3984)
[  4] ∥∇∥:   44274.7461, ∥Δ∥: 0.8143, η: 1.00, ∥Δᶜ∥: 0.8143, log 𝓛: -26065.9766 (+45779.0781)
[  5] ∥∇∥:   16330.8594, ∥Δ∥: 0.7587, η: 1.00, ∥Δᶜ∥: 0.7587, log 𝓛: -9457.3301 (+16608.6465)
[  6] ∥∇∥:    6001.5020, ∥Δ∥: 0.7142, η: 1.00, ∥Δᶜ∥: 0.7142, log 𝓛: -3392.9321 (+6064.3979)
[  7] ∥∇∥:    2194.9365, ∥Δ∥: 0.7031, η: 1.00, ∥Δᶜ∥: 0.7031, log 𝓛: -1181.3032 (+2211.6289)
[  8] ∥∇∥:     793.7987, ∥Δ∥: 0.6952, η: 1.00, ∥Δᶜ∥: 0.6952, log 𝓛: -382.4990 (+798.8043)
[  9] ∥∇∥:     278.6925, ∥Δ∥: 0.6715, η: 1.00, ∥Δᶜ∥: 0.6715, log 𝓛: -102.3015 (+280.1974)
[ 10] ∥∇∥:      90.4133, ∥Δ∥: 0.5957, η: 1.00, ∥Δᶜ∥: 0.5957, log 𝓛: -11.4285 (+90.8730)
[ 11] ∥∇∥:      23.9223, ∥Δ∥: 0.4006, η: 1.00, ∥Δᶜ∥: 0.4006, log 𝓛:  12.6798 (+24.1083)
[ 12] ∥∇∥:       3.8675, ∥Δ∥: 0.1392, η: 1.00, ∥Δᶜ∥: 0.1392, log 𝓛:  16.5909 (+3.9112)
[ 13] ∥∇∥:       0.2448, ∥Δ∥: 0.0157, η: 1.00, ∥Δᶜ∥: 0.0157, log 𝓛:  16.8083 (+0.2174)
[ 14] ∥∇∥:       0.0104, ∥Δ∥: 0.0007, η: 1.00, ∥Δᶜ∥: 0.0007, log 𝓛:  16.8100 (+0.0017)
[ 15] ∥∇∥:       0.0004, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8095 (-0.0004)
[ 16] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8101 (+0.0005)
[ 17] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8102 (+0.0001)
[ 18] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8099 (-0.0003)
[ 19] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8102 (+0.0003)
[ 20] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8100 (-0.0002)
[ 21] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8104 (+0.0004)
[ 22] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8103 (-0.0002)
[ 23] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8097 (-0.0006)
[ 24] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛:  16.8097 (-0.0000)

[∇: score, Δ: 𝐉⁻¹∇, η: learning rate, Δᶜ: clip(𝛉 + ηΔ, lb, ub) - 𝛉, 𝓛: restricted likelihood]

✓ Converged at iteration 24, total time: 0.02s

[1] -0.9727983 -1.3280730 -1.2528642
[1] 0.14290193 0.07021832 0.08161612
[1]  4.5182486  0.2978466 -0.4140441
 [1]  0.501183033  0.004962160 -0.784563184  0.006125156  0.474949241
 [6]  0.044639699 -0.308948159  0.062228836 -0.809932053 -0.089373551
[11] -0.196435124  0.225757748  0.231664374  0.243399158  0.424699038
[16]  0.200964108  0.078077391 -0.110181324  0.289575398 -0.338969529
[21]  0.256131411  0.024088217 -0.176998153 -0.253058344  0.123135120
[26] -0.141226426 -0.150395423 -0.106756583  0.073703296  0.201533794
[31] -0.532640934 -0.301233232  0.243239060  0.134878263  0.275336683
[36]  0.180418730  0.050569884 -0.047783885  0.151079014  0.053760849
[41] -0.008047517 -0.199577749
         G01          G02          G03          G04          G05          G06 
 0.501183033  0.004962160 -0.784563184  0.006125156  0.474949241  0.044639699 
         G07          G08          G09          G10          G11          G12 
-0.308948159  0.062228836 -0.809932053 -0.089373551 -0.196435124  0.225757748 
         G13          G14          G15          G16          G17          G18 
 0.231664374  0.243399158  0.424699038  0.200964108  0.078077391 -0.110181324 
         G19          G20          G21          G22          G23          G24 
 0.289575398 -0.338969529  0.256131411  0.024088217 -0.176998153 -0.253058344