xtremax

Extreme value theory for JAX — NumPyro-native EVT distributions, temporal / spatial / spatiotemporal point processes, and xarray-native extremes extraction.

Built on JAX, NumPyro, equinox, and xarray.

New here? Start with the Vision to understand why xtremax exists. The Architecture design doc is a draft roadmap — it sketches where the library is headed, not the current module layout; the API Reference reflects what ships today.

Installation¶

pip install xtremax

Or with uv:

uv add xtremax

Quickstart¶

import jax.random as jr
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

import xtremax

# NumPyro-native EVT distributions
gev = xtremax.GeneralizedExtremeValueDistribution(
    loc=30.0, scale=2.0, concentration=0.1
)
maxima = gev.sample(jr.key(0), (100,))
z100 = gev.return_level(100)  # 100-year return level

# Bayesian inference with NUTS
def model(y):
    loc = numpyro.sample("loc", dist.Normal(30.0, 10.0))
    scale = numpyro.sample("scale", dist.LogNormal(0.0, 1.0))
    conc = numpyro.sample("concentration", dist.Normal(0.0, 0.3))
    numpyro.sample(
        "obs",
        xtremax.GeneralizedExtremeValueDistribution(loc, scale, conc),
        obs=y,
    )

mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=1000)
mcmc.run(jr.key(1), maxima)

Extraction utilities stay in xarray — coordinates and metadata are preserved:

import xtremax.extraction as xtx

# da: any xr.DataArray with a "time" dimension, e.g. daily station temperatures
annual_max = xtx.temporal_block_maxima(da, freq="YE")
u = xtx.quantile_threshold(da, quantile=0.95)
peaks = xtx.decluster_runs(da, threshold=u)

Examples¶

Extreme value distributions — pick the tail behaviour:

Distributions overview — the whole family side by side
GEVD: annual block maxima — the canonical block-maxima workflow
Fréchet: heavy-tailed maxima and Weibull: bounded extremes — fixed-type fits

Peaks over threshold — use more of the data than block maxima:

GPD: peaks-over-threshold — threshold selection, declustering, and GPD fits
GPD: quantile-regression threshold — trend-following thresholds

Non-stationary EVT — when the climate moves under you:

Parametric trend — covariate-dependent parameters
P-splines — smooth non-parametric trends
Neural ODE — learned parameter dynamics

Point processes — each built from scratch with primitives, then fit through NumPyro:

Temporal: homogeneous Poisson, inhomogeneous Poisson, renewal, Hawkes, marked, thinning
Spatial: homogeneous, inhomogeneous, marked
Spatiotemporal: homogeneous, inhomogeneous, marked, Hawkes / ETAS
Neural: latent ODE for irregular IPPs

Tutorials — a longer, real-data curriculum:

Spatial extremes — from one station’s annual maxima to spatial GEV fields and non-stationary trends, on CDS station data over Spain

Installation¶

Quickstart¶

Examples¶

Links¶