torch_openreml.MarginalREML

class torch_openreml.MarginalREML(v)[source]

Bases: object

Marginal REML estimator for generalised least squares with a parametric marginal covariance matrix.

Fits covariance parameters \(\boldsymbol{\theta}\) by maximising the restricted log-likelihood

\[\ell_R(\boldsymbol{\theta}) = -\frac{1}{2} \left( \log |\symbf{V}(\boldsymbol{\theta})| + \log |\symbf{X}^\top \symbf{V}(\boldsymbol{\theta})^{-1} \symbf{X}| + \symbf{y}^\top \symbf{P} \symbf{y} \right)\]

where \(\symbf{P} = \symbf{V}(\boldsymbol{\theta})^{-1} - \symbf{V}(\boldsymbol{\theta})^{-1}\symbf{X} (\symbf{X}^\top\symbf{V}(\boldsymbol{\theta})^{-1}\symbf{X})^{-1}\symbf{X}^\top\symbf{V}(\boldsymbol{\theta})^{-1}\) is the projection matrix and \(\symbf{V}(\boldsymbol{\theta})\) is the marginal covariance matrix of \(\symbf{y}\).

Optimisation uses the average information (AI) algorithm, which forms a quasi-Newton step \(\Delta = \symbf{AI}^{-1} \symbf{s}\) from the score vector \(\symbf{s}\) and the AI matrix at each iteration.

The covariance model \(\symbf{V}(\boldsymbol{\theta})\) is supplied as a Matrix instance via v. Gradients are handled internally by the matrix.

Initialize a MarginalREML estimator.

Parameters:

v (Matrix) – A Matrix instance that constructs \(\symbf{V}(\boldsymbol{\theta})\) and its Jacobian.

Raises:

TypeError – If v is not a Matrix instance.

Example:

import torch
from torch_openreml import MarginalREML
from torch_openreml.covariance import ScalarMatrix

n, p = 50, 2
y = torch.randn(n)
x = torch.randn(n, p)
theta = torch.tensor([0.0])

mat = ScalarMatrix(n)
reml = MarginalREML(mat)
theta_hat, beta_hat, n_iter = reml.optimize(y, x, theta, verbose=1)
theta_hat, beta_hat

[  1] ∥∇∥:       9.9403, ∥Δ∥: 0.1306, η: 1.00, ∥Δᶜ∥: 0.1306, log 𝓛: -22.9222
[  2] ∥∇∥:       1.4189, ∥Δ∥: 0.0144, η: 1.00, ∥Δᶜ∥: 0.0144, log 𝓛: -22.3336 (+0.5887)
[  3] ∥∇∥:       0.0202, ∥Δ∥: 0.0002, η: 1.00, ∥Δᶜ∥: 0.0002, log 𝓛: -22.3233 (+0.0103)
[  4] ∥∇∥:       0.0000, ∥Δ∥: 0.0000, η: 1.00, ∥Δᶜ∥: 0.0000, log 𝓛: -22.3233 (+0.0000)

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

✓ Converged at iteration 4, total time: 0.00s


(tensor([-0.1160]), tensor([-0.0768, -0.1107]))

Methods

ai_step(y, x, theta[, require_loglik, ...])

Perform a single average information (AI) algorithm step.

blue(y, x, theta)

Compute the best linear unbiased estimator (BLUE) of \(\boldsymbol{\beta}\).

compute_v_dv(theta)

Compute the covariance matrix and its Jacobian.

get_beta([select])

Retrieve a coefficient estimate from the optimisation history.

get_theta([select])

Retrieve a covariance parameter estimate from the optimisation history.

is_converged()

Check whether the optimisation has converged.

loglik(y, x, theta)

Evaluate the REML log-likelihood.

optimize(y, x[, theta, max_iter, eta, ...])

Run the AI-REML optimisation loop.

predict(y, x, theta)

Compute fitted values.

residual(y, x, theta)

Compute residuals.

update(theta, delta, eta[, lb, ub])

Apply a damped AI step and clip covariance parameters to bounds.
blue(y, x, theta)[source]

Compute the best linear unbiased estimator (BLUE) of \(\boldsymbol{\beta}\).

Solves for \(\hat{\boldsymbol{\beta}}\) via generalised least squares:

\[\hat{\boldsymbol{\beta}} = (\symbf{X}^\top \symbf{V}(\boldsymbol{\theta})^{-1} \symbf{X})^{-1} \symbf{X}^\top \symbf{V}(\boldsymbol{\theta})^{-1} \symbf{y}\]
Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor) – Flat covariance parameter tensor.

Returns:

Coefficient estimate \(\hat{\boldsymbol{\beta}}\) of shape (p,).

Return type:

torch.Tensor

predict(y, x, theta)[source]

Compute fitted values.

\[\hat{\symbf{y}} = \symbf{X} \hat{\boldsymbol{\beta}}\]
Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor) – Flat covariance parameter tensor.

Returns:

Fitted values of shape (n,).

Return type:

torch.Tensor

residual(y, x, theta)[source]

Compute residuals.

\[\hat{\symbf{e}} = \symbf{y} - \symbf{X}\hat{\boldsymbol{\beta}}\]
Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor) – Flat covariance parameter tensor.

Returns:

Residuals of shape (n,).

Return type:

torch.Tensor

loglik(y, x, theta)[source]

Evaluate the REML log-likelihood.

\[\ell_R(\boldsymbol{\theta}) = -\frac{1}{2} \left( \log |\symbf{V}(\boldsymbol{\theta})| + \log |\symbf{X}^\top \symbf{V}(\boldsymbol{\theta})^{-1} \symbf{X}| + \symbf{y}^\top \symbf{P} \symbf{y} \right)\]
Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor) – Flat covariance parameter tensor.

Returns:

Scalar REML log-likelihood value.

Return type:

torch.Tensor

compute_v_dv(theta)[source]

Compute the covariance matrix and its Jacobian.

Calls v to build \(\symbf{V}(\boldsymbol{\theta})\) and grad() for the Jacobian \(\partial\symbf{V}(\boldsymbol{\theta})/\partial\boldsymbol{\theta}\).

Parameters:

theta (torch.Tensor) – Flat covariance parameter tensor.

Returns:

(v, dv), where v is the covariance matrix of shape (n, n) and dv is the Jacobian of shape (num_free_params, n, n).

Return type:

tuple

ai_step(y, x, theta, require_loglik=True, require_beta=True, trace_approx=False, subspace_fraction=0.01)[source]

Perform a single average information (AI) algorithm step.

Computes the score vector \(\symbf{s}\), AI matrix, and optionally the REML log-likelihood and coefficient estimate at the current \(\boldsymbol{\theta}\).

The score vector and AI matrix are:

\[\begin{split}s_k &= \frac{1}{2}\left( \symbf{y}^\top \symbf{P} \frac{\partial\symbf{V}(\boldsymbol{\theta})}{\partial\theta_k} \symbf{P} \symbf{y} - \mathrm{tr}\!\left(\symbf{P} \frac{\partial\symbf{V}(\boldsymbol{\theta})}{\partial\theta_k}\right) \right) \\ \mathrm{AI}_{kj} &= \frac{1}{2} \symbf{y}^\top \symbf{P} \frac{\partial\symbf{V}(\boldsymbol{\theta})}{\partial\theta_k} \symbf{P} \frac{\partial\symbf{V}(\boldsymbol{\theta})}{\partial\theta_j} \symbf{P} \symbf{y}\end{split}\]

When trace_approx is True, the trace term \(\mathrm{tr}(\symbf{P}\,\partial\symbf{V}/\partial\theta_k)\) is replaced by a stochastic probe estimate, which avoids forming \(\symbf{V}^{-1}\partial\symbf{V}\) in full. With

\[\symbf{A}_k = \symbf{V}^{-1} \frac{\partial \symbf{V}}{\partial \theta_k}, \qquad \symbf{B}_k = \symbf{V}^{-1} \symbf{X} (\symbf{X}^\top \symbf{V}^{-1} \symbf{X})^{-1} \symbf{X}^\top \symbf{V}^{-1} \frac{\partial \symbf{V}}{\partial \theta_k},\]

the trace splits as

\[\mathrm{tr}\!\left(\symbf{P} \frac{\partial \symbf{V}}{\partial \theta_k}\right) = \mathrm{tr}(\symbf{A}_k) - \mathrm{tr}(\symbf{B}_k),\]

and each of the two terms is estimated on its own, so that neither \(\symbf{A}_k\) nor \(\symbf{B}_k\) is ever formed in full.

The estimator is Hutch++ (Meyer, Musco, Musco & Woodruff, Hutch++: Optimal stochastic trace estimation, SOSA 2021, pp. 142–155, doi:10.1137/1.9781611976496.16). For a matrix \(\symbf{C}\), two independent Rademacher probe matrices \(\symbf{S}, \symbf{G} \in \{\pm 1\}^{n \times m}\) are drawn once per call and shared across parameters, with \(m = \lfloor n f \rfloor + 1\) probes each and \(f\) the subspace_fraction, and

\[\begin{split}\symbf{Q} &= \mathrm{qr}(\symbf{C} \symbf{S}), \\ \hat{h}(\symbf{C}) &= \mathrm{tr}(\symbf{Q}^\top \symbf{C} \symbf{Q}) + \frac{1}{m} \mathrm{tr}\!\left( \symbf{G}^\top (\symbf{I} - \symbf{Q}\symbf{Q}^\top) \symbf{C} (\symbf{I} - \symbf{Q}\symbf{Q}^\top) \symbf{G} \right).\end{split}\]

\(\symbf{Q}\) spans the dominant range of \(\symbf{C}\), so the first term is exact there; the second applies Hutchinson’s estimator to the deflated residual.

When trace_approx is False the traces are evaluated exactly. The average information matrix and the log-likelihood are computed exactly either way; only the trace term of the score is approximated.

Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor) – Flat covariance parameter tensor.

  • require_loglik (bool, optional) – Whether to evaluate the REML log-likelihood. Defaults to True.

  • require_beta (bool, optional) – Whether to compute the coefficient estimate. Defaults to True.

  • trace_approx (bool, optional) – Whether to estimate the trace term of the score stochastically. False computes the exact score, which is more expensive per call. This is a plain switch with no internal state: optimize() owns the decision of when to turn it off. Defaults to False.

  • subspace_fraction (float, optional) – Fraction of the n observations used as the dimension of the random probe subspace in the stochastic trace estimate. Each of the two probe matrices holds int(n * subspace_fraction) + 1 vectors, so the total probe budget is about twice that. Must lie in [0, 1]; larger values trade time for accuracy of the approximate score. Only used when trace_approx is True. Defaults to 0.01.

Raises:

ValueError – If subspace_fraction is outside [0, 1].

Returns:

(beta, score, ai, loglik), where beta is of shape (p,), score is of shape (num_free_params,), ai is of shape (num_free_params, num_free_params), and loglik is a scalar tensor. beta and loglik are torch.nan if their respective require_* flag is False.

Return type:

tuple

get_theta(select='last')[source]

Retrieve a covariance parameter estimate from the optimisation history.

Parameters:

select (str, optional) – "last" returns the final iterate; any other value returns the iterate with the highest log-likelihood. Defaults to "last".

Returns:

Selected covariance parameter tensor \(\boldsymbol{\theta}\).

Return type:

torch.Tensor

get_beta(select='last')[source]

Retrieve a coefficient estimate from the optimisation history.

Parameters:

select (str, optional) – "last" returns the final iterate; any other value returns the iterate with the highest log-likelihood. Defaults to "last".

Returns:

Selected coefficient estimate \(\hat{\boldsymbol{\beta}}\).

Return type:

torch.Tensor

is_converged()[source]

Check whether the optimisation has converged.

Convergence is declared when all enabled criteria fall below their respective tolerances. At least two iterations must have completed before any criterion can be satisfied.

Returns:

True if all enabled criteria are satisfied, False otherwise.

Return type:

bool

update(theta, delta, eta, lb=-inf, ub=inf)[source]

Apply a damped AI step and clip covariance parameters to bounds.

\[\boldsymbol{\theta} \leftarrow \mathrm{clip}(\boldsymbol{\theta} + \eta \Delta,\, \text{lb},\, \text{ub})\]
Parameters:

  • theta (torch.Tensor) – Current covariance parameter tensor.

  • delta (torch.Tensor) – AI step \(\Delta = \symbf{AI}^{-1}\symbf{s}\).

  • eta (float) – Step size (learning rate).

  • lb (float, optional) – Lower bound for clipping. Defaults to -inf.

  • ub (float, optional) – Upper bound for clipping. Defaults to inf.

Returns:

(theta, update), where theta is the updated covariance parameter tensor and update is the actual change after clipping.

Return type:

tuple

optimize(y, x, theta=None, max_iter=200, eta=1.0, require_loglik=True, lb=-inf, ub=inf, verbose=0, check_score=True, check_delta=True, check_loglik=True, tol_score=0.0001, tol_delta=0.0001, tol_loglik=0.0001, trace_approx=False, subspace_fraction=0.01, exact_score_threshold=10.0)[source]

Run the AI-REML optimisation loop.

Iterates ai_step() and update() until convergence or max_iter is reached. Optimisation history is stored in history and can be queried afterwards via get_theta() and get_beta().

Parameters:

  • y (torch.Tensor) – Response vector of shape (n,).

  • x (torch.Tensor) – Design matrix of shape (n, p).

  • theta (torch.Tensor or dict) – Initial covariance parameter tensor, or a dictionary mapping free parameter names to value tensors. If omitted, default values are used. Defaults to None.

  • max_iter (int, optional) – Maximum number of iterations. Defaults to 200.

  • eta (float, optional) – Step size applied to each AI update. Defaults to 1.0.

  • require_loglik (bool, optional) – Whether to evaluate the REML log-likelihood at each iteration. Defaults to True.

  • lb (float, optional) – Lower bound for covariance parameter clipping. Defaults to -inf.

  • ub (float, optional) – Upper bound for covariance parameter clipping. Defaults to inf.

  • verbose (int, optional) – Verbosity level. 0 suppresses all output; 1 prints per-iteration diagnostics, whether the optimisation converged or hit max_iter, and the total optimisation time. Defaults to 0.

  • check_score (bool, optional) – Include score norm in convergence check. Defaults to True.

  • check_delta (bool, optional) – Include covariance parameter update norm in convergence check. Defaults to True.

  • check_loglik (bool, optional) – Include log-likelihood change in convergence check. Defaults to True.

  • tol_score (float, optional) – Score norm tolerance. Defaults to 1e-4.

  • tol_delta (float, optional) – Covariance parameter update norm tolerance. Defaults to 1e-4.

  • tol_loglik (float, optional) – Log-likelihood change tolerance. Defaults to 1e-4.

  • trace_approx (bool, optional) –

    Whether to start the optimisation with the stochastic probe estimate of ai_step() for the score, which follows Hutch++ (Meyer, Musco, Musco & Woodruff, Hutch++: Optimal stochastic trace estimation, SOSA 2021, pp. 142–155, doi:10.1137/1.9781611976496.16). While it is on, the optimiser monitors the score norm, and once that norm drops below exact_score_threshold the exact score is used for that and every remaining iteration. The probe estimate has a noise floor that does not decay, so an optimisation left on the approximation cannot satisfy tol_score and will run to max_iter; the switch is what makes the approximation usable. Defaults to False, which runs the whole optimisation with the exact score.

  • subspace_fraction (float, optional) – Fraction of the n observations used as the dimension of the random probe subspace in the stochastic trace estimate of ai_step(). Must lie in [0, 1]; larger values trade time for accuracy of the score on iterations that are still far from the optimum. Defaults to 0.01.

  • exact_score_threshold (float, optional) – Score norm below which the stochastic probe estimate is switched off for the remainder of the optimisation. Since the first score is always computed with the current setting, inf switches off after the first iteration while 0.0 keeps the approximation on throughout. Defaults to 10.0.

Returns:

(theta, beta, n_iter), where theta is the final covariance parameter estimate, beta is the corresponding coefficient estimate, and n_iter is the number of iterations completed.

Return type:

tuple