ai4sci-climate-emulation

AI for ScienceClimSimrigorous codebase

Description

Climate Physics Emulation: Neural Network Architecture

Research Question

Design an improved neural network architecture for emulating sub-grid atmospheric physics processes in climate models. Your architecture should achieve lower Normalized MSE (NMSE) than the default MLP baseline on the ClimSim low-resolution dataset.

Background

Global climate models divide the atmosphere into grid cells, but many critical physical processes (radiation, convection, cloud formation) occur at scales smaller than these grid cells. Traditionally, these sub-grid processes are approximated by parameterization schemes -- handcrafted physics-based approximations. Neural network emulators can learn these mappings from high-resolution simulation data, potentially improving both accuracy and computational efficiency.

ClimSim provides data from the E3SM multi-scale climate model, where each sample maps an atmospheric column state to the corresponding sub-grid physics tendencies computed by the high-resolution physics module.

Task

Modify the Custom model class (lines 86-118 in custom_emulator.py) to implement a better neural network architecture. The model must:

Accept input_dim and output_dim in __init__
Implement forward(x) where x has shape (batch_size, input_dim)
Return predictions of shape (batch_size, output_dim)

Interface

Input structure (556-dim vector per atmospheric column):

9 multi-level variables x 60 vertical levels = 540 features: temperature (state_t), specific humidity (state_q0001), cloud ice (state_q0002), cloud liquid (state_q0003), zonal wind (state_u), meridional wind (state_v), ozone (pbuf_ozone), methane (pbuf_CH4), nitrous oxide (pbuf_N2O)
16-17 single-level (surface/TOA) scalar variables: surface pressure, solar insolation, heat fluxes, wind stress, albedos, surface type fractions, snow depths

Output structure (368-dim vector):

6 multi-level tendency variables x 60 levels = 360 features: temperature tendency (ptend_t), humidity tendencies (ptend_q0001-q0003), wind tendencies (ptend_u, ptend_v)
8 single-level diagnostic outputs: net shortwave, longwave down, snow/rain precipitation, direct/diffuse solar

Evaluation

Primary metric: Normalized MSE (NMSE = MSE / Var(target), lower is better)
Secondary metrics: R-squared (R2, higher is better), RMSE
Training budgets: 10 epochs (fast convergence), 30 epochs (standard), 60 epochs (extended)
All three training budgets run in parallel; improvements should be consistent across all three

Code

custom_emulator.py

EditableRead-only

1"""Custom Climate Physics Emulator
2Trained on ClimSim low-resolution E3SM data to predict sub-grid physics tendencies.
3
4Input: 556-dim atmospheric state vector (V2 variables)
5Output: 368-dim sub-grid physics tendencies
6"""
7
8import math
9import os
10import time
11from dataclasses import dataclass
12
13import numpy as np
14import torch
15import torch.nn as nn

Additional context files (read-only):

ClimSim/climsim_utils/data_utils.py

Results

Model	Type	nmse short-10ep ↓	rmse short-10ep ↓	ml nmse short-10ep ↓	sl nmse short-10ep ↓	nmse medium-30ep ↓	rmse medium-30ep ↓	ml nmse medium-30ep ↓	sl nmse medium-30ep ↓	nmse long-60ep ↓	rmse long-60ep ↓	ml nmse long-60ep ↓	sl nmse long-60ep ↓	r2 short-10ep ↑	r2 medium-30ep ↑	r2 long-60ep ↑
cnn	baseline	0.437	0.504	0.447	0.092	0.379	0.463	0.388	0.065	0.359	0.448	0.368	0.057	0.563	0.621	0.641
ed	baseline	0.535	0.552	0.541	0.338	0.487	0.526	0.492	0.326	0.464	0.513	0.469	0.321	0.465	0.513	0.535
hsr	baseline	0.504	0.535	0.510	0.323	0.448	0.503	0.452	0.311	0.423	0.489	0.427	0.307	0.495	0.552	0.577
mlp	baseline	0.406	0.476	0.409	0.302	0.381	0.458	0.383	0.295	0.374	0.453	0.376	0.293	0.594	0.619	0.626
unet	baseline	0.349	0.442	0.347	0.393	0.342	0.432	0.341	0.385	0.343	0.433	0.342	0.385	0.651	0.658	0.657