Likelihoods

Observation-noise models for the Tier I likelihoods: GaussianNoise (the default, today's numbers), StudentTNoise for heavy-tailed retrievals near the detection floor, and HeteroscedasticNoise for a per-pixel retrieval-error map. Pass one as noise= to the plume and puff inference models; log_prob gives the same likelihood in plain JAX for a training loss. Tier I.

likelihoods

Observation-noise models for the Tier I likelihoods (implementation in plumax.gauss_plume._src.likelihoods).

GaussianNoise

Bases: Noise

Homoscedastic Gaussian noise, ε ~ 𝒩(0, σ²).

log p(y | ŷ) = −½ ((y − ŷ)/σ)² − log σ − ½ log 2π.

The Tier I models default to this with σ = obs_noise_std, so their numbers are unchanged.

Attributes:

Name	Type	Description
std	float | Float[Array, '*#obs']	Noise standard deviation σ > 0, in the units of y; a scalar or an array that broadcasts against the observations.

Raises:

Type	Description
ParameterError	If std is not positive.

Examples:

The log-density at the mean of a unit Gaussian is −½ log 2π:

>>> import jax.numpy as jnp
>>> noise = GaussianNoise(std=1.0)
>>> value = noise.log_prob(jnp.zeros(3), jnp.zeros(3))
>>> value.shape
(3,)
>>> bool(jnp.allclose(value, -0.5 * jnp.log(2.0 * jnp.pi)))
True

A non-positive scale is rejected:

>>> GaussianNoise(std=0.0)
Traceback (most recent call last):
...
ParameterError: std must be > 0, got 0.0

HeteroscedasticNoise

Bases: Noise

Per-observation Gaussian noise, ε_i ~ 𝒩(0, σ_i²), from a retrieval-error map.

log p(y_i | ŷ_i) = −½ ((y_i − ŷ_i)/σ_i)² − log σ_i − ½ log 2π.

Unlike :class:GaussianNoise, σ must give one value per observation, and a mismatch with the observations is an error rather than a broadcast.

Attributes:

Name	Type	Description
std	Float[Array, ' obs']	Per-observation standard deviation σ_i > 0, in the units of y.

Raises:

Type	Description
ParameterError	If any std is not positive, or (from log_prob and the NumPyro site) if its shape differs from the observations'.

Examples:

Doubling σ_i halves the weight of a residual:

>>> import jax.numpy as jnp
>>> noise = HeteroscedasticNoise(std=jnp.array([1.0, 2.0]))
>>> value = noise.log_prob(jnp.array([2.0, 4.0]), jnp.zeros(2))
>>> bool(jnp.allclose(value[0] - value[1], jnp.log(2.0)))
True

Build it from variances, e.g. Instrument.observation_variance():

>>> HeteroscedasticNoise.from_variance(jnp.array([4.0, 9.0])).std.tolist()
[2.0, 3.0]

from_variance classmethod

from_variance(variance: Float[Array, ' obs']) -> HeteroscedasticNoise

The noise with σ_i = √(variance_i).

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@classmethod
def from_variance(cls, variance: Float[Array, " obs"]) -> HeteroscedasticNoise:
    """The noise with σ_i = √(variance_i)."""
    # σ_i = √(R_ii)
    std = jnp.sqrt(jnp.asarray(variance))
    return cls(std=std)

Noise

Bases: Module

Observation-noise model p(y | ŷ) for a Tier I likelihood.

Subclasses define :meth:distribution (the NumPyro distribution, for models) and :meth:log_prob (plain JAX, for losses); they must agree.

distribution abstractmethod

distribution(mean: Float[Array, '*obs']) -> dist.Distribution

The NumPyro distribution of y given the prediction ŷ = mean.

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@abc.abstractmethod
def distribution(self, mean: Float[Array, "*obs"]) -> dist.Distribution:
    """The NumPyro distribution of y given the prediction ŷ = ``mean``."""

log_prob abstractmethod

log_prob(obs: Float[Array, '*obs'], mean: Float[Array, '*obs']) -> Float[Array, '*obs']

Pointwise log-density log p(y | ŷ), without NumPyro.

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@abc.abstractmethod
def log_prob(
    self, obs: Float[Array, "*obs"], mean: Float[Array, "*obs"]
) -> Float[Array, "*obs"]:
    """Pointwise log-density log p(y | ŷ), without NumPyro."""

StudentTNoise

Bases: Noise

Heavy-tailed Student-t noise, ε ~ t_ν(0, s), for outliers near the detection floor.

log p(y | ŷ) = log Γ((ν+1)/2) − log Γ(ν/2) − ½ log(ν π) − log s − (ν+1)/2 · log(1 + ((y − ŷ)/s)² / ν).

ν ≈ 5 is the Tier I design's default for low-SNR retrievals; as ν → ∞ it tends to :class:GaussianNoise with σ = s.

Attributes:

Name	Type	Description
scale	float | Float[Array, '*#obs']	Scale s > 0, in the units of y; a scalar or an array that broadcasts against the observations. The standard deviation is s √(ν/(ν − 2)) for ν > 2.
df	float	Degrees of freedom ν > 0.

Raises:

Type	Description
ParameterError	If scale or df is not positive.

Examples:

A 10-scale outlier costs far less than under a Gaussian:

>>> import jax.numpy as jnp
>>> obs, mean = jnp.array([10.0]), jnp.array([0.0])
>>> student = StudentTNoise(scale=1.0).log_prob(obs, mean)
>>> gaussian = GaussianNoise(std=1.0).log_prob(obs, mean)
>>> bool(student[0] > gaussian[0] + 30.0)
True

Large ν recovers the Gaussian:

>>> near = StudentTNoise(scale=1.0, df=1.0e8).log_prob(obs, mean)
>>> bool(jnp.allclose(near, gaussian, rtol=1.0e-6))
True