torch_openreml.covariance.UnconstrainedMatrix

class torch_openreml.covariance.UnconstrainedMatrix(n, param_specs=None)[source]

Bases: Matrix

Full symmetric covariance matrix parameterised by lower-triangular entries.

The matrix is built from the lower triangle (including the diagonal), then mirrored to the upper triangle to ensure symmetry:

\[\begin{split}\symbf{V}_{ij} = \begin{cases} \theta_{ij} & i \ge j \\ \theta_{ji} & i < j \end{cases}\end{split}\]

Diagonal entries (\(i = j\)) are transformed to positive values via TransformExpPow2 by default. Off-diagonal entries (\(i > j\)) are unconstrained and use TransformIdentity.

Note

This parameterisation ensures symmetry but does not guarantee positive definiteness. For a positive-definite covariance matrix, consider a Cholesky-based parameterisation.

Initialize an unconstrained symmetric matrix of size n x n.

By default, the matrix has \(n(n+1)/2\) free parameters: one for each lower-triangular entry. Diagonal entries are parameterised on the positive real line (via TransformExpPow2); off-diagonal entries are unconstrained (via TransformIdentity).

Parameters:
  • 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 (Transform), respectively.

Example:

import torch
from torch_openreml.covariance import UnconstrainedMatrix

mat = UnconstrainedMatrix(3)
mat
UnconstrainedMatrix(shape=(3, 3), param_specs={'sigma^2_0_0': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}, ..., 'sigma^2_2_2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}})
free_params = torch.tensor([0.0, 0.5, 1.0, 0.2, -0.3, 0.4])
mat(free_params)
tensor([[ 1.0000,  0.5000,  0.2000],
        [ 0.5000,  7.3891, -0.3000],
        [ 0.2000, -0.3000,  2.2255]])
mat.grad(free_params)
(tensor([[[ 2.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000]],
 
         [[ 0.0000,  1.0000,  0.0000],
          [ 1.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000]],
 
         [[ 0.0000,  0.0000,  0.0000],
          [ 0.0000, 14.7781,  0.0000],
          [ 0.0000,  0.0000,  0.0000]],
 
         [[ 0.0000,  0.0000,  1.0000],
          [ 0.0000,  0.0000,  0.0000],
          [ 1.0000,  0.0000,  0.0000]],
 
         [[ 0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  1.0000],
          [ 0.0000,  1.0000,  0.0000]],
 
         [[ 0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  4.4511]]]),
 ['sigma^2_0_0',
  'sigma^2_1_0',
  'sigma^2_1_1',
  'sigma^2_2_0',
  'sigma^2_2_1',
  'sigma^2_2_2'])

Methods

__call__([free_params])

Construct the matrix from a flat parameter tensor.

auto_grad([free_params])

Compute the Jacobian of build() with respect to free parameters using automatic differentiation.

build_params([free_params, include_fixed, ...])

Construct the full parameter tensor from free parameters.

disable_cache()

Disable caching of intermediate computation results.

enable_cache()

Enable caching of intermediate computation results.

get_default_dtype_device()

Get the dtype and the device of the parameter defaults.

get_intermediates(params)

Retrieve cached intermediate computation results if still valid.

grad([free_params])

Compute the Jacobian of __call__() with respect to trainable parameters.

inverse([free_params])

Construct the matrix and return its inverse.

manual_grad([free_params])

Compute the Jacobian of __call__() with respect to trainable parameters using a closed-form analytic expression.

reset_intermediates()

Clear the intermediate computation cache.

resolve_dim(n)

Resolve a matrix dimension from an integer or a vector of levels.

set_intermediates(params, intermediates)

Cache intermediate computation results keyed by the parameter values.

set_param_specs(key, *[, fixed, default, trans])

Update the fields of one or more parameter specifications in place.

trans_grad([free_params])

Compute the element-wise derivative of the free parameter transforms.

Attributes

fixed_param_defaults

Fixed parameter defaults.

fixed_param_index

Index of fixed parameters.

fixed_param_names

Fixed parameter names.

fixed_param_trans

Transforms for fixed parameters.

free_param_defaults

Free parameter defaults.

free_param_index

Index of free parameters.

free_param_names

Free parameter names.

free_param_trans

Transforms for free parameters.

num_fixed_params

Total number of fixed parameters.

num_free_params

Total number of free parameters.

num_params

Total number of parameters.

param_defaults

Parameter defaults.

param_names

Parameter names.

param_specs

Parameter specifications, returned as stored rather than copied.

param_trans

Parameter transforms.

repr_dict

Key-value pairs used to build the string representation.

shape

Output matrix shape.

grad_mode

Gradient computation mode.

jacobian_method

Jacobian method used by auto_grad(), one of "jacrev", "jacfwd", or "jacobian".

jacobian_chunk_size

Number of outputs held per batch when auto_grad() differentiates with "jacrev", or None to compute the Jacobian in a single batch.

cache

Whether intermediate computation results are cached, toggled by enable_cache() and disable_cache().

__call__(free_params=None)[source]

Construct the matrix from a flat parameter tensor.

Must be implemented by subclasses. Implementations should convert free_params via build_params() to validate, include fixed parameters, and apply transforms before any computation.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.

Returns:

Constructed matrix of shape shape.

Return type:

torch.Tensor

manual_grad(free_params=None)[source]

Compute the Jacobian of __call__() with respect to trainable parameters using a closed-form analytic expression.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.

Returns:

(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 build_params().

Return type:

tuple