"""
Diagonal covariance matrix.
This module provides a diagonal covariance matrix with one variance
parameter per diagonal entry, for use in linear mixed-effects models.
Classes:
DiagonalMatrix:
A diagonal covariance matrix :math:`V = \\mathrm{diag}(\\sigma^2)`.
"""
from torch_openreml.covariance.matrix import Matrix
from torch_openreml.covariance.transform import TransformExpPow2
import torch
[docs]
class DiagonalMatrix(Matrix):
r"""
Diagonal covariance matrix with one variance parameter per entry.
.. math::
\symbf{V} = \mathrm{diag}(\sigma^2_0, \ldots, \sigma^2_{n-1})
Each diagonal entry is parameterised by a single unconstrained scalar
transformed to a positive variance via :class:`~torch_openreml.covariance.transform.TransformExpPow2`
by default. Off-diagonal entries are always zero.
"""
def __init__(self, n, param_specs=None):
"""
Initialize a diagonal covariance matrix of size ``n x n``.
Args:
n (int, list, tuple or pandas.Series): Matrix dimension, or a
vector of levels whose distinct elements define it.
param_specs (dict): Parameter specifications. Keys should be strings
representing parameter names. Values should be dictionaries
containing the specification for each parameter. Each specification
dictionary should contain the keys ``"fixed"``, ``"default"``, and ``"trans"``,
representing whether the parameter is fixed or free (bool), the
default value (1D torch.Tensor), and the transform (:class:`~torch_openreml.covariance.transform.Transform`),
respectively.
Example:
.. jupyter-execute::
import torch
from torch_openreml.covariance import DiagonalMatrix
mat = DiagonalMatrix(3)
mat
.. jupyter-execute::
free_params = torch.tensor([0.0, 0.5, 1.0])
mat(free_params)
.. jupyter-execute::
mat.grad(free_params)
"""
n = self.resolve_dim(n)
param_specs = param_specs or {
f"sigma^2_{i}": {
"fixed": False,
"default": torch.tensor([0.0]),
"trans": TransformExpPow2()
} for i in range(n)
}
super().__init__((n, n), param_specs)
[docs]
def __call__(self, free_params=None):
sigma2 = self.build_params(free_params)
return torch.diag(sigma2)
[docs]
def inverse(self, free_params=None):
"""
Return the inverse of the matrix from its closed form.
The matrix is
:math:`\\symbf{V} = \\mathrm{diag}(\\sigma^2_0, \\ldots, \\sigma^2_{n-1})`,
so its inverse is
:math:`\\symbf{V}^{-1} = \\mathrm{diag}(1 / \\sigma^2_0, \\ldots, 1 / \\sigma^2_{n-1})`.
Args:
free_params (torch.Tensor or dict): Flat 1D parameter tensor or
parameter dictionary. If omitted, default values are used.
Default: ``None``.
Returns:
torch.Tensor: The ``n x n`` inverse of the constructed matrix.
Example:
.. jupyter-execute::
import torch
from torch_openreml.covariance import DiagonalMatrix
mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
mat.inverse(free_params)
"""
sigma2 = self.build_params(free_params)
return torch.diag(1 / sigma2)
[docs]
def manual_grad(self, free_params=None):
"""
Compute the Jacobian of :meth:`__call__` with respect to trainable
parameters using a closed-form analytic expression.
Args:
free_params (torch.Tensor or dict): Flat 1D parameter tensor or
parameter dictionary.
If omitted, default values are used. Default: ``None``.
Returns:
tuple: ``(grad, grad_names)``, where ``grad`` is a 3D tensor of
shape ``(num_free_params, *shape)`` and
``grad_names`` is a list of the corresponding parameter names.
Returns ``(None, [])`` when the matrix has no free parameters,
after ``free_params`` is validated via
:meth:`~torch_openreml.covariance.matrix.Matrix.build_params`.
"""
free_params = self.build_params(free_params, include_fixed=False, trans=False)
if len(free_params) == 0:
return None, []
device = free_params.device
dtype = free_params.dtype
grad = torch.zeros(free_params.shape[0], self.shape[0], self.shape[0], device=device, dtype=dtype)
idx = torch.arange(self.shape[0], device=device)
mask = torch.zeros(self.shape[0], dtype=torch.bool, device=device)
mask[self.free_param_index] = True
idx = idx[mask]
grad[torch.arange(free_params.shape[0], device=device), idx, idx] = self.trans_grad(free_params)
return grad, self.free_param_names