Project: Vision Transformer (ViT) for Spatial Heatwave Downscaling and Calibration

This project demonstrates a key competency for the "AI in S2S Research Contractor" position: Using Vision Transformers (ViTs) to correct and downscale spatial climate fields (e.g., temperature grids during extreme heatwaves).

We implement a custom PyTorch ViT that takes coarse, biased forecast grids (8x8) and downscales them to high-resolution calibrated grids (32x32) matching ERA5 observations.

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
from torch.utils.data import DataLoader, TensorDataset

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

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

1. Synthetic Spatial Data Generation

We simulate spatial daily temperature fields over a regional domain:

• Truth (ERA5 Mock): $32 \times 32$ grid. Includes a spatial gradient (latitude effect), topography features (elevation-related cooling), and random heatwave anomalies (Gaussian hot-spots) peaking during summer.
• Raw Forecast (CFSv2/NASA Mock): Coarse $8 \times 8$ grid. Created by spatially averaging the truth grid, then applying a systematic negative bias (cold bias), underestimating peak heatwave temperatures (smoothing anomalies), and adding coarse noise.

# Domain configuration
GRID_SIZE_HR = 32
GRID_SIZE_LR = 8
NUM_DAYS = 500

def generate_spatial_data(num_days, hr_size, lr_size):
    x = np.linspace(0, 1, hr_size)
    y = np.linspace(0, 1, hr_size)
    X, Y = np.meshgrid(x, y)
    
    # 1. Geographic effect: Warm in south, cool in north + topographic ridge
    topo = 1.5 * np.exp(-((X - 0.5)**2 + (Y - 0.5)**2) / 0.1) # central ridge (cooling)
    base_field = 25 - 10 * Y - topo # Temp gradient
    
    hr_data = []
    lr_data = []
    
    for day in range(num_days):
        # Seasonal cycle
        season = 8 * np.sin(2 * np.pi * day / 365)
        day_base = base_field + season
        
        # Heatwave anomalies (1-3 Gaussian blobs)
        anomalies = np.zeros((hr_size, hr_size))
        if day % 30 < 10: # simulate periodic heatwaves (summer-like pulses)
            num_blobs = np.random.randint(1, 4)
            for _ in range(num_blobs):
                cx, cy = np.random.uniform(0.2, 0.8, 2)
                intensity = np.random.uniform(5, 12) # high intensity hot spot
                scale = np.random.uniform(0.1, 0.25)
                anomalies += intensity * np.exp(-((X - cx)**2 + (Y - cy)**2) / (2 * scale**2))
        
        # Fine-scale daily noise
        noise_hr = np.random.randn(hr_size, hr_size) * 0.5
        truth_day = day_base + anomalies + noise_hr
        
        # Generate biased coarse forecast (8x8)
        # Average into 8x8 blocks
        hr_tensor = torch.tensor(truth_day).unsqueeze(0).unsqueeze(0)
        lr_tensor = nn.functional.avg_pool2d(hr_tensor, kernel_size=hr_size // lr_size)
        lr_np = lr_tensor.squeeze().numpy()
        
        # Apply forecast bias: cold bias (-3 degrees) and underestimate heatwave amplitude (multiplied by 0.7)
        forecast_day = (lr_np - 3.0) + (lr_np - 15.0) * -0.2 # non-linear bias
        # Add coarse forecast noise
        forecast_day += np.random.randn(lr_size, lr_size) * 0.8
        
        hr_data.append(truth_day)
        lr_data.append(forecast_day)
        
    return np.array(lr_data, dtype=np.float32), np.array(hr_data, dtype=np.float32)

lr_forecasts, hr_truths = generate_spatial_data(NUM_DAYS, GRID_SIZE_HR, GRID_SIZE_LR)
print(f"Generated {NUM_DAYS} days of forecast {lr_forecasts.shape} and truth {hr_truths.shape}")

Generated 500 days of forecast (500, 8, 8) and truth (500, 32, 32)

2. Data Preprocessing for ViT

We downscale the coarse forecast to $32 \times 32$ using bilinear interpolation as our initial guess (raw baseline) and feed it into the Vision Transformer. The model's task is to predict a high-resolution residual map to correct the interpolated forecast.

# Convert to tensors
lr_tensor = torch.tensor(lr_forecasts).unsqueeze(1) # (N, 1, 8, 8)
hr_tensor = torch.tensor(hr_truths).unsqueeze(1)    # (N, 1, 32, 32)

# Upsample LR forecasts to 32x32 using bilinear interpolation (Raw baseline)
raw_baselines = nn.functional.interpolate(lr_tensor, size=(GRID_SIZE_HR, GRID_SIZE_HR), mode='bilinear', align_corners=True)

# Train/Test Split (80/20)
split = int(NUM_DAYS * 0.8)
X_train, y_train = raw_baselines[:split], hr_tensor[:split]
X_test, y_test = raw_baselines[split:], hr_tensor[split:]

train_dataset = TensorDataset(X_train, y_train)
test_dataset = TensorDataset(X_test, y_test)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

3. Vision Transformer (ViT) Architecture

We build a custom patch-based Vision Transformer for spatial grid calibration:

• Patch Extraction: Splits the $32 \times 32$ upsampled forecast into 64 patches of size $4 \times 4$.
• Linear Projection: Flattens patches (16 channels) and projects them to d_model dimension (32).
• Transformer Encoder: 2 blocks of Multi-Head Self-Attention (nhead=4, dim_feedforward=64) to capture long-range spatial correlations across the entire regional domain.
• Reconstruction Decoder: Linear projection back to flat patch sizes, followed by a fold operation to reconstruct the $32 \times 32$ residual correction map.
• Residual Connection: Adds the output of the transformer decoder back to the input upsampled forecast.

class ClimateViTDownscaler(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_channels=1, d_model=32, nhead=4, num_layers=2):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        patch_dim = in_channels * patch_size * patch_size
        
        # Patch projection
        self.patch_proj = nn.Linear(patch_dim, d_model)
        
        # Position embedding
        self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, d_model))
        
        # Transformer Encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 2, dropout=0.1, batch_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # Decoder projection
        self.decoder_proj = nn.Linear(d_model, patch_dim)
        
    def forward(self, x):
        # x: (batch, 1, img_size, img_size)
        batch_size = x.shape[0]
        
        # 1. Extract patches
        # Reshape to (batch, 1, num_patches_y, patch_size, num_patches_x, patch_size)
        p = self.patch_size
        h_p = self.img_size // p
        w_p = self.img_size // p
        
        # Extract non-overlapping patches
        patches = x.unfold(2, p, p).unfold(3, p, p) # (batch, 1, h_p, w_p, p, p)
        patches = patches.contiguous().view(batch_size, self.num_patches, -1) # (batch, num_patches, p*p)
        
        # 2. Project patches + position embedding
        x_proj = self.patch_proj(patches) + self.pos_embed
        
        # 3. Apply Transformer Encoder
        x_enc = self.transformer_encoder(x_proj)
        
        # 4. Decode patches
        x_dec = self.decoder_proj(x_enc) # (batch, num_patches, p*p)
        
        # 5. Reconstruct grid (Fold)
        x_dec = x_dec.view(batch_size, h_p, w_p, p, p)
        x_dec = x_dec.permute(0, 1, 3, 2, 4).contiguous() # (batch, h_p, p, w_p, p)
        residual = x_dec.view(batch_size, 1, self.img_size, self.img_size)
        
        # Final output is input + residual correction
        return x + residual

model = ClimateViTDownscaler()
print("Climate ViT Downscaler model initialized successfully.")

Climate ViT Downscaler model initialized successfully.

4. Model Training

We train the ViT downscaler using Mean Squared Error (MSE) loss against the high-resolution ground truth.

criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)

EPOCHS = 30
train_losses = []

for epoch in range(EPOCHS):
    model.train()
    batch_losses = []
    for batch_X, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        batch_losses.append(loss.item())
        
    avg_loss = np.mean(batch_losses)
    train_losses.append(avg_loss)
    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1}/{EPOCHS} - Loss: {avg_loss:.5f}")

# Plot and save training loss
plt.figure(figsize=(8, 4))
plt.plot(train_losses, label='ViT Training Loss', color='darkblue', linewidth=2)
plt.title("Vision Transformer Downscaling Training Loss")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.legend()
plt.grid(True)
plt.savefig("assets/loss.png")
plt.show()

Epoch 5/30 - Loss: 6.58749

Epoch 10/30 - Loss: 2.74546

Epoch 15/30 - Loss: 1.56045

Epoch 20/30 - Loss: 1.19593

Epoch 25/30 - Loss: 1.08767

Epoch 30/30 - Loss: 1.02767

5. Evaluation and Spatial Visualizations

We run evaluation on the test set and plot spatial temperature maps comparing:

1. Coarse Raw Forecast (upsampled)
2. Vision Transformer Calibrated Forecast
3. High-resolution Observations (ERA5 Truth)

model.eval()
test_preds = []
test_actuals = []
test_raws = []

with torch.no_grad():
    for batch_X, batch_y in test_loader:
        preds = model(batch_X)
        test_preds.append(preds.numpy())
        test_actuals.append(batch_y.numpy())
        test_raws.append(batch_X.numpy())

test_preds = np.concatenate(test_preds)
test_actuals = np.concatenate(test_actuals)
test_raws = np.concatenate(test_raws)

# Calculate Spatial RMSE
def spatial_rmse(truth, pred):
    return np.sqrt(np.mean((truth - pred) ** 2))

rmse_raw = spatial_rmse(test_actuals, test_raws)
rmse_vit = spatial_rmse(test_actuals, test_preds)

print(f"Overall Test Spatial RMSE (Raw Interpolated): {rmse_raw:.4f}°C")
print(f"Overall Test Spatial RMSE (ViT Calibrated): {rmse_vit:.4f}°C")

Overall Test Spatial RMSE (Raw Interpolated): 5.5728°C
Overall Test Spatial RMSE (ViT Calibrated): 1.1877°C

Spatial Field Comparison (Heatwave Event Day)

Let's inspect a day in the test set containing a strong heatwave anomaly.

# Find a test day with high heatwave anomaly peak temperature
heatwave_day_idx = np.argmax([test_actuals[i].max() for i in range(len(test_actuals))])

fig, axes = plt.subplots(1, 3, figsize=(18, 5))
vmin = min(test_actuals[heatwave_day_idx].min(), test_raws[heatwave_day_idx].min())
vmax = max(test_actuals[heatwave_day_idx].max(), test_raws[heatwave_day_idx].max())

im0 = axes[0].imshow(test_raws[heatwave_day_idx, 0], cmap='YlOrRd', vmin=vmin, vmax=vmax)
axes[0].set_title(f"Coarse Raw Forecast (Bilinear)\nRMSE = {spatial_rmse(test_actuals[heatwave_day_idx], test_raws[heatwave_day_idx]):.2f}°C")
axes[0].axis('off')

im1 = axes[1].imshow(test_preds[heatwave_day_idx, 0], cmap='YlOrRd', vmin=vmin, vmax=vmax)
axes[1].set_title(f"ViT Calibrated Forecast\nRMSE = {spatial_rmse(test_actuals[heatwave_day_idx], test_preds[heatwave_day_idx]):.2f}°C")
axes[1].axis('off')

im2 = axes[2].imshow(test_actuals[heatwave_day_idx, 0], cmap='YlOrRd', vmin=vmin, vmax=vmax)
axes[2].set_title("Ground Truth (ERA5 Mock)")
axes[2].axis('off')

fig.colorbar(im2, ax=axes.ravel().tolist(), label="Temperature (°C)", orientation='horizontal', shrink=0.6, pad=0.1)
plt.suptitle(f"Spatial Comparison of Heatwave Event downscaling (Test Day {heatwave_day_idx})")
plt.savefig("assets/spatial_comparison.png", bbox_inches='tight')
plt.show()

Error distribution and correlation plots

# Flatten predictions and actuals for pixel-level analysis
flat_actuals = test_actuals.flatten()
flat_raws = test_raws.flatten()
flat_preds = test_preds.flatten()

# 1. Error distribution (KDE)
plt.figure(figsize=(10, 5))
sns.kdeplot(flat_raws - flat_actuals, label="Raw Interpolated Error", fill=True, color="red", alpha=0.3)
sns.kdeplot(flat_preds - flat_actuals, label="ViT Calibrated Error", fill=True, color="green", alpha=0.4)
plt.axvline(0, color='black', linestyle='--')
plt.title("Pixel-Level Temperature Error Distribution (Forecast - Truth)")
plt.xlabel("Error (°C)")
plt.ylabel("Density")
plt.legend()
plt.grid(True)
plt.savefig("assets/error_dist.png")
plt.show()

# 2. Scatter plot (Actual vs Predicted)
plt.figure(figsize=(8, 8))
plt.scatter(flat_actuals[::10], flat_raws[::10], alpha=0.1, label="Raw Forecast", color="red")
plt.scatter(flat_actuals[::10], flat_preds[::10], alpha=0.1, label="ViT Calibrated", color="green")
min_val = min(flat_actuals.min(), flat_raws.min(), flat_preds.min())
max_val = max(flat_actuals.max(), flat_raws.max(), flat_preds.max())
plt.plot([min_val, max_val], [min_val, max_val], color='black', linestyle='--', label="Ideal (1:1)")
plt.title("Pixel-Level Actual vs Forecasted Temperature")
plt.xlabel("Actual Temperature (°C)")
plt.ylabel("Forecasted Temperature (°C)")
plt.legend()
plt.grid(True)
plt.savefig("assets/scatter_comparison.png")
plt.show()

Conclusion

Vision Transformers are capable of performing spatial downscaling and correction of systematic forecast biases. By splitting regional climate fields into patches and processing them using self-attention, the model learns physical relationships (such as geographic gradients and topographical cooling) and resolves localized extreme heatwaves.