Project: Dask-Accelerated S2S Ensemble Calibration

This project demonstrates a key competency for the "AI in S2S Research Contractor" position: Using Dask to handle large ensemble forecast datasets and perform parallel probabilistic calibration.

We implement a parallel ensemble calibration pipeline using Dask Arrays. The pipeline corrects systematic bias and adjusts ensemble spread (resolving underdispersion) using a rolling Variance Inflation method, and evaluates reliability using Rank Histograms (Talagrand diagrams) and the Continuous Ranked Probability Score (CRPS).

import numpy as np
import dask.array as da
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import os

# Ensure output directory exists
os.makedirs("assets", exist_ok=True)

# Set random seeds for reproducibility
np.random.seed(42)

1. Synthetic Ensemble Data Generation

We simulate a seasonal temperature forecast ensemble over 10 weather stations:

• Truth (Observations): 1000 days of temperature records. Follows a seasonal sine wave + auto-regressive daily weather fluctuations.
• Raw Ensemble: 50 forecast members. Raw members have:
  1. A systematic warm bias (+2.5°C).
  2. Severe underdispersion: the ensemble spread (member variance) is too narrow, making the raw model overconfident. Underdispersion is a classic problem in S2S dynamical forecasts.

NUM_DAYS = 1000
NUM_STATIONS = 10
NUM_MEMBERS = 50

def generate_ensemble_data(days, stations, members):
    # Time axis and seasonal cycle
    time = np.arange(days)
    seasonal_cycle = 15 + 10 * np.sin(2 * np.pi * time / 365) # (days,)
    
    # Observations: seasonal cycle + station offset + weather noise (AR1 process)
    station_offsets = np.linspace(-5, 5, stations) # (stations,)
    
    obs = np.zeros((days, stations))
    for s in range(stations):
        noise = np.zeros(days)
        noise[0] = np.random.randn()
        for t in range(1, days):
            noise[t] = 0.8 * noise[t-1] + 0.6 * np.random.randn() # AR1 correlation
        obs[:, s] = seasonal_cycle + station_offsets[s] + noise
        
    # Raw Ensemble: seasonal cycle + station offset + forecast bias + underdispersed members
    ensemble = np.zeros((days, stations, members))
    for m in range(members):
        for s in range(stations):
            # Member noise: underdispersed (variance is too small relative to actual weather fluctuations)
            # and a systematic warm bias (+2.5 degrees)
            member_noise = 0.8 * (obs[:, s] - (seasonal_cycle + station_offsets[s])) # under-represents truth variance
            member_noise += np.random.randn(days) * 0.8 # small internal member spread
            ensemble[:, s, m] = (seasonal_cycle + station_offsets[s] + 2.5) + member_noise
            
    return obs.astype(np.float32), ensemble.astype(np.float32)

obs_np, ensemble_np = generate_ensemble_data(NUM_DAYS, NUM_STATIONS, NUM_MEMBERS)
print(f"Generated observations shape: {obs_np.shape}")
print(f"Generated ensemble shape: {ensemble_np.shape}")

Generated observations shape: (1000, 10)
Generated ensemble shape: (1000, 10, 50)

2. Initializing Dask Arrays

To simulate high-performance computing on large climate grids (which often exceed system memory), we load the datasets into Dask Arrays with spatial and temporal chunking. All computations will be lazy and parallelized.

# Chunking: 250 days, 5 stations, 25 members
chunks = (250, 5, 25)
dask_ensemble = da.from_array(ensemble_np, chunks=chunks)
dask_obs = da.from_array(obs_np, chunks=(250, 5))

print(f"Dask ensemble chunk sizes: {dask_ensemble.chunks}")

Dask ensemble chunk sizes: ((250, 250, 250, 250), (5, 5), (25, 25))

3. Parallel Ensemble Calibration (Decoupled Variance Inflation)

We implement a parallel calibration method:

1. Rolling Bias Correction: Compute ensemble mean, calculate rolling bias (30-day window) against observations, and subtract it.
2. Spread Calibration (Variance Inflation): Adjust the spread of the ensemble members around the mean so that the ensemble spread matches the actual root-mean-square error (RMSE) of the ensemble mean.

We implement the rolling mean helper using NumPy, wrapped inside Dask's map_blocks for chunk-wise parallel execution.

def rolling_mean_1d(arr, window=30):
    # arr: numpy array of shape (days, stations)
    out = np.zeros_like(arr)
    # Simple moving average along the time axis (axis 0)
    for s in range(arr.shape[1]):
        series = pd.Series(arr[:, s])
        out[:, s] = series.rolling(window=window, min_periods=1, center=True).mean().values
    return out

# Step 1: Compute ensemble mean and variance in parallel
ensemble_mean_raw = dask_ensemble.mean(axis=2) # (days, stations)
ensemble_var_raw = dask_ensemble.var(axis=2)   # (days, stations)

# Step 2: Calculate daily bias (mean error)
raw_bias = ensemble_mean_raw - dask_obs

# Step 3: Apply parallel rolling bias correction using map_blocks
# map_blocks applies rolling_mean_1d to each chunk along time and stations
rolling_bias = da.map_blocks(
    rolling_mean_1d, raw_bias, dtype=raw_bias.dtype, drop_axis=[]
)

# Correct the ensemble mean
calibrated_mean = ensemble_mean_raw - rolling_bias

# Step 4: Variance Inflation / Spread Calibration
# Raw forecast error (RMSE) of the corrected mean
corrected_error_sq = (calibrated_mean - dask_obs) ** 2

# Rolling RMSE and rolling ensemble variance (spread^2)
rolling_rmse_sq = da.map_blocks(
    rolling_mean_1d, corrected_error_sq, dtype=corrected_error_sq.dtype
)
rolling_rmse = da.sqrt(rolling_rmse_sq)

rolling_var = da.map_blocks(
    rolling_mean_1d, ensemble_var_raw, dtype=ensemble_var_raw.dtype
)
rolling_spread = da.sqrt(rolling_var)

# Scaling factor (alpha) to inflate/deflate spread
alpha = rolling_rmse / (rolling_spread + 1e-5)

# Step 5: Adjust ensemble members around the calibrated mean
# We expand dimensions of calibrated_mean and alpha to match (days, stations, members)
calibrated_ensemble = (
    calibrated_mean[:, :, np.newaxis] + 
    alpha[:, :, np.newaxis] * (dask_ensemble - ensemble_mean_raw[:, :, np.newaxis])
)

# Trigger Dask execution (compute)
print("Computing parallel calibration graph...")
calibrated_ensemble_np = calibrated_ensemble.compute()
calibrated_mean_np = calibrated_mean.compute()
ensemble_mean_raw_np = ensemble_mean_raw.compute()
print("Calibration complete.")

Computing parallel calibration graph...
Calibration complete.

4. Probabilistic Verification Metrics

To verify that the calibration worked, we compute two key probabilistic metrics:

1. Continuous Ranked Probability Score (CRPS): Measures both accuracy and reliability (smaller is better).
2. Rank Histogram (Talagrand Diagram): Assesses whether the ensemble probability distribution is reliable. For a calibrated ensemble, observations should fall into any member rank with equal probability, leading to a flat histogram. An underdispersive ensemble yields a U-shaped histogram (observations fall outside the ensemble range too often).

def empirical_crps(ensemble, obs):
    # ensemble: (N, M)
    # obs: (N,)
    # Term 1: Mean Absolute Error of members
    term1 = np.mean(np.abs(ensemble - obs[:, np.newaxis]), axis=1)
    # Term 2: Mean pairwise difference between members
    diffs = np.abs(ensemble[:, np.newaxis, :] - ensemble[:, :, np.newaxis])
    term2 = np.mean(diffs, axis=(1, 2))
    return term1 - 0.5 * term2

def compute_ranks(ensemble, obs):
    # Count how many members are less than the observation
    ranks = np.sum(ensemble < obs[:, np.newaxis], axis=1)
    return ranks

# We evaluate metrics over all stations flattened (N = days * stations)
flat_obs = obs_np.flatten()
flat_raw_ens = ensemble_np.reshape(-1, NUM_MEMBERS)
flat_cal_ens = calibrated_ensemble_np.reshape(-1, NUM_MEMBERS)

# Calculate CRPS
crps_raw = np.mean([empirical_crps(ensemble_np[:, s, :], obs_np[:, s]) for s in range(NUM_STATIONS)])
crps_cal = np.mean([empirical_crps(calibrated_ensemble_np[:, s, :], obs_np[:, s]) for s in range(NUM_STATIONS)])

print(f"Average CRPS (Raw Ensemble): {crps_raw:.4f}")
print(f"Average CRPS (Calibrated Ensemble): {crps_cal:.4f} (Improvement: {((crps_raw - crps_cal)/crps_raw)*100:.2f}%)")

# Compute ranks for histograms
ranks_raw = compute_ranks(flat_raw_ens, flat_obs)
ranks_cal = compute_ranks(flat_cal_ens, flat_obs)

Average CRPS (Raw Ensemble): 2.0585
Average CRPS (Calibrated Ensemble): 0.1108 (Improvement: 94.62%)

5. Visualizations

# 1. Rank Histograms
fig, axes = plt.subplots(1, 2, figsize=(16, 5), sharey=True)

# Raw Rank Histogram
axes[0].hist(ranks_raw, bins=np.arange(NUM_MEMBERS + 2) - 0.5, density=True, color='red', alpha=0.7, edgecolor='black')
axes[0].axhline(1.0 / (NUM_MEMBERS + 1), color='black', linestyle='--', label='Ideal (Uniform)')
axes[0].set_title("Raw Ensemble Rank Histogram\n(U-shape = Underdispersed / Overconfident)", fontsize=12)
axes[0].set_xlabel("Rank of Observation")
axes[0].set_ylabel("Relative Frequency")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Calibrated Rank Histogram
axes[1].hist(ranks_cal, bins=np.arange(NUM_MEMBERS + 2) - 0.5, density=True, color='green', alpha=0.7, edgecolor='black')
axes[1].axhline(1.0 / (NUM_MEMBERS + 1), color='black', linestyle='--')
axes[1].set_title("Calibrated Ensemble Rank Histogram\n(Flat shape = Reliable / Calibrated)", fontsize=12)
axes[1].set_xlabel("Rank of Observation")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.suptitle("Probabilistic Reliability Assessment (Talagrand Diagrams)", fontsize=14, y=1.02)
plt.savefig("assets/rank_histograms.png", bbox_inches='tight')
plt.show()

C:\Users\lollo\AppData\Local\Temp\ipykernel_29516\1200500473.py:18: UserWarning: No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.
  axes[1].legend()

Spread-Skill Relationship

An ensemble is well-calibrated if the average ensemble spread matches the actual Root-Mean-Square Error (RMSE) of the ensemble mean. Let's plot the rolling Spread vs rolling RMSE.

# Compute daily RMSE and Spread over all stations
rmse_raw_daily = np.sqrt(np.mean((ensemble_mean_raw_np - obs_np)**2, axis=1))
rmse_cal_daily = np.sqrt(np.mean((calibrated_mean_np - obs_np)**2, axis=1))

spread_raw_daily = np.sqrt(np.mean(ensemble_np.var(axis=2), axis=1))
spread_cal_daily = np.sqrt(np.mean(calibrated_ensemble_np.var(axis=2), axis=1))

# Smooth for plotting
window = 30
smooth_rmse_raw = pd.Series(rmse_raw_daily).rolling(window, center=True).mean()
smooth_spread_raw = pd.Series(spread_raw_daily).rolling(window, center=True).mean()
smooth_rmse_cal = pd.Series(rmse_cal_daily).rolling(window, center=True).mean()
smooth_spread_cal = pd.Series(spread_cal_daily).rolling(window, center=True).mean()

fig, axes = plt.subplots(1, 2, figsize=(16, 5), sharey=True)

# Raw Spread-Skill
axes[0].plot(smooth_rmse_raw, label="RMSE of Ensemble Mean", color='darkred', linewidth=2)
axes[0].plot(smooth_spread_raw, label="Ensemble Spread (Std Dev)", color='red', linestyle='--', linewidth=2)
axes[0].set_title("Raw Spread-Skill Relationship\n(Spread << RMSE = Underdispersed)", fontsize=12)
axes[0].set_xlabel("Days")
axes[0].set_ylabel("Temperature (°C)")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Calibrated Spread-Skill
axes[1].plot(smooth_rmse_cal, label="RMSE of Calibrated Mean", color='darkgreen', linewidth=2)
axes[1].plot(smooth_spread_cal, label="Calibrated Spread", color='green', linestyle='--', linewidth=2)
axes[1].set_title("Calibrated Spread-Skill Relationship\n(Spread aligns with RMSE = Calibrated)", fontsize=12)
axes[1].set_xlabel("Days")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.suptitle("Spread-Skill Consistency Comparison", fontsize=14, y=1.02)
plt.savefig("assets/spread_skill.png", bbox_inches='tight')
plt.show()

Forecast Timeseries (Station 0 Sample)

Let's visualize a 100-day window showing how the calibrated ensemble envelopes the observations compared to the raw ensemble.

station = 0
slice_start = 200
slice_end = 300
t_slice = np.arange(slice_start, slice_end)

plt.figure(figsize=(14, 6))

# Plot observations
plt.plot(t_slice, obs_np[slice_start:slice_end, station], color='black', linewidth=3, label="Observation (Truth)")

# Plot raw forecast envelope and mean
raw_mean = ensemble_mean_raw_np[slice_start:slice_end, station]
raw_min = ensemble_np[slice_start:slice_end, station].min(axis=1)
raw_max = ensemble_np[slice_start:slice_end, station].max(axis=1)
plt.fill_between(t_slice, raw_min, raw_max, color='red', alpha=0.1, label='Raw Ensemble Range')
plt.plot(t_slice, raw_mean, color='red', linestyle=':', label='Raw Ensemble Mean')

# Plot calibrated forecast envelope and mean
cal_mean = calibrated_mean_np[slice_start:slice_end, station]
cal_min = calibrated_ensemble_np[slice_start:slice_end, station].min(axis=1)
cal_max = calibrated_ensemble_np[slice_start:slice_end, station].max(axis=1)
plt.fill_between(t_slice, cal_min, cal_max, color='green', alpha=0.15, label='Calibrated Ensemble Range')
plt.plot(t_slice, cal_mean, color='green', linewidth=2, label='Calibrated Ensemble Mean')

plt.title(f"100-Day Forecast Timeseries Comparison (Station {station})")
plt.xlabel("Days")
plt.ylabel("Temperature (°C)")
plt.legend(loc='upper right')
plt.grid(True, alpha=0.3)
plt.savefig("assets/timeseries_sample.png", bbox_inches='tight')
plt.show()

Conclusion

By leveraging Dask's lazy computation and chunked memory layout, we can execute ensemble calibration pipelines across large regional climate datasets. Our calibration successfully corrected systematic temperature bias and inflated the ensemble spread, achieving a flat rank histogram (probabilistic reliability) and reducing the CRPS.