Simulations

Synthetic extremes generators for testing, teaching, and method development — data with a known ground truth, so you can check whether an estimator recovers the parameters you put in.

Unlike the rest of the package these are NumPy-based and take an integer seed rather than a JAX PRNG key.

Extreme-value simulators

Per-variable simulators with physically sensible defaults: temperature block maxima from a GEV, precipitation extremes, and wind gusts from a Weibull. Each salts its RNG with a per-variable constant, so the same seed across variables still yields independent fields rather than rank-correlated ones.

Synthetic extreme value data generators.

Temporal trajectories (GMST, energy balance), spatial domain/terrain generators, and variable-specific extreme generators for temperature, precipitation, and wind.

simulate_temp_extremes(mu: xr.DataArray, scale: float = 1.5, shape: float = -0.1, seed: int = 42) -> xr.Dataset

Simulates Block Maxima Temperature via the GEV distribution.

Parameters:

Name	Type	Description	Default
mu	DataArray	Location parameter (varying over time/space).	required
scale	float	Scale parameter (assumed constant here, but could be array).	1.5
shape	float	Shape parameter (xi).	-0.1
seed	int	Reproducibility seed. The stream is salted per variable, so temperature, precipitation, and wind simulators called with equal seeds still draw independently of each other.	42

Source code in src/xtremax/simulations/extremes.py

def simulate_temp_extremes(
    mu: xr.DataArray,
    scale: float = 1.5,
    # Negative shape = bounded (Weibull-like); positive = heavy tail (Frechet).
    shape: float = -0.1,
    seed: int = 42,
) -> xr.Dataset:
    """
    Simulates Block Maxima Temperature via the GEV distribution.

    Args:
        mu: Location parameter (varying over time/space).
        scale: Scale parameter (assumed constant here, but could be array).
        shape: Shape parameter (xi).
        seed: Reproducibility seed. The stream is salted per variable, so
            temperature, precipitation, and wind simulators called with
            equal seeds still draw independently of each other.
    """
    rng = _variable_rng("temp", seed)

    # Scipy uses shape 'c' where c = -xi (sign flip vs typical EVT notation).
    # However, standard GEV: mu + (sigma/xi)* ((p^-xi) - 1)
    # We will use simple inversion sampling or scipy's rvs

    # Note: Scipy genextreme args are (c, loc, scale). c = -shape parameter (xi)
    # if xi < 0 (Weibull domain), scipy c > 0
    c = -shape

    data = genextreme.rvs(
        c, loc=mu.values, scale=scale, size=mu.shape, random_state=rng
    )

    ds = mu.to_dataset(name="mu_tmax")
    ds["tmax"] = (mu.dims, data)
    ds.attrs["distribution"] = "GEV"
    return ds

simulate_precip_extremes(spatial_ds: xr.Dataset, gmst_da: xr.DataArray, seed: int = 42) -> xr.Dataset

Simulates Precipitation Extremes: 1. Intensity (Rx1day): Gamma or GEV distribution. 2. Duration (Consecutive Wet Days - CWD): Poisson/Geometric approach.

The seed is salted per variable (see :func:simulate_temp_extremes).

Note: CWD is modelled as a Poisson draw whose mean is floored at 1.0 — a count model, not a true maximum over wet runs; it captures the trend in duration, not the extremal behaviour of run lengths.

Source code in src/xtremax/simulations/extremes.py

def simulate_precip_extremes(
    spatial_ds: xr.Dataset, gmst_da: xr.DataArray, seed: int = 42
) -> xr.Dataset:
    """
    Simulates Precipitation Extremes:
    1. Intensity (Rx1day): Gamma or GEV distribution.
    2. Duration (Consecutive Wet Days - CWD): Poisson/Geometric approach.

    The seed is salted per variable (see :func:`simulate_temp_extremes`).

    Note: CWD is modelled as a Poisson draw whose mean is floored at 1.0
    — a count model, not a true maximum over wet runs; it captures the
    trend in duration, not the extremal behaviour of run lengths.
    """
    if "year" not in gmst_da.sizes:
        raise ValueError("gmst_da must include a 'year' dimension.")
    if "site" not in spatial_ds.sizes:
        raise ValueError("spatial_ds must include a 'site' dimension.")
    rng = _variable_rng("precip", seed)

    # --- A. Intensity (Rx1day - Annual Max Precip) ---
    # Physical intuition: Warmer air holds more moisture (Clausius-Clapeyron ~7%/K)
    # Base precip depends on elevation
    base_intensity = 40.0 + 0.01 * spatial_ds["elevation"]  # mm
    cc_scaling = 1.0 + 0.07 * gmst_da  # Simple 7% per degree scaling

    # Expected intensity location parameter
    loc_intensity = base_intensity * cc_scaling

    # Generate using Gamma distribution (Shape/Scale parameterization)
    # Mean = shape * scale. Let's fix shape, vary scale.
    gamma_shape = 4.0
    gamma_scale = loc_intensity / gamma_shape

    # Values must be broadcast manually for numpy sampling if using arrays
    # But Xarray math handles the broadcasting for parameters
    # We sample (transpose by *name* so inputs with unexpected dims error
    # clearly instead of being silently mislabeled):
    rx1day = rng.gamma(
        shape=gamma_shape,
        scale=gamma_scale.transpose("year", "site").values,
    )

    # --- B. Duration (CWD - Max consecutive wet days) ---
    # Duration may decrease slightly with GMST in the Mediterranean (drier summers).
    base_duration = 10.0 + 0.001 * spatial_ds["elevation"]  # days
    trend_duration = 1.0 - 0.2 * gmst_da

    lambda_param = base_duration * trend_duration
    # Floor the Poisson mean at 1.0 (not just positivity — it clamps the
    # low-duration tail so every site averages at least one wet day).
    lambda_param = np.maximum(1.0, lambda_param)

    # Poisson for count data
    cwd = rng.poisson(lambda_param.transpose("year", "site").values)

    ds = xr.Dataset(
        coords={"year": gmst_da.year, "site": spatial_ds.site},
        data_vars={
            "rx1day": (("year", "site"), rx1day),
            "cwd": (("year", "site"), cwd),
            "lon": spatial_ds.lon,
            "lat": spatial_ds.lat,
            "elevation": spatial_ds.elevation,
        },
    )
    return ds

simulate_wind_extremes(spatial_ds: xr.Dataset, gmst_da: xr.DataArray, seed: int = 42) -> xr.Dataset

Simulates Extreme Wind Speeds (Gusts) using Weibull distribution.

The seed is salted per variable (see :func:simulate_temp_extremes).

Source code in src/xtremax/simulations/extremes.py

def simulate_wind_extremes(
    spatial_ds: xr.Dataset, gmst_da: xr.DataArray, seed: int = 42
) -> xr.Dataset:
    """
    Simulates Extreme Wind Speeds (Gusts) using Weibull distribution.

    The seed is salted per variable (see :func:`simulate_temp_extremes`).
    """
    rng = _variable_rng("wind", seed)
    # Wind often higher at higher elevation and near coast (ignored coast for now)
    base_wind = 15.0 + 0.01 * spatial_ds["elevation"]  # m/s

    # Trend: Maybe slight increase or noise. Let's assume just noise around trend.
    # Note: Weibull is defined by Shape (k) and Scale (lambda)
    # Mean = lambda * Gamma(1 + 1/k)

    k_shape = 2.0  # Rayleigh-like
    w_scale = base_wind  # Scale parameter varies by site

    # Broadcast to time
    w_scale_time = w_scale + (0.5 * gmst_da)  # Slight increase with warming?

    # Sample
    # Scipy weibull_min takes 'c' as shape parameter
    wind_max = weibull_min.rvs(
        c=k_shape,
        scale=w_scale_time.transpose("year", "site").values,
        loc=0,
        random_state=rng,
    )

    ds = xr.Dataset(
        coords={"year": gmst_da.year, "site": spatial_ds.site},
        data_vars={
            "wind_max": (("year", "site"), wind_max),
            "lon": spatial_ds.lon,
            "lat": spatial_ds.lat,
            "elevation": spatial_ds.elevation,
        },
    )
    return ds

Climate signals & GMST

Covariate trajectories for non-stationary EVT — a synthetic global mean surface temperature series to regress the GEV parameters against. generate_gmst_trajectory is the cheap statistical version; generate_physical_gmst runs a 0-D energy-balance model for a trajectory with physical structure rather than a fitted curve.

Synthetic extreme value data generators.

Temporal trajectories (GMST, energy balance), spatial domain/terrain generators, and variable-specific extreme generators for temperature, precipitation, and wind.

generate_gmst_trajectory(n_years: int = 40, start_year: int = 1981, trend_type: Literal['linear', 'exponential', 'logistic'] = 'linear', noise_std: float = 0.05, seed: int = 42) -> xr.DataArray

Generates a synthetic Global Mean Surface Temperature (GMST) anomaly curve.

Parameters:

Name	Type	Description	Default
trend_type	Literal['linear', 'exponential', 'logistic']	Shape of the warming curve.	'linear'

Source code in src/xtremax/simulations/temporal.py

def generate_gmst_trajectory(
    n_years: int = 40,
    start_year: int = 1981,
    trend_type: Literal["linear", "exponential", "logistic"] = "linear",
    noise_std: float = 0.05,
    seed: int = 42,
) -> xr.DataArray:
    """
    Generates a synthetic Global Mean Surface Temperature (GMST) anomaly curve.

    Args:
        trend_type: Shape of the warming curve.
    """
    rng = np.random.default_rng(seed)
    years = np.arange(start_year, start_year + n_years)
    t = np.linspace(0, 1, n_years)  # Normalized time 0 to 1

    # Deterministic Trend Component
    if trend_type == "linear":
        # Simple linear warming: +0.8C over period
        trend = 0.1 + 0.8 * t
    elif trend_type == "exponential":
        # Accelerating warming
        trend = 0.1 + 0.8 * (t**2)
    elif trend_type == "logistic":
        # S-curve (stabilization scenario)
        trend = 1.0 / (1 + np.exp(-10 * (t - 0.5)))
    else:
        raise ValueError(
            f"Unknown trend_type: {trend_type!r}. "
            "Choose from: 'linear', 'exponential', 'logistic'"
        )

    # Add AR(1) red noise to mimic internal climate variability:
    # noise[i] = alpha * noise[i-1] + epsilon[i], expressed as the IIR
    # filter 1 / (1 - alpha z^-1) applied to the innovations.
    epsilon = rng.normal(0, noise_std, n_years)
    alpha = 0.6  # Autocorrelation factor
    noise = lfilter([1.0], [1.0, -alpha], epsilon)

    gmst = trend + noise

    return xr.DataArray(
        gmst,
        coords={"year": years},
        dims="year",
        name="gmst",
        attrs={"description": f"Synthetic GMST ({trend_type})"},
    )

generate_physical_gmst(n_years: int = 100, start_year: int = 1900, climate_sensitivity: float = 3.0, ocean_heat_capacity: float = 10.0, seed: int = 42) -> xr.Dataset

Simulates Global Mean Surface Temperature (GMST) using a 0-D Energy Balance Model.

The evolution of temperature anomaly T(t) is governed by the ODE:

C * dT/dt = F(t) - lambda * T(t)

Where

C : Effective heat capacity of the system (ocean mixed layer). lambda : Climate feedback parameter (Watts / m2 / K). F(t) : Total Radiative Forcing (Watts / m2).

The feedback parameter lambda is derived from Equilibrium Climate Sensitivity (ECS): lambda = F_2xCO2 / ECS (F_2xCO2 is approx 3.7 W/m2)

Parameters:

Name	Type	Description	Default
n_years	int	Duration of simulation.	100
climate_sensitivity	float	Equilibrium warming for doubled CO2.	3.0
ocean_heat_capacity	float	Thermal inertia (higher = slower response).	10.0

Returns:

Type	Description
Dataset	xr.Dataset containing Temperature, Total Forcing, and Components.

Source code in src/xtremax/simulations/temporal.py

def generate_physical_gmst(
    n_years: int = 100,
    start_year: int = 1900,
    climate_sensitivity: float = 3.0,  # Equilibrium CS (deg C for 2xCO2)
    ocean_heat_capacity: float = 10.0,  # Effective capacity (Watt-year / m2 / K)
    seed: int = 42,
) -> xr.Dataset:
    """
    Simulates Global Mean Surface Temperature (GMST) using a 0-D Energy Balance Model.

    The evolution of temperature anomaly T(t) is governed by the ODE:

      C * dT/dt = F(t) - lambda * T(t)

    Where:
      C      : Effective heat capacity of the system (ocean mixed layer).
      lambda : Climate feedback parameter (Watts / m2 / K).
      F(t)   : Total Radiative Forcing (Watts / m2).

    The feedback parameter lambda is derived from Equilibrium Climate Sensitivity (ECS):
      lambda = F_2xCO2 / ECS
      (F_2xCO2 is approx 3.7 W/m2)

    Args:
        n_years: Duration of simulation.
        climate_sensitivity: Equilibrium warming for doubled CO2.
        ocean_heat_capacity: Thermal inertia (higher = slower response).

    Returns:
        xr.Dataset containing Temperature, Total Forcing, and Components.
    """
    rng = np.random.default_rng(seed)

    # --------------------------------------------------------------------------
    # 1. Physics Constants & Setup
    # --------------------------------------------------------------------------
    # Exclude the terminal endpoint: `np.linspace(0, n_years, n_years * 12)`
    # includes `t = n_years`, which falls in a floor(t) = n_years bin and
    # produces an extra final-year row with a single sample. Using
    # endpoint=False keeps exactly `n_years` monthly bins per integer year.
    t_eval = np.linspace(0, n_years, n_years * 12, endpoint=False)
    F_2xCO2 = 3.7  # Radiative forcing for doubling CO2 (W/m2)

    # Calculate Feedback Parameter (lambda)
    # At equilibrium: 0 = F_2xCO2 - lambda * ECS  =>  lambda = 3.7 / ECS
    lam = F_2xCO2 / climate_sensitivity

    # --------------------------------------------------------------------------
    # 2. Define Forcing Components F(t)
    # --------------------------------------------------------------------------

    # A. Greenhouse Gases (Logarithmic relation to CO2, Logistic CO2 growth)
    #    CO2(t) ~ Logistic curve from 280ppm to 560ppm (doubling)
    #    F_ghg(t) = 5.35 * ln(CO2(t) / CO2_ref)
    def forcing_ghg(t):
        # Center the logistic rise at year 50 (relative)
        sigmoid = 1 / (1 + np.exp(-0.1 * (t - (n_years / 2))))
        # Scale forcing from 0 to F_2xCO2 approx
        return F_2xCO2 * sigmoid

    # B. Solar Cycles (11-year Schwabe cycle)
    #    Amplitude approx 0.1 W/m2
    def forcing_solar(t):
        return 0.1 * np.sin(2 * np.pi * t / 11.0)

    # C. Volcanic Eruptions (Stochastic Spikes)
    #    modeled as discrete negative impulses decaying exponentially
    n_eruptions = int(n_years / 10)  # Approx 1 per decade
    # Reserve a 5-year buffer at both ends; if the run is shorter than that
    # window (or n_eruptions < 1) fall back to an empty schedule. Without
    # this guard `rng.uniform(5, n_years - 5, ...)` raises ValueError for
    # `n_years < 10` since `high < low`.
    if n_eruptions >= 1 and n_years >= 10:
        eruption_times = np.sort(rng.uniform(5, n_years - 5, n_eruptions))
        eruption_magnitudes = rng.gamma(shape=2.0, scale=1.5, size=n_eruptions)  # W/m2
    else:
        eruption_times = np.empty(0, dtype=float)
        eruption_magnitudes = np.empty(0, dtype=float)

    def forcing_volcano(t):
        # Vectorised over t: sum the pulse of every eruption already in
        # the past. Eruptions cause cooling (negative forcing) with a
        # rapid onset and slow decay (~2 year lifetime).
        t_arr = np.asarray(t, dtype=float)
        dt = t_arr[..., None] - eruption_times  # (..., n_eruptions)
        active = dt > 0
        # Zero out inactive entries *before* exp so future eruptions
        # (dt < 0) can't overflow exp(-dt/2).
        dt_safe = np.where(active, dt, 0.0)
        pulse = eruption_magnitudes * np.exp(-dt_safe / 2.0) * (dt_safe * 2.0)
        val = -np.sum(np.where(active, pulse, 0.0), axis=-1)
        return val if val.ndim else float(val)

    # D. Stochastic Weather/Internal Variability (Ornstein-Uhlenbeck / Red Noise)
    #    Since ODE solvers need continuous functions, we pre-generate noise
    #    and interpolate it.
    noise_dt = 0.1  # High res noise generation
    noise_steps = int(n_years / noise_dt)
    noise_t = np.linspace(0, n_years, noise_steps)
    white_noise = rng.normal(0, 0.2, size=noise_steps)

    # Generate Red Noise (AR1): red[i] = alpha*red[i-1] + (1-alpha)*white[i]
    # with red[0] = 0, via the IIR filter (1-alpha) / (1 - alpha z^-1).
    # Zeroing the first innovation reproduces the red[0] = 0 start.
    alpha = 0.95
    innovations = white_noise.copy()
    innovations[0] = 0.0
    red_noise = lfilter([1.0 - alpha], [1.0, -alpha], innovations)

    # Create continuous noise function (0 outside the sampled range).
    def forcing_noise_func(t):
        return np.interp(t, noise_t, red_noise, left=0.0, right=0.0)

    # --------------------------------------------------------------------------
    # 3. Solve ODE
    # --------------------------------------------------------------------------

    def system_dynamics(t, y):
        """
        dy/dt = (1/C) * (F_total(t) - lambda * y)
        """
        T = y[0]

        # Aggregate Forcings
        F_g = forcing_ghg(t)
        F_s = forcing_solar(t)
        F_v = forcing_volcano(t)
        F_n = forcing_noise_func(t)

        F_total = F_g + F_s + F_v + F_n

        dT_dt = (F_total - lam * T) / ocean_heat_capacity
        # solve_ivp expects an array-like of the same shape as y (length 1).
        return [dT_dt]

    # Initial Condition: Start at equilibrium (0 anomaly)
    y0 = [0.0]

    sol = solve_ivp(
        system_dynamics, t_span=(0, n_years), y0=y0, t_eval=t_eval, method="RK45"
    )
    if not sol.success:
        raise RuntimeError(f"EBM ODE integration failed: {sol.message}")

    # --------------------------------------------------------------------------
    # 4. Packaging Results
    # --------------------------------------------------------------------------

    # Reconstruct forcing components for the output dataset (all four
    # forcing functions are vectorised over t).
    f_ghg = forcing_ghg(sol.t)
    f_volc = forcing_volcano(sol.t)
    f_solar = forcing_solar(sol.t)
    f_noise = forcing_noise_func(sol.t)
    f_total = f_ghg + f_volc + f_solar + f_noise

    years_abs = start_year + sol.t

    ds = xr.Dataset(
        coords={"time": years_abs},
        data_vars={
            "gmst": (("time",), sol.y[0]),
            "forcing_total": (("time",), f_total),
            "forcing_ghg": (("time",), f_ghg),
            "forcing_volcanic": (("time",), f_volc),
            "forcing_solar": (("time",), f_solar),
            "forcing_stochastic": (("time",), f_noise),
        },
        attrs={
            "description": "Zero-dimensional Energy Balance Model (EBM)",
            "equation": "C * dT/dt = F(t) - lambda * T",
            "climate_sensitivity": f"{climate_sensitivity} K per 2xCO2",
            "heat_capacity": f"{ocean_heat_capacity} W-yr/m2/K",
        },
    )

    # Downsample to annual means (monthly output grouped by integer year).
    # Build an explicit integer year coordinate so the resulting group
    # dimension is named "year" regardless of xarray's groupby-anonymous
    # naming behaviour.
    year_coord = np.floor(ds["time"].values).astype(int)
    ds_annual = ds.assign_coords(year=("time", year_coord)).groupby("year").mean()

    return ds_annual

compute_climate_signal(spatial_ds: xr.Dataset, gmst_da: xr.DataArray, base_val: float, coeffs: dict[str, float]) -> xr.DataArray

Combines Spatial and Temporal fields to create a mean climate signal (Mu).

Formula: Mu = Base + C1Elev + C2Lat + C3GMST + (C4GMST*Elev)

Source code in src/xtremax/simulations/extremes.py

def compute_climate_signal(
    spatial_ds: xr.Dataset,
    gmst_da: xr.DataArray,
    base_val: float,
    coeffs: dict[str, float],
) -> xr.DataArray:
    """
    Combines Spatial and Temporal fields to create a mean climate signal (Mu).

    Formula: Mu = Base + C1*Elev + C2*Lat + C3*GMST + (C4*GMST*Elev)
    """
    # Broadcast GMST to (year, site) and Spatial to (year, site)
    # Xarray handles broadcasting automatically by dimension name

    # 1. Static Effects
    elev_effect = coeffs.get("elevation", 0.0) * (spatial_ds["elevation"] / 1000.0)
    lat_effect = coeffs.get("lat", 0.0) * (spatial_ds["lat"] - spatial_ds["lat"].mean())

    # 2. Temporal Effects (The Signal)
    time_effect = coeffs.get("gmst", 0.0) * gmst_da

    # 3. Interaction Effects (e.g., higher elevations warm faster — elevation-dependent)
    interaction = (
        coeffs.get("interaction", 0.0) * gmst_da * (spatial_ds["elevation"] / 1000.0)
    )

    # Sum components
    mu = base_val + elev_effect + lat_effect + time_effect + interaction

    # Ensure correct dimension order
    return mu.transpose("year", "site")

compute_advanced_climate_signal(spatial_ds: xr.Dataset, gmst_da: xr.DataArray) -> xr.Dataset

Generates Tmax and Precip Mean fields using advanced physics proxies.

New Physics: 1. Continentality: Tmax increases with dist_to_coast. 2. Orographic Lift: Precip increases with Slope * Wind_Direction alignment. 3. Insolation: Tmax higher on South-facing slopes (Aspect ~ 180).

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

def compute_advanced_climate_signal(
    spatial_ds: xr.Dataset,
    gmst_da: xr.DataArray,
) -> xr.Dataset:
    """
    Generates Tmax and Precip Mean fields using advanced physics proxies.

    New Physics:
    1. Continentality: Tmax increases with dist_to_coast.
    2. Orographic Lift: Precip increases with Slope * Wind_Direction alignment.
    3. Insolation: Tmax higher on South-facing slopes (Aspect ~ 180).
    """
    # Create empty container

    # --- 1. Temperature Mean (Mu) ---
    # Base: 30C
    # Lapse Rate: -6.5 C / km
    # Continentality: +1.5 C per 100km from coast (Summer effect)
    # Latitude: -0.5 C per degree North
    # Aspect: +1.0 C if South Facing (180 deg), -1.0 if North
    # GMST: +1.5 * GMST

    aspect_factor = -np.cos(
        np.radians(spatial_ds["aspect"])
    )  # +1 at 180 (S), -1 at 0 (N)

    mu_temp_space = (
        30.0
        - 6.5 * (spatial_ds["elevation"] / 1000.0)
        + 0.015 * spatial_ds["dist_to_coast"]
        - 0.5 * (spatial_ds["lat"] - 36)
        + 1.0 * aspect_factor
    )

    # Combine with Time (GMST)
    # Continentality effect may *intensify* with GMST (land warms faster than ocean).
    land_ocean_contrast = 1.0 + (0.2 * spatial_ds["dist_to_coast"] / 100.0)

    mu_temp = mu_temp_space + (gmst_da * 1.5 * land_ocean_contrast)

    # --- 2. Precipitation Intensity Scale (Mu_Precip) ---
    # Base: 30mm
    # Elevation: Increase with height
    # Orographic: Increase if Slope is high (simple proxy)
    # Drying Trend: -5% per degree GMST

    mu_precip_space = (
        30.0
        + 0.02 * spatial_ds["elevation"]
        + 0.5 * spatial_ds["slope"]  # Steeper slopes = more orographic rain
    )

    # Thermodynamic scaling (Clausius Clapeyron + Dynamics)
    # Warmer = more intensity (CC), but circulation changes (dynamics) might reduce it.
    precip_scaling = 1.0 + 0.06 * gmst_da

    mu_precip = mu_precip_space * precip_scaling

    # Pack into Dataset
    ds_out = xr.Dataset()
    ds_out["mu_tmax"] = mu_temp.where(spatial_ds["mask"])
    ds_out["mu_precip"] = mu_precip.where(spatial_ds["mask"])

    return ds_out

Spatial domains & terrain

A worked spatial testbed: a procedural Iberian Peninsula with fractal Brownian-motion terrain, plus the feature extractors (elevation, distance to coast, and friends) that turn that geometry into covariates for spatial pooling.

Synthetic extreme value data generators.

Temporal trajectories (GMST, energy balance), spatial domain/terrain generators, and variable-specific extreme generators for temperature, precipitation, and wind.

SpatialFeatureExtractor

Static methods to derive physical features from a Digital Elevation Model (DEM).

get_distance_to_coast(ds: xr.Dataset) -> xr.DataArray staticmethod

Calculates Euclidean distance to the nearest coast (mask boundary). Uses SciPy's Distance Transform.

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

@staticmethod
def get_distance_to_coast(ds: xr.Dataset) -> xr.DataArray:
    """
    Calculates Euclidean distance to the nearest coast (mask boundary).
    Uses SciPy's Distance Transform.
    """
    # Transpose by name so datasets storing (lon, lat) — or any other
    # order — are processed and labeled correctly.
    mask = ds["mask"].transpose("lat", "lon").values
    # distance_transform_edt calculates distance to background (0)
    # We want distance from land (1) to ocean (0)
    dist_grid = distance_transform_edt(mask)

    # Convert pixels to approx km (assuming 0.1 deg ~ 11km)
    # This is an approximation, real geodesy would require Haversine
    res_deg = ds.lat[1] - ds.lat[0]
    dist_km = dist_grid * float(res_deg) * 111.0

    return xr.DataArray(
        dist_km, coords=ds.coords, dims=("lat", "lon"), name="dist_to_coast"
    )

get_slope_and_aspect(ds: xr.Dataset) -> tuple[xr.DataArray, xr.DataArray] staticmethod

Calculates local Slope (magnitude) and Aspect (direction). Slope is critical for orographic precipitation. Aspect is critical for solar insolation (South vs North facing).

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

@staticmethod
def get_slope_and_aspect(ds: xr.Dataset) -> tuple[xr.DataArray, xr.DataArray]:
    """
    Calculates local Slope (magnitude) and Aspect (direction).
    Slope is critical for orographic precipitation.
    Aspect is critical for solar insolation (South vs North facing).
    """
    # Transpose by name so the row axis is latitude regardless of the
    # dataset's stored dim order (the per-row cos(lat) scaling below
    # depends on it).
    elev = ds["elevation"].transpose("lat", "lon").values
    res_deg = float(ds.lat[1] - ds.lat[0])

    # Metric spacing differs per axis: one degree of longitude spans
    # 111 km · cos(latitude), so using the N–S spacing for both axes
    # overstates E–W gradients by ~30% at 40°N. Take gradients in
    # degree units, then rescale each axis by its own metres-per-degree
    # factor (per-row for longitude).
    grad_y_deg, grad_x_deg = np.gradient(elev, res_deg, res_deg)
    m_per_deg_lat = 111000.0
    m_per_deg_lon = m_per_deg_lat * np.cos(np.radians(ds["lat"].values))[:, None]
    grad_y = grad_y_deg / m_per_deg_lat
    grad_x = grad_x_deg / m_per_deg_lon

    # Slope (Magnitude) in degrees
    slope_rad = np.arctan(np.sqrt(grad_x**2 + grad_y**2))
    slope_deg = np.degrees(slope_rad)

    # Aspect (Direction)
    # 0 = North, 90 = East, 180 = South, 270 = West
    aspect_rad = np.arctan2(-grad_x, grad_y)  # Mathematical convention to compass
    aspect_deg = (np.degrees(aspect_rad) + 360) % 360

    da_slope = xr.DataArray(
        slope_deg, coords=ds.coords, dims=("lat", "lon"), name="slope"
    )
    da_aspect = xr.DataArray(
        aspect_deg, coords=ds.coords, dims=("lat", "lon"), name="aspect"
    )

    return da_slope, da_aspect

get_roughness(ds: xr.Dataset, window_size: int = 3) -> xr.DataArray staticmethod

Calculates Terrain Roughness Index (TRI). Roughness = Std Dev of elevation in a local window.

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

@staticmethod
def get_roughness(ds: xr.Dataset, window_size: int = 3) -> xr.DataArray:
    """
    Calculates Terrain Roughness Index (TRI).
    Roughness = Std Dev of elevation in a local window.
    """
    # Transpose by name so the output labels match the array layout
    # for any stored dim order.
    elev = ds["elevation"].transpose("lat", "lon").values
    # Uniform (boxcar) window of `window_size` pixels, matching the
    # TRI definition — previously this used a Gaussian filter with
    # sigma=window_size, i.e. a much wider, weighted window.
    local_mean = uniform_filter(elev, size=window_size)
    local_sq_mean = uniform_filter(elev**2, size=window_size)

    # Var = E[x^2] - (E[x])^2
    roughness = np.sqrt(np.maximum(0, local_sq_mean - local_mean**2))

    return xr.DataArray(
        roughness, coords=ds.coords, dims=("lat", "lon"), name="roughness"
    )

create_iberian_domain(res_deg: float = 0.1, bounds: tuple[float, float, float, float] = (-10, 5, 36, 44), seed: int = 42) -> xr.Dataset

Creates the master spatial dataset with Lat, Lon, Elevation, and Land Mask.

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

def create_iberian_domain(
    res_deg: float = 0.1,  # Grid resolution (approx 10km)
    bounds: tuple[float, float, float, float] = (-10, 5, 36, 44),
    seed: int = 42,
) -> xr.Dataset:
    """
    Creates the master spatial dataset with Lat, Lon, Elevation, and Land Mask.
    """
    lon_min, lon_max, lat_min, lat_max = bounds
    # `np.linspace(min, max, n)` puts points at both endpoints, giving
    # spacing `(max-min)/(n-1)` not `res_deg` (e.g. 36–44 at 0.1 would
    # yield ~0.1013). Use `np.arange` so realized spacing equals the
    # requested resolution exactly.
    lats = np.arange(lat_min, lat_max, res_deg)
    lons = np.arange(lon_min, lon_max, res_deg)
    n_lat = lats.size
    n_lon = lons.size
    lon_grid, lat_grid = np.meshgrid(lons, lats)

    # 1. Generate Mask
    mask = generate_iberia_mask(lat_grid, lon_grid)

    # 2. Generate Topography (Fractal)
    raw_topo = generate_fractal_terrain((n_lat, n_lon), seed=seed)

    # 3. Shape Topography to Physical Geography
    # Bias towards Pyrenees (North) and Sierra Nevada (South)
    # Simple linear ramp + bumps
    pyrenees_bias = np.exp(-((lat_grid - 42.5) ** 2) / 0.5) * 1.5
    sierra_bias = np.exp(-((lat_grid - 37.0) ** 2 + (lon_grid + 3.0) ** 2) / 0.8) * 1.5
    central_bias = np.exp(-((lat_grid - 40.5) ** 2 + (lon_grid + 4.0) ** 2) / 2.0) * 0.5

    # Apply bias and clip to land
    elevation = (
        raw_topo * 1500
        + (pyrenees_bias * 1500)
        + (sierra_bias * 2000)
        + (central_bias * 800)
    )
    elevation = np.where(mask, elevation, -100)  # Ocean is -100m
    elevation = np.clip(elevation, -100, 3400)  # Max Mulhacen/Aneto height

    ds = xr.Dataset(
        coords={"lat": lats, "lon": lons},
        data_vars={
            "mask": (("lat", "lon"), mask),
            "elevation": (("lat", "lon"), elevation),
        },
    )
    return ds

generate_iberia_mask(lat_grid: np.ndarray, lon_grid: np.ndarray) -> np.ndarray

Procedurally generates a boolean mask approximating the Iberian Peninsula and Balearic Islands using geometrical primitives.

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

def generate_iberia_mask(lat_grid: np.ndarray, lon_grid: np.ndarray) -> np.ndarray:
    """
    Procedurally generates a boolean mask approximating the Iberian Peninsula
    and Balearic Islands using geometrical primitives.
    """
    # 1. Main Peninsula (Rough approximation using rotated ellipse/intersection)
    # Center approx (40, -4)
    dx = (lon_grid - (-3.5)) * 0.85
    dy = lat_grid - 40.0

    # Rotated squared Euclidean distance (simple implicit function)
    # Rotated slightly counter-clockwise
    angle = np.radians(10)
    dx_rot = dx * np.cos(angle) - dy * np.sin(angle)
    dy_rot = dx * np.sin(angle) + dy * np.cos(angle)

    # Super-ellipseish shape for the blocky peninsula
    peninsula = (np.abs(dx_rot / 5.5) ** 4 + np.abs(dy_rot / 3.5) ** 4) <= 1.0

    # 2. Add Balearic Islands (Mallorca, Menorca, Ibiza)
    # Simple circles
    def make_island(lat0, lon0, rad):
        return ((lat_grid - lat0) ** 2 + (lon_grid - lon0) ** 2) < rad**2

    mallorca = make_island(39.6, 2.9, 0.4)
    menorca = make_island(40.0, 4.1, 0.2)
    ibiza = make_island(38.9, 1.4, 0.2)

    # Combine
    land_mask = peninsula | mallorca | menorca | ibiza
    return land_mask

generate_fractal_terrain(shape: tuple[int, int], scale: float = 100.0, octaves: int = 4, persistence: float = 0.5, lacunarity: float = 2.0, seed: int = 42) -> np.ndarray

Generates coherent terrain using Fractal Brownian Motion (fBm). Instead of random elevations, this creates mountain ranges and valleys.

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

def generate_fractal_terrain(
    shape: tuple[int, int],
    scale: float = 100.0,
    octaves: int = 4,
    persistence: float = 0.5,
    lacunarity: float = 2.0,
    seed: int = 42,
) -> np.ndarray:
    """
    Generates coherent terrain using Fractal Brownian Motion (fBm).
    Instead of random elevations, this creates mountain ranges and valleys.
    """
    # Use a local RNG. `np.random.seed(...)` would mutate NumPy's global
    # state and leak reproducibility coupling into any downstream code
    # that draws from NumPy afterwards.
    rng = np.random.default_rng(seed)
    terrain = np.zeros(shape)

    # Add layers of noise (octaves)
    for i in range(octaves):
        freq = lacunarity**i
        amp = persistence**i

        # Generate random noise grid
        noise = rng.normal(0, 1, shape)

        # Smooth it to create correlation (Cheap Perlin proxy via Gaussian)
        # Higher frequency = smaller sigma (rougher features)
        sigma = (shape[0] / freq) / 2.0
        smoothed = gaussian_filter(noise, sigma=sigma)

        # Normalize and add
        smoothed = (smoothed - smoothed.mean()) / (smoothed.std() + 1e-6)
        terrain += smoothed * amp

    # Normalize to 0-1
    terrain = (terrain - terrain.min()) / (terrain.max() - terrain.min())
    return terrain

generate_spatial_field(n_sites: int = 158, bounds: tuple[float, float, float, float] = (-9, 3, 36, 44), seed: int = 42) -> xr.Dataset

Generates static spatial data: Lat, Lon, Elevation, Distance to Coast.

Parameters:

Name	Type	Description	Default
bounds	tuple[float, float, float, float]	(lon_min, lon_max, lat_min, lat_max)	(-9, 3, 36, 44)

Source code in src/xtremax/simulations/extremes.py

def generate_spatial_field(
    n_sites: int = 158,
    bounds: tuple[float, float, float, float] = (-9, 3, 36, 44),  # Spain Box
    seed: int = 42,
) -> xr.Dataset:
    """
    Generates static spatial data: Lat, Lon, Elevation, Distance to Coast.

    Args:
        bounds: (lon_min, lon_max, lat_min, lat_max)
    """
    rng = np.random.default_rng(seed)
    lon_min, lon_max, lat_min, lat_max = bounds

    # Random spatial locations
    lon = rng.uniform(lon_min, lon_max, n_sites)
    lat = rng.uniform(lat_min, lat_max, n_sites)

    # Synthetic Elevation (correlated with distance from center for 'mountain' effect)
    # Simple heuristic: higher in the north/center
    center_lat = (lat_min + lat_max) / 2
    dist_factor = np.abs(lat - center_lat)
    elevation = rng.exponential(500, n_sites) + (dist_factor * 200)
    elevation = np.clip(elevation, 0, 3400)  # Clip to realistic max (Mulhacen approx)

    ds = xr.Dataset(
        coords={"site": np.arange(n_sites)},
        data_vars={
            "lon": (("site",), lon),
            "lat": (("site",), lat),
            "elevation": (("site",), elevation),
        },
    )
    return ds

augment_spatial_features(ds: xr.Dataset) -> xr.Dataset

Pipeline to run all extractors and merge into Dataset.

Returns a new Dataset; the input is not modified.

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

def augment_spatial_features(ds: xr.Dataset) -> xr.Dataset:
    """Pipeline to run all extractors and merge into Dataset.

    Returns a new Dataset; the input is not modified.
    """
    # Shallow copy so assignments below don't mutate the caller's Dataset
    # (repo convention: pure functions, side effects explicit).
    ds = ds.copy()
    extractor = SpatialFeatureExtractor()

    ds["dist_to_coast"] = extractor.get_distance_to_coast(ds)
    ds["slope"], ds["aspect"] = extractor.get_slope_and_aspect(ds)
    ds["roughness"] = extractor.get_roughness(ds)

    # Mask out ocean for clarity (optional, setting features to NaN)
    for var in ["dist_to_coast", "slope", "aspect", "roughness"]:
        ds[var] = ds[var].where(ds["mask"])

    return ds