OrthogonalMixingMatrix

class gpjax.models.OrthogonalMixingMatrix(num_outputs, num_latent_gps, key)

Bases: Module

Mixing matrix H = U S^(1/2) with orthogonal columns.

Parameterizes an orthogonal mixing matrix for OILMM where: - U in R^(p x m) has orthonormal columns (U^T U = I_m) - S > 0 is a diagonal scaling matrix (m x m) - H = U S^(1/2) is the mixing matrix - T = S^(-1/2) U^T is the projection matrix

The orthogonality of U ensures that the projected noise is diagonal:

Sigma_T = T Sigma T^T = sigma^2 S^(-1) + D

where sigma^2 is observation noise and D is latent noise.

Parameters:

• num_outputs (int)

• num_latent_gps (int)

• key (Array)

num_outputs

Number of output dimensions (p)

Type:

int

num_latent_gps

Number of latent GP functions (m)

Type:

int

U_latent

Unconstrained matrix for SVD orthogonalization

Type:

gpjax.parameters.Real

S

Positive diagonal scaling

Type:

gpjax.parameters.PositiveReal

obs_noise_variance

Homogeneous observation noise (sigma^2)

Type:

gpjax.parameters.PositiveReal

latent_noise_variance

Per-latent heterogeneous noise (D), non-negative

Type:

gpjax.parameters.NonNegativeReal

property H: Float[jaxlib._jax.Array, 'P M']

Mixing matrix H = U S^(1/2).

Maps from latent space (m dimensions) to output space (p dimensions). Each column is an orthogonal basis vector scaled by sqrt(S_i).

property H_squared: Float[jaxlib._jax.Array, 'P M']

Element-wise H^2 for fast diagonal variance reconstruction.

When computing marginal variances, we need H^2 @ latent_vars:

var_p = sum_m H^2_pm * var_m

This property caches H^2 to avoid recomputation.

property T: Float[jaxlib._jax.Array, 'M P']

Projection matrix T = S^(-1/2) U^T.

Projects from output space (p dimensions) to latent space (m dimensions). This is the left pseudo-inverse of H: T @ H = I_m.

property U: Float[jaxlib._jax.Array, 'P M']

Orthonormal columns via SVD.

Uses SVD to project U_latent onto the Stiefel manifold (orthonormal columns). This ensures U^T U = I_m exactly.

property inv_sqrt_S: Float[jaxlib._jax.Array, 'M']

S^(-1/2).

Type:

Inverse square root of S diagonal

property projected_noise_variance: Float[jaxlib._jax.Array, 'M']

Sigma_T = sigma^2 S^(-1) + D.

This is the noise variance for each independent latent GP after projection. The orthogonality of U ensures this is diagonal, which is what makes OILMM tractable.

Returns:

Array of shape [M] with noise variance for each latent GP.

Type:

Diagonal projected noise

property sqrt_S: Float[jaxlib._jax.Array, 'M']

S^(1/2).

Type:

Square root of S diagonal