Primitives

The distribution math as standalone pure JAX functions — for custom likelihoods, vmapped parameter grids, and non-stationary fields where a distribution object would be in the way. Every family exposes the same core set (log_prob, cdf, icdf, mean, return_level), so switching tail behaviour is a rename rather than a rewrite.

All functions take parameters positionally as (x, loc, scale, shape) and broadcast over them, so you can evaluate a whole spatial field of parameters in one call.

Generalized Extreme Value (GEV)

The block-maxima workhorse. The concentration \(\xi\) spans all three tail types smoothly, and the \(\xi \to 0\) Gumbel limit is handled with a numerically stable branch rather than a discontinuity. Survival functions come in linear and log forms — use gev_log_survival in the far tail where \(1 - F(x)\) underflows.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

gev_log_prob(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Log PDF of the Generalized Extreme Value distribution.

Source code in src/xtremax/primitives/gev.py

def gev_log_prob(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Log PDF of the Generalized Extreme Value distribution."""
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    z = clamped_standardize(x, loc, scale)
    valid, u_safe = _support(shape, z)

    w = _reduced_exponent(z, u_safe, valid)
    log_t = jnp.log1p(u_safe)
    log_pdf = -jnp.log(scale) - (w + log_t) - safe_exp_neg_density(w)

    # Density vanishes at x = ±inf; NaN input propagates.
    result = jnp.where(jnp.isfinite(x) & valid, log_pdf, -jnp.inf)
    # Exact finite endpoint (1 + ξz = 0): the extended-real density limit is
    # ξ-dependent on the Weibull branch — 1/σ at ξ = -1 (the uniform-like
    # case), +inf for ξ < -1 (density diverges toward the endpoint), and 0
    # for -1 < ξ. The support mask alone would report -inf everywhere, which
    # disagrees with the closed upper endpoint the class layer advertises.
    # Tolerance of a few ulps: standardizing the advertised endpoint
    # x = μ - σ/ξ rounds through several float ops, so ξ·z lands near — not
    # exactly on — -1 for ordinary parameter values. Snapping the band to the
    # endpoint limit is benign: the interior density approaches that same
    # limit continuously on both branches.
    at_endpoint = (jnp.abs(shape * z + 1.0) <= 8.0 * jnp.finfo(z.dtype).eps) & (
        jnp.isfinite(x)
    )
    endpoint_val = jnp.where(
        shape == -1.0,
        -jnp.log(scale),
        jnp.where(shape < -1.0, jnp.inf, -jnp.inf),
    )
    result = jnp.where(at_endpoint, endpoint_val, result)
    return jnp.where(jnp.isnan(x), jnp.nan, result)

gev_cdf(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

CDF of the Generalized Extreme Value distribution.

Source code in src/xtremax/primitives/gev.py

def gev_cdf(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """CDF of the Generalized Extreme Value distribution."""
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    z = clamped_standardize(x, loc, scale)
    valid, u_safe = _support(shape, z)

    w = _reduced_exponent(z, u_safe, valid)
    cdf_inside = jnp.exp(-safe_exp_neg(w))
    # Fréchet tails (ξ > 0) map out-of-support to 0; Weibull tails to 1.
    boundary = jnp.where(shape > 0, 0.0, 1.0)
    cdf = jnp.where(valid, cdf_inside, boundary)

    # Extended-real limits F(-inf) = 0, F(+inf) = 1; NaN input propagates.
    cdf = jnp.where(jnp.isinf(x), jnp.where(x > 0, 1.0, 0.0), cdf)
    return jnp.where(jnp.isnan(x), jnp.nan, cdf)

gev_icdf(q: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Quantile function (inverse CDF) of the GEV distribution.

Source code in src/xtremax/primitives/gev.py

def gev_icdf(
    q: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Quantile function (inverse CDF) of the GEV distribution."""
    return _gev_icdf_from_neg_log_q(-jnp.log(q), loc, scale, shape)

gev_survival(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Survival function :math:S(x) = 1 - F(x) of the GEV distribution.

Uses -expm1(log F) rather than 1 - F so the deep upper tail (where :math:F \approx 1 and :math:S is tiny) stays accurate instead of cancelling against 1.0.

Source code in src/xtremax/primitives/gev.py

def gev_survival(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Survival function :math:`S(x) = 1 - F(x)` of the GEV distribution.

    Uses ``-expm1(log F)`` rather than ``1 - F`` so the deep upper tail
    (where :math:`F \approx 1` and :math:`S` is tiny) stays accurate
    instead of cancelling against 1.0.
    """
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    z = clamped_standardize(x, loc, scale)
    valid, u_safe = _support(shape, z)

    w = _reduced_exponent(z, u_safe, valid)
    s_inside = -jnp.expm1(-safe_exp_neg(w))
    # Below the Fréchet (ξ > 0) lower endpoint S = 1; above the Weibull
    # (ξ < 0) upper endpoint S = 0.
    boundary = jnp.where(shape > 0, 1.0, 0.0)
    s = jnp.where(valid, s_inside, boundary)

    # Extended-real limits S(-inf) = 1, S(+inf) = 0; NaN input propagates.
    s = jnp.where(jnp.isinf(x), jnp.where(x > 0, 0.0, 1.0), s)
    return jnp.where(jnp.isnan(x), jnp.nan, s)

gev_log_survival(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Log survival function :math:\log S(x) = \log(1 - F(x)) of the GEV.

Source code in src/xtremax/primitives/gev.py

def gev_log_survival(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Log survival function :math:`\log S(x) = \log(1 - F(x))` of the GEV."""
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    z = clamped_standardize(x, loc, scale)
    valid, u_safe = _support(shape, z)

    w = _reduced_exponent(z, u_safe, valid)
    ls_inside = jnp.log(-jnp.expm1(-safe_exp_neg(w)))
    boundary = jnp.where(shape > 0, 0.0, -jnp.inf)
    ls = jnp.where(valid, ls_inside, boundary)

    # Extended-real limits log S(-inf) = 0, log S(+inf) = -inf; NaN propagates.
    ls = jnp.where(jnp.isinf(x), jnp.where(x > 0, -jnp.inf, 0.0), ls)
    return jnp.where(jnp.isnan(x), jnp.nan, ls)

gev_mean(loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Mean of the GEV distribution (+inf when :math:\xi \ge 1).

.. math:: \mathbb{E}[X] = \mu + \sigma\,\frac{\Gamma(1 - \xi) - 1}{\xi} = \mu + \sigma\,\frac{\mathrm{expm1}(\log\Gamma(1 - \xi))}{\xi}

The expm1 form avoids the catastrophic Γ(1 - ξ) - 1 cancellation for small ξ (which otherwise makes the mean wildly wrong in float32); a Taylor branch covers the immediate neighbourhood of ξ = 0 where the quotient is 0/0.

Source code in src/xtremax/primitives/gev.py

def gev_mean(
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Mean of the GEV distribution (``+inf`` when :math:`\xi \ge 1`).

    .. math:: \mathbb{E}[X] = \mu + \sigma\,\frac{\Gamma(1 - \xi) - 1}{\xi}
        = \mu + \sigma\,\frac{\mathrm{expm1}(\log\Gamma(1 - \xi))}{\xi}

    The ``expm1`` form avoids the catastrophic ``Γ(1 - ξ) - 1`` cancellation
    for small ``ξ`` (which otherwise makes the mean wildly wrong in float32);
    a Taylor branch covers the immediate neighbourhood of ``ξ = 0`` where the
    quotient is ``0/0``.
    """
    shape = jnp.asarray(shape)
    finite = shape < 1.0
    # Sanitize the divergent branch (ξ ≥ 1, and gammaln(1 - ξ) for ξ ≥ 1) so it
    # contributes neither a nan value nor a nan gradient before being masked.
    xi = jnp.where(finite, shape, 0.0)

    # (Γ(1 - ξ) - 1)/ξ = [expm1(g)/g] · [g/ξ] with g = gammaln(1 - ξ). The first
    # factor is the smooth expm1_over_x; the second, g/ξ = γ + Σ_{k≥2} ζ(k)/k
    # ξ^{k-1}, has a removable singularity at ξ = 0. A direct g/ξ loses float32
    # precision for small ξ (tiny gammaln ÷ tiny ξ), so use the ζ-series there.
    g = gammaln(1.0 - xi)
    small = jnp.abs(xi) < _MEAN_SERIES_THRESHOLD
    xi_exact = jnp.where(small, 1.0, xi)
    ratio_exact = g / xi_exact
    xi_series = jnp.where(small, xi, 0.0)
    ratio_series = EULER_GAMMA + xi_series * (
        _ZETA2_2 + xi_series * (_ZETA3_3 + xi_series * _ZETA4_4)
    )
    ratio = jnp.where(small, ratio_series, ratio_exact)
    coeff = expm1_over_x(g) * ratio

    mean = loc + scale * coeff
    return jnp.where(finite, mean, jnp.inf)

gev_return_level(period: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

T-period return level: :math:z_T = F^{-1}(1 - 1/T).

Parameterized by :math:-\log(1 - 1/T) = -\mathrm{log1p}(-1/T) so the quantile stays accurate for large T (where 1 - 1/T rounds to 1 and log(1 - 1/T) would lose all its precision).

Source code in src/xtremax/primitives/gev.py

def gev_return_level(
    period: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""T-period return level: :math:`z_T = F^{-1}(1 - 1/T)`.

    Parameterized by :math:`-\log(1 - 1/T) = -\mathrm{log1p}(-1/T)` so the
    quantile stays accurate for large ``T`` (where ``1 - 1/T`` rounds to 1 and
    ``log(1 - 1/T)`` would lose all its precision).
    """
    neg_log_q = -jnp.log1p(-1.0 / period)
    return _gev_icdf_from_neg_log_q(neg_log_q, loc, scale, shape)

Generalized Pareto (GPD)

The peaks-over-threshold counterpart, defined on exceedances above a threshold (so loc = 0 by convention). Same structure as the GEV family, including the log-survival form for deep-tail work.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

gpd_log_prob(x: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Log PDF of the Generalized Pareto Distribution (loc = 0).

Source code in src/xtremax/primitives/gpd.py

def gpd_log_prob(
    x: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Log PDF of the Generalized Pareto Distribution (loc = 0)."""
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    finite = jnp.isfinite(x)
    x_calc = jnp.where(finite, x, 0.0)
    r = _standardize_exceedance(x_calc, scale)
    valid, v_safe = _support(r, shape)

    w = _reduced_exponent(r, v_safe, valid)
    log_t = jnp.log1p(v_safe)
    # -log σ - (1/ξ + 1) log(1 + ξx/σ) = -log σ - (w + log_t); → -log σ - x/σ.
    log_pdf = -jnp.log(scale) - (w + log_t)

    # Density vanishes at x = ±inf; NaN input propagates.
    in_support = (x >= 0.0) & valid & finite
    result = jnp.where(in_support, log_pdf, -jnp.inf)
    # Exact finite upper endpoint (1 + ξx/σ = 0, ξ < 0): the density limit is
    # ξ-dependent — 1/σ at ξ = -1 (GPD(-1) IS Uniform(0, σ)), +inf for
    # ξ < -1 (density diverges toward the endpoint), and 0 for -1 < ξ < 0.
    # The support mask alone would report -inf everywhere, disagreeing with
    # the closed endpoint the class-layer support advertises.
    # Tolerance of a few ulps: standardizing the advertised endpoint
    # x = -σ/ξ rounds through several float ops, so ξ·r lands near — not
    # exactly on — -1 for ordinary parameter values. Snapping the band to the
    # endpoint limit is benign: the interior density approaches that same
    # limit continuously on both branches.
    at_endpoint = (
        (jnp.abs(shape * r + 1.0) <= 8.0 * jnp.finfo(r.dtype).eps) & finite & (x >= 0.0)
    )
    endpoint_val = jnp.where(
        shape == -1.0,
        -jnp.log(scale),
        jnp.where(shape < -1.0, jnp.inf, -jnp.inf),
    )
    result = jnp.where(at_endpoint, endpoint_val, result)
    return jnp.where(jnp.isnan(x), jnp.nan, result)

gpd_cdf(x: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

CDF of the Generalized Pareto Distribution.

Source code in src/xtremax/primitives/gpd.py

def gpd_cdf(
    x: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """CDF of the Generalized Pareto Distribution."""
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    # x < 0 (incl. -inf) is below support; x = +inf is the unbounded tail. Both
    # get a finite in-support surrogate so no overflowing / parameter-dependent
    # inf reaches the kernel, then are overridden with their limits below.
    below = x < 0.0
    posinf = jnp.isposinf(x)
    x_calc = jnp.where(below | posinf, 0.0, x)
    r = _standardize_exceedance(x_calc, scale)
    valid, v_safe = _support(r, shape)

    w = _reduced_exponent(r, v_safe, valid)
    # 1 - t^{-1/ξ} = 1 - exp(-w) = -expm1(-w); → 1 - exp(-x/σ) in the limit.
    cdf_inside = -jnp.expm1(-w)
    # Beyond the support, ξ > 0 tails map out-of-support CDF to 0 (lower
    # bound); ξ < 0 tails map it to 1 (finite upper bound, including the
    # endpoint x = -σ/ξ where 1 + ξx/σ = 0 exactly).
    boundary = jnp.where(shape < 0, 1.0, 0.0)
    cdf = jnp.where(valid, cdf_inside, boundary)
    cdf = jnp.where(posinf, 1.0, cdf)  # F(+inf) = 1
    # Clamp the universal lower bound (x < 0, incl. -inf, is out of support).
    cdf = jnp.where(below, 0.0, cdf)
    return jnp.where(jnp.isnan(x), jnp.nan, cdf)  # NaN input propagates

gpd_icdf(q: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Quantile function (inverse CDF) of the GPD.

Source code in src/xtremax/primitives/gpd.py

def gpd_icdf(
    q: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Quantile function (inverse CDF) of the GPD."""
    # log(1 - q) via log1p keeps the upper tail (q → 1) accurate.
    return _gpd_icdf_from_log_exceedance(jnp.log1p(-q), scale, shape)

gpd_survival(x: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Survival function :math:S(x) = (1 + \xi x/\sigma)^{-1/\xi} of the GPD.

Computed as exp(-w) directly (rather than 1 - cdf) so the deep tail stays accurate. Equals exp(-x/σ) in the exponential limit.

Source code in src/xtremax/primitives/gpd.py

def gpd_survival(
    x: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Survival function :math:`S(x) = (1 + \xi x/\sigma)^{-1/\xi}` of the GPD.

    Computed as ``exp(-w)`` directly (rather than ``1 - cdf``) so the deep tail
    stays accurate. Equals ``exp(-x/σ)`` in the exponential limit.
    """
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    below = x < 0.0
    posinf = jnp.isposinf(x)
    x_calc = jnp.where(below | posinf, 0.0, x)
    r = _standardize_exceedance(x_calc, scale)
    valid, v_safe = _support(r, shape)

    w = _reduced_exponent(r, v_safe, valid)
    s_inside = jnp.exp(-w)
    # Above the finite Weibull-type upper bound (ξ < 0) S = 0; ξ > 0 has no
    # upper bound so `valid` never fails there for x ≥ 0.
    s = jnp.where(valid, s_inside, 0.0)
    s = jnp.where(posinf, 0.0, s)  # S(+inf) = 0
    # Below x = 0 (incl. -inf) the exceedance is certain, S = 1.
    s = jnp.where(below, 1.0, s)
    return jnp.where(jnp.isnan(x), jnp.nan, s)  # NaN input propagates

gpd_log_survival(x: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Log survival :math:\log S(x) = -\tfrac{1}{\xi}\log(1 + \xi x/\sigma).

Equals -w inside the support, so the deep tail is exact instead of passing through log(1 - cdf).

Source code in src/xtremax/primitives/gpd.py

def gpd_log_survival(
    x: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Log survival :math:`\log S(x) = -\tfrac{1}{\xi}\log(1 + \xi x/\sigma)`.

    Equals ``-w`` inside the support, so the deep tail is exact instead of
    passing through ``log(1 - cdf)``.
    """
    shape = jnp.asarray(shape)
    x = jnp.asarray(x)
    below = x < 0.0
    posinf = jnp.isposinf(x)
    x_calc = jnp.where(below | posinf, 0.0, x)
    r = _standardize_exceedance(x_calc, scale)
    valid, v_safe = _support(r, shape)

    w = _reduced_exponent(r, v_safe, valid)
    ls = jnp.where(valid, -w, -jnp.inf)
    ls = jnp.where(posinf, -jnp.inf, ls)  # log S(+inf) = -inf
    ls = jnp.where(below, 0.0, ls)
    return jnp.where(jnp.isnan(x), jnp.nan, ls)  # NaN input propagates

gpd_mean(scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

GPD mean: :math:\sigma / (1 - \xi) for :math:\xi < 1, else +inf.

Source code in src/xtremax/primitives/gpd.py

def gpd_mean(
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""GPD mean: :math:`\sigma / (1 - \xi)` for :math:`\xi < 1`, else ``+inf``."""
    shape = jnp.asarray(shape)
    # Substitute a safe divisor where shape >= 1 so the masked-out branch
    # traces without a concrete Python ZeroDivisionError.
    denom = jnp.where(shape < 1.0, 1.0 - shape, 1.0)
    mean_val = scale / denom
    return jnp.where(shape < 1.0, mean_val, jnp.inf)

gpd_return_level(period: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

T-period return level: :math:Q(1 - 1/T).

Parameterized by the exceedance probability 1/T directly (log p = -log T) so it stays accurate for large T.

Source code in src/xtremax/primitives/gpd.py

def gpd_return_level(
    period: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    r"""T-period return level: :math:`Q(1 - 1/T)`.

    Parameterized by the exceedance probability ``1/T`` directly (``log p =
    -log T``) so it stays accurate for large ``T``.
    """
    return _gpd_icdf_from_log_exceedance(-jnp.log(period), scale, shape)

Gumbel — GEV Type I

The \(\xi = 0\) case: exponential tails, closed-form moments (\(\mu + \sigma\gamma_E\)), and no shape parameter to estimate.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

gumbel_log_prob(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...]) -> Float[Array, ...]

Gumbel log PDF: :math:-\log\sigma - z - e^{-z}, :math:z=(x-\mu)/\sigma.

Source code in src/xtremax/primitives/gumbel.py

def gumbel_log_prob(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Gumbel log PDF: :math:`-\log\sigma - z - e^{-z}`, :math:`z=(x-\mu)/\sigma`."""
    z = (x - loc) / scale
    return -jnp.log(scale) - z - jnp.exp(-z)

gumbel_cdf(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...]) -> Float[Array, ...]

Gumbel CDF: :math:F(x) = \exp(-\exp(-z)).

Source code in src/xtremax/primitives/gumbel.py

def gumbel_cdf(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Gumbel CDF: :math:`F(x) = \exp(-\exp(-z))`."""
    z = (x - loc) / scale
    return jnp.exp(-jnp.exp(-z))

gumbel_icdf(q: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...]) -> Float[Array, ...]

Gumbel quantile: :math:Q(p) = \mu - \sigma\log(-\log p).

Source code in src/xtremax/primitives/gumbel.py

def gumbel_icdf(
    q: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Gumbel quantile: :math:`Q(p) = \mu - \sigma\log(-\log p)`."""
    return loc - scale * jnp.log(-jnp.log(q))

gumbel_mean(loc: Float[Array, ...], scale: Float[Array, ...]) -> Float[Array, ...]

Gumbel mean: :math:\mu + \sigma\gamma_E.

Source code in src/xtremax/primitives/gumbel.py

def gumbel_mean(
    loc: Float[Array, ...],
    scale: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Gumbel mean: :math:`\mu + \sigma\gamma_E`."""
    return loc + scale * EULER_GAMMA

gumbel_return_level(period: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...]) -> Float[Array, ...]

T-period return level: :math:Q(1 - 1/T).

Source code in src/xtremax/primitives/gumbel.py

def gumbel_return_level(
    period: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
) -> Float[Array, ...]:
    r"""T-period return level: :math:`Q(1 - 1/T)`."""
    return gumbel_icdf(1.0 - 1.0 / period, loc, scale)

Fréchet — GEV Type II

The \(\xi > 0\) case: heavy polynomial tails and lower-bounded support. The mean diverges once \(\xi \ge 1\), and frechet_mean returns +inf there rather than a silently wrong number.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

frechet_log_prob(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Fréchet log PDF (assumes shape > 0).

Source code in src/xtremax/primitives/frechet.py

def frechet_log_prob(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Fréchet log PDF (assumes ``shape > 0``)."""
    return gev_log_prob(x, loc, scale, shape)

frechet_cdf(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Fréchet CDF.

Source code in src/xtremax/primitives/frechet.py

def frechet_cdf(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Fréchet CDF."""
    return gev_cdf(x, loc, scale, shape)

frechet_icdf(q: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Fréchet quantile function.

Source code in src/xtremax/primitives/frechet.py

def frechet_icdf(
    q: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Fréchet quantile function."""
    return gev_icdf(q, loc, scale, shape)

frechet_mean(loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Fréchet mean (+inf when shape >= 1).

Source code in src/xtremax/primitives/frechet.py

def frechet_mean(
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Fréchet mean (``+inf`` when ``shape >= 1``)."""
    return gev_mean(loc, scale, shape)

frechet_return_level(period: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

T-period return level.

Source code in src/xtremax/primitives/frechet.py

def frechet_return_level(
    period: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """T-period return level."""
    return gev_return_level(period, loc, scale, shape)

Weibull (reverse) — GEV Type III

The \(\xi < 0\) case: a bounded upper tail, for quantities with a physical ceiling. Moments are always finite.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

weibull_log_prob(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Reverse-Weibull log PDF (assumes shape < 0).

Source code in src/xtremax/primitives/weibull.py

def weibull_log_prob(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Reverse-Weibull log PDF (assumes ``shape < 0``)."""
    return gev_log_prob(x, loc, scale, shape)

weibull_cdf(x: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Reverse-Weibull CDF.

Source code in src/xtremax/primitives/weibull.py

def weibull_cdf(
    x: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Reverse-Weibull CDF."""
    return gev_cdf(x, loc, scale, shape)

weibull_icdf(q: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Reverse-Weibull quantile function.

Source code in src/xtremax/primitives/weibull.py

def weibull_icdf(
    q: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Reverse-Weibull quantile function."""
    return gev_icdf(q, loc, scale, shape)

weibull_mean(loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

Reverse-Weibull mean (always finite for shape < 0).

Source code in src/xtremax/primitives/weibull.py

def weibull_mean(
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """Reverse-Weibull mean (always finite for ``shape < 0``)."""
    return gev_mean(loc, scale, shape)

weibull_return_level(period: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...]) -> Float[Array, ...]

T-period return level.

Source code in src/xtremax/primitives/weibull.py

def weibull_return_level(
    period: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
) -> Float[Array, ...]:
    """T-period return level."""
    return gev_return_level(period, loc, scale, shape)

Non-stationary EVT

Under a changing climate the return level is no longer a single number — the GEV parameters become functions of covariates, and "the 100-year event" has to be re-stated per year or per scenario. These helpers build covariate design matrices, assemble spatio-temporal parameter fields, and invert between return levels and return periods for both stationary and non-stationary fields.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

design_matrix(covariates: Float[Array, ...]) -> Float[Array, ...]

Prepend an intercept column to a static-covariate matrix.

.. math:: X(s) = [1, z_1(s), \ldots, z_D(s)]

Parameters:

Name	Type	Description	Default
covariates	Float[Array, ...]	Static site covariates of shape (n_sites, n_covariates). A 1-D array of shape (n_sites,) is treated as a single column.	required

Returns:

Type	Description
Float[Array, ...]	Design matrix of shape (n_sites, 1 + n_covariates) whose first
Float[Array, ...]	column is all ones.

Source code in src/xtremax/primitives/spatial.py

def design_matrix(covariates: Float[Array, ...]) -> Float[Array, ...]:
    r"""Prepend an intercept column to a static-covariate matrix.

    .. math:: X(s) = [1, z_1(s), \ldots, z_D(s)]

    Args:
        covariates: Static site covariates of shape ``(n_sites, n_covariates)``.
            A 1-D array of shape ``(n_sites,)`` is treated as a single column.

    Returns:
        Design matrix of shape ``(n_sites, 1 + n_covariates)`` whose first
        column is all ones.
    """
    covariates = jnp.asarray(covariates)
    if covariates.ndim == 1:
        covariates = covariates[:, None]
    n_sites = covariates.shape[0]
    intercept = jnp.ones((n_sites, 1), dtype=covariates.dtype)
    return jnp.concatenate([intercept, covariates], axis=1)

assemble_nonstationary_gev_fields(mu0: Float[Array, ...], mu1: Float[Array, ...], log_sigma: Float[Array, ...], shape: Float[Array, ...], time_covariate: Float[Array, ...]) -> tuple[Float[Array, ...], Float[Array, ...], Float[Array, ...]]

Build non-stationary spatio-temporal GEV parameter fields.

Combines per-site intercept, trend, and log-scale fields with a temporal covariate into block-by-site GEV parameters:

.. math::

\mu_t(s) &= \mu_0(s) + \sum_k x_{t,k}\,\mu_{1,k}(s) \\
\sigma_t(s) &= \exp(\eta_\sigma(s)) \\
\xi_t(s) &= \xi

The scale is constant in time and the shape is broadcast from a scalar or per-site value; only the location varies with the covariate. This is the parameterization underneath hierarchical non-stationary GEV models, decoupled from any particular prior / GP / regression machinery.

Parameters:

Name	Type	Description	Default
mu0	Float[Array, ...]	Spatial intercept :math:\mu_0(s) of shape (n_sites,).	required
mu1	Float[Array, ...]	Trend coefficient field of shape (n_sites,) for a single temporal covariate, or (n_sites, n_covariates) for several.	required
log_sigma	Float[Array, ...]	Log-scale field :math:\log \sigma(s) of shape (n_sites,).	required
shape	Float[Array, ...]	GEV shape :math:\xi, a scalar or per-site array of shape (n_sites,).	required
time_covariate	Float[Array, ...]	Temporal covariate of shape (n_blocks,) for a single covariate, or (n_blocks, n_covariates) for several.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (loc, scale, shape) of fields, each of shape
Float[Array, ...]	(n_blocks, n_sites).

Source code in src/xtremax/primitives/nonstationary.py

def assemble_nonstationary_gev_fields(
    mu0: Float[Array, ...],
    mu1: Float[Array, ...],
    log_sigma: Float[Array, ...],
    shape: Float[Array, ...],
    time_covariate: Float[Array, ...],
) -> tuple[Float[Array, ...], Float[Array, ...], Float[Array, ...]]:
    r"""Build non-stationary spatio-temporal GEV parameter fields.

    Combines per-site intercept, trend, and log-scale fields with a
    temporal covariate into block-by-site GEV parameters:

    .. math::

        \mu_t(s) &= \mu_0(s) + \sum_k x_{t,k}\,\mu_{1,k}(s) \\
        \sigma_t(s) &= \exp(\eta_\sigma(s)) \\
        \xi_t(s) &= \xi

    The scale is constant in time and the shape is broadcast from a scalar
    or per-site value; only the location varies with the covariate. This is
    the parameterization underneath hierarchical non-stationary GEV models,
    decoupled from any particular prior / GP / regression machinery.

    Args:
        mu0: Spatial intercept :math:`\mu_0(s)` of shape ``(n_sites,)``.
        mu1: Trend coefficient field of shape ``(n_sites,)`` for a single
            temporal covariate, or ``(n_sites, n_covariates)`` for several.
        log_sigma: Log-scale field :math:`\log \sigma(s)` of shape
            ``(n_sites,)``.
        shape: GEV shape :math:`\xi`, a scalar or per-site array of shape
            ``(n_sites,)``.
        time_covariate: Temporal covariate of shape ``(n_blocks,)`` for a
            single covariate, or ``(n_blocks, n_covariates)`` for several.

    Returns:
        Tuple ``(loc, scale, shape)`` of fields, each of shape
        ``(n_blocks, n_sites)``.
    """
    mu0 = jnp.asarray(mu0)
    mu1 = jnp.asarray(mu1)
    log_sigma = jnp.asarray(log_sigma)
    time = jnp.asarray(time_covariate)

    if time.ndim == 1:
        time = time[:, None]
    if mu1.ndim == 1:
        mu1 = mu1[:, None]

    # x_{t,k} mu_{1,k}(s) summed over covariates k -> (n_blocks, n_sites).
    time_comp = jnp.einsum("tk,nk->tn", time, mu1)
    loc = mu0[None, :] + time_comp

    scale = jnp.broadcast_to(jnp.exp(log_sigma)[None, :], loc.shape)
    shape_field = jnp.broadcast_to(jnp.asarray(shape), loc.shape)
    return loc, scale, shape_field

nonstationary_return_level(return_period: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...], time_axis: int | None = None, num_bisection_steps: int = 64, lower: Float[Array, ...] | None = None, upper: Float[Array, ...] | None = None) -> Float[Array, ...]

Return level for stationary or non-stationary GEV parameter fields.

Stationary (time_axis is None) uses the closed-form quantile :math:z_T = F^{-1}(1 - 1/T).

Non-stationary (time_axis given): there is no single closed-form quantile because the parameters vary by block, so z_T is found by bisection such that the expected number of exceedances over the supplied blocks equals :math:N_\text{blocks} / T:

.. math:: \sum_t \Pr(Y_t > z_T) = N_\text{blocks} / T.

Parameters:

Name	Type	Description	Default
return_period	Float[Array, ...]	Return period :math:T > 1, broadcastable to the non-time dimensions.	required
loc	Float[Array, ...]	GEV location field.	required
scale	Float[Array, ...]	GEV scale field.	required
shape	Float[Array, ...]	GEV shape field.	required
time_axis	int | None	Axis indexing blocks, or None for the stationary case.	None
num_bisection_steps	int	Bisection iterations for the non-stationary solve.	64
lower	Float[Array, ...] | None	Optional lower bracket for the non-stationary solve. Defaults to a conservative value derived from loc and scale.	None
upper	Float[Array, ...] | None	Optional starting upper bracket. It is grown automatically until it brackets the root, so an underestimate (including the default) is corrected rather than silently returned.	None

Returns:

Type	Description
Float[Array, ...]	Return level :math:z_T with the shape of the non-time dimensions.

Source code in src/xtremax/primitives/nonstationary.py

def nonstationary_return_level(
    return_period: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
    time_axis: int | None = None,
    num_bisection_steps: int = 64,
    lower: Float[Array, ...] | None = None,
    upper: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    r"""Return level for stationary or non-stationary GEV parameter fields.

    Stationary (``time_axis is None``) uses the closed-form quantile
    :math:`z_T = F^{-1}(1 - 1/T)`.

    Non-stationary (``time_axis`` given): there is no single closed-form
    quantile because the parameters vary by block, so ``z_T`` is found by
    bisection such that the expected number of exceedances over the supplied
    blocks equals :math:`N_\text{blocks} / T`:

    .. math:: \sum_t \Pr(Y_t > z_T) = N_\text{blocks} / T.

    Args:
        return_period: Return period :math:`T > 1`, broadcastable to the
            non-time dimensions.
        loc: GEV location field.
        scale: GEV scale field.
        shape: GEV shape field.
        time_axis: Axis indexing blocks, or ``None`` for the stationary case.
        num_bisection_steps: Bisection iterations for the non-stationary solve.
        lower: Optional lower bracket for the non-stationary solve. Defaults
            to a conservative value derived from ``loc`` and ``scale``.
        upper: Optional starting upper bracket. It is grown automatically until
            it brackets the root, so an underestimate (including the default)
            is corrected rather than silently returned.

    Returns:
        Return level :math:`z_T` with the shape of the non-time dimensions.
    """
    loc = jnp.asarray(loc)
    scale = jnp.asarray(scale)
    shape = jnp.asarray(shape)

    if time_axis is None:
        # Route through the stable return-level path (parameterized by
        # -log1p(-1/T)) so large periods do not lose precision to the
        # 1 - 1/T rounding that a direct gev_icdf(1 - 1/T) would incur.
        return gev_return_level(return_period, loc, scale, shape)

    n_blocks = loc.shape[time_axis]
    target = n_blocks / return_period

    def residual(level):
        return (
            expected_exceedances(level, loc, scale, shape, time_axis=time_axis) - target
        )

    if lower is None:
        lower = jnp.min(loc - 8.0 * scale, axis=time_axis)
    if upper is None:
        upper = jnp.max(loc + 20.0 * scale, axis=time_axis)

    lower, upper = jnp.broadcast_arrays(jnp.asarray(lower), jnp.asarray(upper))

    def solve(fn, initial_guess):
        del initial_guess
        # The residual is monotone decreasing in ``level`` and tends to
        # ``-target`` as ``level -> +inf`` (every block's survival vanishes,
        # including past the finite Weibull endpoint where ``gev_survival``
        # saturates to 0). A fixed ``loc + 20*scale`` underestimates heavy
        # Fréchet tails, so grow the upper bracket until it brackets the root.
        up = _expand_upper_bracket(fn, lower, upper)
        return _bisect_monotone_decreasing(fn, lower, up, num_steps=num_bisection_steps)

    # ``custom_root`` supplies the derivative via the implicit function theorem
    # (dz/dθ = -(∂R/∂θ)/(∂R/∂z)) instead of differentiating the bracketing
    # ``while_loop`` (which reverse-mode AD cannot handle) or the discrete
    # bisection comparisons (whose gradient is identically zero). The linear
    # tangent problem is diagonal, so it is solved elementwise.
    def tangent_solve(g, y):
        return y / g(jnp.ones_like(y))

    return jax.lax.custom_root(residual, lower, solve, tangent_solve)

nonstationary_return_period(threshold: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...], time_axis: int | None = None) -> Float[Array, ...]

Return period associated with a threshold under (non-)stationary GEV.

Stationary (time_axis is None):

.. math:: T(z) = 1 / \Pr(Y > z).

Non-stationary (time_axis given): the effective return period from the average exceedance rate over the supplied blocks,

.. math:: T_\text{eff}(z) = N_\text{blocks} / \sum_t \Pr(Y_t > z).

Parameters:

Name	Type	Description	Default
threshold	Float[Array, ...]	Threshold z, broadcastable to the non-time dimensions.	required
loc	Float[Array, ...]	GEV location field.	required
scale	Float[Array, ...]	GEV scale field.	required
shape	Float[Array, ...]	GEV shape field.	required
time_axis	int | None	Axis indexing blocks, or None for the stationary case.	None

Returns:

Type	Description
Float[Array, ...]	Return period with the shape of the non-time dimensions.

Source code in src/xtremax/primitives/nonstationary.py

def nonstationary_return_period(
    threshold: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
    time_axis: int | None = None,
) -> Float[Array, ...]:
    r"""Return period associated with a threshold under (non-)stationary GEV.

    Stationary (``time_axis is None``):

    .. math:: T(z) = 1 / \Pr(Y > z).

    Non-stationary (``time_axis`` given): the *effective* return period from
    the average exceedance rate over the supplied blocks,

    .. math:: T_\text{eff}(z) = N_\text{blocks} / \sum_t \Pr(Y_t > z).

    Args:
        threshold: Threshold ``z``, broadcastable to the non-time dimensions.
        loc: GEV location field.
        scale: GEV scale field.
        shape: GEV shape field.
        time_axis: Axis indexing blocks, or ``None`` for the stationary case.

    Returns:
        Return period with the shape of the non-time dimensions.
    """
    count = expected_exceedances(threshold, loc, scale, shape, time_axis=time_axis)
    if time_axis is None:
        return 1.0 / count
    n_blocks = jnp.asarray(loc).shape[time_axis]
    return n_blocks / count

expected_exceedances(threshold: Float[Array, ...], loc: Float[Array, ...], scale: Float[Array, ...], shape: Float[Array, ...], time_axis: int | None = None) -> Float[Array, ...]

Expected number of block exceedances of a threshold.

For GEV block maxima the per-block exceedance probability is the survival function :math:\Pr(Y_t > z) = 1 - F(z). Summing over blocks gives the expected number of exceedances:

.. math:: \Lambda(z) = \sum_t \Pr(Y_t > z).

Parameters:

Name	Type	Description	Default
threshold	Float[Array, ...]	Threshold / candidate return level z, broadcastable to the non-time dimensions of the parameter fields.	required
loc	Float[Array, ...]	GEV location field.	required
scale	Float[Array, ...]	GEV scale field.	required
shape	Float[Array, ...]	GEV shape field.	required
time_axis	int | None	Axis indexing blocks. If None, the single-block exceedance probability is returned instead of a summed count.	None

Returns:

Name	Type	Description
Either	Float[Array, ...]	math:\Pr(Y > z) (time_axis is None) or
	Float[Array, ...]	math:\sum_t \Pr(Y_t > z) (summed over time_axis).

Source code in src/xtremax/primitives/nonstationary.py

def expected_exceedances(
    threshold: Float[Array, ...],
    loc: Float[Array, ...],
    scale: Float[Array, ...],
    shape: Float[Array, ...],
    time_axis: int | None = None,
) -> Float[Array, ...]:
    r"""Expected number of block exceedances of a threshold.

    For GEV block maxima the per-block exceedance probability is the
    survival function :math:`\Pr(Y_t > z) = 1 - F(z)`. Summing over blocks
    gives the expected number of exceedances:

    .. math:: \Lambda(z) = \sum_t \Pr(Y_t > z).

    Args:
        threshold: Threshold / candidate return level ``z``, broadcastable
            to the non-time dimensions of the parameter fields.
        loc: GEV location field.
        scale: GEV scale field.
        shape: GEV shape field.
        time_axis: Axis indexing blocks. If ``None``, the single-block
            exceedance probability is returned instead of a summed count.

    Returns:
        Either :math:`\Pr(Y > z)` (``time_axis is None``) or
        :math:`\sum_t \Pr(Y_t > z)` (summed over ``time_axis``).
    """
    threshold = jnp.asarray(threshold)
    if time_axis is None:
        return gev_survival(threshold, loc, scale, shape)
    # Honor the documented contract that ``threshold`` broadcasts against the
    # non-time dimensions for *any* ``time_axis`` (not just a leading one) by
    # inserting the block axis. A scalar already broadcasts everywhere.
    if threshold.ndim > 0:
        threshold = jnp.expand_dims(threshold, time_axis)
    survival = gev_survival(threshold, loc, scale, shape)
    return jnp.sum(survival, axis=time_axis)

Spatial helpers

Building blocks for spatial pooling: a distance matrix between sites and a two-range exponential correlation used as a GP covariance for parameter fields.

Pure-JAX primitive functions for extreme value distributions.

Every function takes arrays and returns arrays. No NumPyro, no class state, no side effects — compatible with jax.jit, jax.grad, and jax.vmap.

The Distribution classes in :mod:xtremax.distributions are thin wrappers that route log_prob / cdf / icdf / mean / return_level to these primitives.

pairwise_distances(coords: Float[Array, ...]) -> Float[Array, ...]

Euclidean distance matrix between sites.

.. math:: D_{ij} = \lVert s_i - s_j \rVert_2

Parameters:

Name	Type	Description	Default
coords	Float[Array, ...]	Site coordinates of shape (n_sites, n_dims). For spatial models these are typically normalized longitude / latitude, so the resulting distances are dimensionless.	required

Returns:

Type	Description
Float[Array, ...]	Dense distance matrix of shape (n_sites, n_sites) with a zero
Float[Array, ...]	diagonal.

Source code in src/xtremax/primitives/spatial.py

def pairwise_distances(coords: Float[Array, ...]) -> Float[Array, ...]:
    r"""Euclidean distance matrix between sites.

    .. math:: D_{ij} = \lVert s_i - s_j \rVert_2

    Args:
        coords: Site coordinates of shape ``(n_sites, n_dims)``. For spatial
            models these are typically normalized longitude / latitude, so
            the resulting distances are dimensionless.

    Returns:
        Dense distance matrix of shape ``(n_sites, n_sites)`` with a zero
        diagonal.
    """
    coords = jnp.asarray(coords)
    deltas = coords[:, None, :] - coords[None, :, :]
    # ‖·‖ has derivative Δ/‖Δ‖, which is 0/0 on the zero diagonal and yields NaN
    # gradients. Take sqrt of a masked-positive squared distance, then restore
    # the exact zero diagonal — value and gradient are both finite.
    sq = jnp.sum(deltas**2, axis=-1)
    positive = sq > 0.0
    dist = jnp.where(positive, jnp.sqrt(jnp.where(positive, sq, 1.0)), 0.0)
    # A NaN coordinate makes its squared distances NaN, which ``sq > 0`` reads as
    # False — masking it to a zero (coincident) distance and hiding the bad
    # input. Propagate the NaN instead of collapsing it onto the diagonal.
    return jnp.where(jnp.isnan(sq), jnp.nan, dist)

two_range_correlation(dist_matrix: Float[Array, ...], weight: Float[Array, ...], range_short: Float[Array, ...], range_long: Float[Array, ...], jitter: float = 1e-06) -> Float[Array, ...]

Two-range exponential correlation matrix with unit diagonal.

Convex mixture of two exponential dependence ranges, used as the correlation structure of a Gaussian copula over sites:

.. math:: \rho(h) = c_0 \exp(-h / c_1) + (1 - c_0) \exp(-h / c_2)

where h is the inter-site distance, c_0 = weight mixes the short and long ranges, and c_1 = range_short / c_2 = range_long are the two correlation length scales. The diagonal is renormalized to exactly one (after adding jitter for positive-definiteness) so the result is a valid correlation matrix.

Parameters:

Name	Type	Description	Default
dist_matrix	Float[Array, ...]	Pairwise distances of shape (n_sites, n_sites), e.g. from :func:pairwise_distances.	required
weight	Float[Array, ...]	Mixture weight :math:c_0 \in (0, 1) on the short range.	required
range_short	Float[Array, ...]	Short correlation length scale :math:c_1 > 0.	required
range_long	Float[Array, ...]	Long correlation length scale :math:c_2 > 0.	required
jitter	float	Small positive diagonal inflation added before renormalization for numerical stability.	1e-06

Returns:

Type	Description
Float[Array, ...]	Correlation matrix of shape (n_sites, n_sites) with unit diagonal.

Source code in src/xtremax/primitives/spatial.py

def two_range_correlation(
    dist_matrix: Float[Array, ...],
    weight: Float[Array, ...],
    range_short: Float[Array, ...],
    range_long: Float[Array, ...],
    jitter: float = 1e-6,
) -> Float[Array, ...]:
    r"""Two-range exponential correlation matrix with unit diagonal.

    Convex mixture of two exponential dependence ranges, used as the
    correlation structure of a Gaussian copula over sites:

    .. math:: \rho(h) = c_0 \exp(-h / c_1) + (1 - c_0) \exp(-h / c_2)

    where ``h`` is the inter-site distance, ``c_0 = weight`` mixes the
    short and long ranges, and ``c_1 = range_short`` / ``c_2 = range_long``
    are the two correlation length scales. The diagonal is renormalized to
    exactly one (after adding ``jitter`` for positive-definiteness) so the
    result is a valid correlation matrix.

    Args:
        dist_matrix: Pairwise distances of shape ``(n_sites, n_sites)``,
            e.g. from :func:`pairwise_distances`.
        weight: Mixture weight :math:`c_0 \in (0, 1)` on the short range.
        range_short: Short correlation length scale :math:`c_1 > 0`.
        range_long: Long correlation length scale :math:`c_2 > 0`.
        jitter: Small positive diagonal inflation added before
            renormalization for numerical stability.

    Returns:
        Correlation matrix of shape ``(n_sites, n_sites)`` with unit diagonal.
    """
    dist_matrix = jnp.asarray(dist_matrix)
    n_sites = dist_matrix.shape[0]

    base = weight * jnp.exp(-dist_matrix / range_short) + (1.0 - weight) * jnp.exp(
        -dist_matrix / range_long
    )
    jittered = base + jitter * jnp.eye(n_sites)

    # Renormalize to unit diagonal so the matrix is a valid correlation matrix.
    diag = jnp.sqrt(jnp.clip(jnp.diag(jittered), min=1e-12))
    return jittered / (diag[:, None] * diag[None, :])