HeteroscedasticGEV#

class gpjax.likelihoods.HeteroscedasticGEV(shape=0.1, noise_transform=LogNormalTransform(), num_points=20)[source]#

Bases: AbstractHeteroscedasticLikelihood

GEV likelihood with a latent location and a latent scale.

Two independent latent Gaussian processes give the GEV location and the transformed GEV scale; the shape \(\xi\) is a parameter that all inputs share:

\[y \mid f(x), g(x) \sim \operatorname{GEV}\big(\mu = f(x),\, \sigma = h(g(x)),\, \xi\big).\]

The transform \(h\) is noise_transform. The default, LogNormalTransform, gives \(\sigma = e^{g}\), the log link that is usual for the GEV scale. Note the difference from HeteroscedasticGaussian, where the transform gives a variance: here it gives the GEV scale itself. Only the transform is used, not its moments.

Build the model with HeteroscedasticModel and fit a HeteroscedasticVariationalFamily with heteroscedastic_elbo(). This likelihood has no closed-form bound, so heteroscedastic_elbo uses the chained bound, with expected_log_likelihood().

link_function() uses the support continuation of GeneralisedExtremeValue and predict() is exact, as for GEV.

Fitting. HeteroscedasticVariationalFamily is not whitened, and its fit depends strongly on the start and on the step size. If you use it, start the variational means at the Gumbel moment estimates of GEV (variational_mean_f at \(\mu_0\) and variational_mean_g at \(\log\sigma_0\) for \(\sigma = e^{g}\)), with constant mean functions at the same values, and use a small step size. Two whitened families, whose latents start at their mean functions, are easier to fit. Give their diagonal predictions to expected_log_likelihood():

def negative_elbo(model, data):  # model holds q_f, q_g and likelihood
    f = model.q_f.condition(data)(data.X, covariance="diagonal")
    g = model.q_g.condition(data)(data.X, covariance="diagonal")
    expected = model.likelihood.expected_log_likelihood(
        data.y, f.mean[:, None], f.variance[:, None],
        mean_g=g.mean[:, None], variance_g=g.variance[:, None],
    )
    kl = model.q_f.prior_kl() + model.q_g.prior_kl()
    return kl - jnp.sum(expected) * data.full_size / data.n

Here q_f and q_g are WhitenedVariationalGaussian, with constant mean functions at \(\mu_0\) and \(\log\sigma_0\) and with fixed inducing inputs, as for GEV. With many observations at each input (for example many years at each station), gather the moments of each station by a station index. With one observation at each input, a location and a scale that both vary are weakly identified.

Example

>>> import gpjax as gpx
>>> prior = gpx.gps.Prior(
...     mean_function=gpx.mean_functions.Constant(40.0),
...     kernel=gpx.kernels.RBF(),
... )
>>> noise_prior = gpx.gps.Prior(
...     mean_function=gpx.mean_functions.Constant(2.5),
...     kernel=gpx.kernels.RBF(),
... )
>>> model = gpx.gps.HeteroscedasticModel(
...     prior=prior,
...     likelihood=gpx.likelihoods.HeteroscedasticGEV(shape=0.1),
...     noise_prior=noise_prior,
... )

Expand for references to gpjax.likelihoods.HeteroscedasticGEV

HeteroscedasticGEV

Parameters:

shape (Any)
noise_transform (AbstractNoiseTransform)
num_points (int)

expected_log_likelihood(y, mean, variance, mean_g=None, variance_g=None, **_)[source]#

The expected log-likelihood under independent Gaussians of f and g.

\[\mathbb{E}_{q(f)\,q(g)}\left[\log p(y \mid f, g)\right], \quad q(f) = \mathcal{N}(m_f, s_f),\ q(g) = \mathcal{N}(m_g, s_g),\]

by a tensor-product Gauss–Hermite rule with num_points**2 nodes, with the continued log density. Other keyword arguments, for example the noise_stats that heteroscedastic_elbo passes, are ignored.

The function takes the marginal moments directly, for any number \(N\) of observations, and it can be used under jax.jit and jax.grad. For example, with many years of maxima at each station, gather the moments of each station by a station index: mean[station_index]. The memory is of order N * D * num_points**2.

Parameters:

y (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D']) – The observations, shape (N, D).
mean (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D']) – The means of the latent location, shape (N, D).
variance (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D']) – The variances of the latent location, shape (N, D).
mean_g (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D'] | None) – The means of the latent scale process, shape (N, D).
variance_g (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D'] | None) – The variances of the latent scale process, shape (N, D).
**_ (Any) – Ignored.

Returns:

The expected log-likelihood of each observation, summed over the D columns, shape (N,).

Raises:

ValueError – If mean_g or variance_g is not given.

Return type:

Float[jaxlib._jax.Array, ‘N’] | Float[ndarray, ‘N’]

link_function(f, g=None)[source]#

The GEV distribution of the observations, given f and g.

Parameters:

f (Float[Array, ...]) – The latent location values.
g (tp.Optional[Float[Array, ...]]) – The latent scale values. Required.

Returns:

The distribution of \(y\) given \(f\) and \(g\),: with the support continuation.

Return type:

GeneralisedExtremeValue

Raises:

ValueError – If g is not given.

predict(dist, noise_dist=None)[source]#

The plug-in predictive distribution at the latent means.

The result is the exact GEV with the location at the mean of dist and the scale at the transform of the mean of noise_dist (for \(\sigma = e^{g}\), this is the posterior median of \(\sigma\)). The latent uncertainty is not integrated out; use return_level_samples() for that.

Parameters:

dist (MultivariateNormal | GaussianDistribution) – The latent location at a finite set of test inputs.
noise_dist (MultivariateNormal | GaussianDistribution | None) – The latent scale process at the same inputs.

Returns:

The exact GEV, one for each test input.

Return type:

GeneralisedExtremeValue

Raises:

ValueError – If noise_dist is not given.

return_level_samples(dist, noise_dist, return_period, *, key, num_samples=1000)[source]#

Posterior samples of the return level at each test input.

As GEV.return_level_samples(), but each sample draws the latent location from dist and, independently, the latent scale process from noise_dist, and uses the scale \(h(g)\) of that draw. The samples include the uncertainty of both latent processes, not of the shape, which is a point estimate, or of the kernel hyperparameters.

With a fitted HeteroscedasticVariationalFamily q, use the two predictive distributions of q and the likelihood of q.model, which holds the fitted shape:

dist_f, dist_g = q.predict_latents(test_inputs)
levels = q.model.likelihood.return_level_samples(
    dist_f, dist_g, jnp.array([10.0, 100.0]), key=key
)

Parameters:

dist (MultivariateNormal | GaussianDistribution) – The posterior of the latent location at \(N\) test inputs.
noise_dist (MultivariateNormal | GaussianDistribution) – The posterior of the latent scale process at the same inputs.
return_period (float | int | Float[jaxlib._jax.Array, '...'] | Float[ndarray, '...'] | Int[jaxlib._jax.Array, '...'] | Int[ndarray, '...']) – One return period \(T > 1\), or an array of them.
key (UInt32[jaxlib._jax.Array, '2'] | Key[jaxlib._jax.Array, '']) – The random key for the latent draws.
num_samples (int) – The number of samples \(S\).

Returns:

The return levels, of shape (S, N) for one return period and (S, *jnp.shape(return_period), N) for an array of them.

Return type:

Float[jaxlib._jax.Array, ‘S …’] | Float[ndarray, ‘S …’]

return_period_samples(dist, noise_dist, level, *, key, num_samples=1000)[source]#

Posterior samples of the return period of level at each test input.

As GEV.return_period_samples() (including inf samples and paired samples), with the latent scale drawn as in return_level_samples().

Parameters:

dist (MultivariateNormal | GaussianDistribution) – The posterior of the latent location at \(N\) test inputs.
noise_dist (MultivariateNormal | GaussianDistribution) – The posterior of the latent scale process at the same inputs.
level (float | int | Float[jaxlib._jax.Array, '...'] | Float[ndarray, '...'] | Int[jaxlib._jax.Array, '...'] | Int[ndarray, '...']) – The level: a number, or an array that broadcasts to (N,).
key (UInt32[jaxlib._jax.Array, '2'] | Key[jaxlib._jax.Array, '']) – The random key for the latent draws.
num_samples (int) – The number of samples \(S\).

Returns:

The return periods, of shape (S, N).

Return type:

Float[jaxlib._jax.Array, ‘S N’] | Float[ndarray, ‘S N’]