torch_openreml.covariance.EquicorrelationMatrix¶

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

Bases: Matrix

Equicorrelation matrix with a single shared correlation parameter.

\[\symbf{V} = (1 - \rho)\symbf{I}_n + \rho\symbf{J}_n\]

where \(\symbf{I}_n\) is the identity matrix and \(\symbf{J}_n\) is the matrix of ones. All diagonal entries equal one and all off-diagonal entries equal \(\rho\).

For \(\symbf{V}\) to be positive definite, the correlation parameter must satisfy \(\rho > -1/(n-1)\). The default transform enforces this by mapping an unconstrained scalar through a sigmoid scaled to \((-1/(n-1),\, 1)\).

Unlike CompoundSymmetricMatrix, this matrix has no variance parameter.

Initialize an equicorrelation matrix of size n x n.

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 EquicorrelationMatrix

mat = EquicorrelationMatrix(3)
mat

EquicorrelationMatrix(shape=(3, 3), param_specs={'rho': {'fixed': False, 'default': tensor([0.]), 'trans': TransformChain([TransformSigmoid(), TransformScaleShift(a=1.5, b=-0.5)])}})

free_params = torch.tensor([0.0])
mat(free_params)

tensor([[1.0000, 0.2500, 0.2500],
        [0.2500, 1.0000, 0.2500],
        [0.2500, 0.2500, 1.0000]])

mat.grad(free_params)

(tensor([[[0.0000, 0.3750, 0.3750],
          [0.3750, 0.0000, 0.3750],
          [0.3750, 0.3750, 0.0000]]]),
 ['rho'])

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])	Return the inverse of the matrix from its closed form.
`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

inverse(free_params=None)[source]¶

Return the inverse of the matrix from its closed form.

The matrix is \(\symbf{V} = (1 - \rho)\symbf{I}_n + \rho\symbf{J}_n\), a rank-one update of a scaled identity, so its inverse follows from the Sherman-Morrison formula:

\[\symbf{V}^{-1} = \frac{\symbf{I}_n}{1 - \rho} - \frac{\rho\,\symbf{J}_n}{(1 - \rho)(1 + (n - 1)\rho)}\]

Parameters:: free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.
Returns:: The n x n inverse of the constructed matrix.
Return type:: torch.Tensor

Example:

import torch
from torch_openreml.covariance import EquicorrelationMatrix

mat = EquicorrelationMatrix(3)
free_params = torch.tensor([0.5])
mat.inverse(free_params)

tensor([[ 1.3557, -0.4101, -0.4101],
        [-0.4101,  1.3557, -0.4101],
        [-0.4101, -0.4101,  1.3557]])

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