Add O96 Mesh Analysis Output and Gridded Verification Against NWP Forecasts

[azadeh-gh]

Add O96 Mesh Analysis Output and Gridded Verification Against NWP Forecasts

Summary

Enable the model to save evolved mesh state (O96 grid) for gridded verification against traditional NWP model forecasts. Currently, the model only outputs predictions decoded to observation space, but the mesh state itself represents a gridded analysis that can be directly compared with operational forecast models.

Background

• The O96 icosahedral mesh (~40,962 nodes, ~100km resolution) represents the analyzed atmospheric state
• Current validation only compares decoded predictions against observations (irregular locations)
• For operational use and comprehensive validation, we need gridded comparisons on the O96 mesh
• This enables standard DA metrics (analysis-forecast differences, ACC, RMSE on grid)

Proposed Changes

1. Modify Model Forward Pass

Add capability to return mesh state before decoding

2. Save Mesh State During Validation

Capture and save mesh state for gridded analysis

3. Implement NWP Comparison on O96 Mesh

Add utilities to:

• Interpolate NWP gridded forecasts to O96 mesh nodes
• Compute gridded verification metrics (RMSE, bias, ACC)
• Compare GNN analysis vs NWP forecast on common O96 grid

4. Output Format

Save mesh analyses in netCDF format compatible with operational workflows

Benefits

1. Dual validation: Both observation-space and gridded metrics
2. Operational compatibility: Direct comparison with NWP models on same grid
3. Standard verification: Enables traditional DA metrics (O-A, O-F statistics)
4. Interpretability: Gridded output easier to visualize and understand than scattered obs

Implementation Tasks

• Modify forward() to optionally return mesh state
• Update validation_step() to capture mesh state
• Implement NWP interpolation to O96 mesh
• Add gridded metric computation
• Create netCDF output utility for mesh analyses
• Add gridded metrics to logging/tensorboard
• Update documentation with gridded validation workflow
• Add example notebook showing O96 mesh analysis vs NWP comparison

Mesh Resolution Details

The current model uses mesh_resolution=6, which creates an icosahedral mesh with the following properties:

Technical Specifications

• Number of nodes: 40,962 (calculated as 10 × 4^6 + 2)
• Spatial resolution: ~100-110 km average node spacing
• Global coverage: Quasi-uniform distribution over Earth's surface
• Area per node: ~12,450 km²

Resolution Comparison Table

Resolution	Splits	Nodes	Approx. Spacing	Use Case
4	4	2,562	~280 km	Coarse testing
5	5	10,242	~140 km	Medium resolution
6	6	40,962	~100 km	current
7	7	163,842	~50 km	High resolution
8	8	655,362	~25 km	Very high resolution

[mlaurako]

Update on this part:

1. Modify Model Forward Pass - Add capability to return mesh state before decoding

Problem:

The mesh state before decoding contains intermediate abstract features for all variables, which does NOT represent physical values of target variables (e.g., 288K, 5.2 m/s).

To capture physical values on a mesh grid, we considered the modified workflow proposed in Issue #38 (add projection layers during training). However, since we do not want to alter the training behavior, we've decided to instead:

• Keep training unchanged
• During validation/inference: decode mesh state to mesh grid locations to obtain physical predictions

Finalized workflow:

TRAINING (unchanged):
Mesh [abstract] → Decoder → Sparse Observations → Loss

VALIDATION/INFERENCE (new):
Mesh [abstract] → Decoder → Physical Predictions on Grid Locations (default: O96 mesh node location) → Save physical predictions on grid

Implementation will support:

• Decoding to O96 mesh node locations (default; using self-connections/1NN instead of 3NN)
• Optional: Decoding to custom grids (FV3, Gaussian, regional lat-lon locations, any resolution of icosahedral mesh)

[azadeh-gh]

@mlaurako - Thanks for the update!

Your proposed validation/inference workflow:
Mesh [abstract, 40,962, hidden_dim]
↓
Decoder (1NN self-connections to same mesh)
↓
Mesh [abstract, 40,962, hidden_dim] ← Still abstract!
↓
Output_mapper
↓
Physical Mesh [40,962, num_vars]

My concern: This creates a circular path where the mesh decodes to itself. The decoder was trained to transform mesh → observation locations (handling scan angles, spatial context, etc.), not mesh → same mesh. Using it for self-connections doesn't leverage what it learned during training and adds unnecessary computational overhead.

Alternative: Direct Mesh Projection
Instead of routing through the decoder, we can add a dedicated mesh-to-physical projection layer.

Mesh [abstract, 40,962, hidden_dim]
→ Mesh Projection (linear layer)
→ Physical Mesh [40,962, num_vars]
→ Save to grid