A65_DYAMOND_v7: Deep Learning Super-Resolution for Climate Model Output

Experiment Overview

This project implements a deep learning-based super-resolution model to upscale low-resolution climate model output to high-resolution using the DYAMOND (DYnamics of the Atmospheric general circulation Modeled On Non-hydrostatic Domains) winter dataset. The model performs 8× spatial upsampling from coarse-grid (1080×180) to fine-grid (8640×1440) resolution on a cubed-sphere grid geometry.

Key Objectives:

• Model atmospheric variables at high spatial resolution using limited computational resources
• Preserve physical consistency and sharp features in upsampled fields
• Enable efficient multi-node distributed training for production-scale deployment

---

Model Architecture

ProgressiveUNetv7

The model employs a ConvNeXt-based progressive upsampling architecture designed specifically for 8× super-resolution tasks. This architecture choice addresses common challenges in climate data upsampling, including edge artifacts, spatial discontinuities, and feature preservation across large scaling factors.

Architecture Components

1. Feature Extraction Stage (30×30)

• Stem Layer: 3×3 convolution with reflection padding → 96 channels
• Deep Feature Extraction: 3 ConvNeXt blocks at base resolution
• Reflection padding prevents boundary artifacts critical for global climate grids

2. Progressive Upsampling Stages

• Stage 1 (30→60): PixelShuffle 2× + 3 ConvNeXt blocks
• Stage 2 (60→120): PixelShuffle 2× + 9 ConvNeXt blocks
• Stage 3 (120→240): PixelShuffle 2× + 3 ConvNeXt blocks

3. Output Layer

• 3×3 convolution with reflection padding → 2 output channels

ConvNeXt Block Design

Each ConvNeXt block consists of:

1. Depthwise 7×7 convolution (reflection padding, groups=channels)
2. LayerNorm (channel-wise normalization)
3. Pointwise expansion (1×1 → 4× channels)
4. GELU activation
5. Pointwise projection (4× channels → 1×)
6. Layer scale (learnable scaling parameter)
7. Residual connection

Model Parameters: 44.19M

Architectural Rationale

1. ConvNeXt over Transformers:

   • Better inductive bias for spatial data
   • More parameter-efficient for high-resolution outputs
   • Faster training on limited GPU memory

2. Progressive Upsampling over Single-Step:

   • Prevents checkerboard artifacts common in large upsampling factors
   • Allows gradual feature refinement across scales
   • More stable training dynamics

3. Reflection Padding:

   • Critical for cubed-sphere grids to avoid edge discontinuities
   • Maintains physical continuity across patch boundaries
   • Reduces artifacts at domain edges

4. PixelShuffle over Transposed Convolution:

   • Avoids checkerboard artifacts inherent to strided transposed convolutions
   • More parameter-efficient
   • Better preserves high-frequency details

---

Input/Output Specification

Input Variables (6 Channels, 30×30 patches)

Variable	Description	Physical Unit	Notes
T2M	2-meter air temperature	K	Primary meteorological variable
SWGNET	Net shortwave radiation at ground	W/m²	Radiative forcing
PRECTOT	Total precipitation	kg/m²/s	Convective + stratiform
PHIS	Surface geopotential	m²/s²	Topography (static)
Spatial Encoding	2D coordinate encoding	normalized	Grid position awareness
Hour Encoding	Time-of-day encoding	[0,1]	Diurnal cycle information

Input Spatial Resolution: 1080×180 global grid (30×30 patches)

Output Variables (2 Channels, 240×240 patches)

Variable	Description	Physical Unit
T2M	2-meter air temperature	K
PRECTOT	Total precipitation	kg/m²/s

Output Spatial Resolution: 8640×1440 global grid (240×240 patches)

Upsampling Factor: 8× in each spatial dimension

Data Preprocessing

• Normalization: Per-variable z-score normalization using training set statistics
• Patch Strategy: 216 patches per timestep (36×6 grid decomposition)
• Grid Structure: Cubed-sphere geometry (6 faces)

---

Loss Function

Composite Multi-Objective Loss

The training employs a weighted combination of three loss components to balance pixel-wise accuracy, spatial gradients, and spectral fidelity:

L_total = λ_charb · L_charbonnier + λ_grad · L_gradient + λ_fft · L_spectral

Loss Weights: λ_charb = 1.0, λ_grad = 1.0, λ_fft = 0.02

1. Charbonnier Loss (Robust L1)

L_charb = mean(√((pred - target)² + ε²))

• Purpose: Pixel-wise accuracy with robustness to outliers
• ε = 1e-6: Smoothing parameter for numerical stability
• Advantage: Less sensitive to extreme values than L2, smoother than L1

2. Gradient Loss (Spatial Sharpness)

L_grad = L_charb(∇_x pred, ∇_x target) + L_charb(∇_y pred, ∇_y target)

• Purpose: Preserve sharp features and spatial gradients
• Rationale: Prevents over-smoothing common in super-resolution
• Implementation: First-order finite differences with Charbonnier metric

3. Spectral Loss (FFT-based Texture)

L_spectral = L1(|FFT(pred)|, |FFT(target)|)

• Purpose: Match high-frequency spectral content
• Weight: Small (0.02) to "nudge" texture without dominating training
• Rationale: Encourages realistic fine-scale structures

Loss Design Rationale

• Charbonnier: Provides robust baseline accuracy
• Gradient: Addresses super-resolution-specific challenge of feature sharpness
• Spectral: Ensures realistic high-frequency content without artifacts
• Weighting: Empirically tuned for balance between smoothness and detail

---

Training Configuration

Dataset: DYAMOND Winter

• Training Period: 2020-11-01 to 2021-01-14 (8,640 timesteps)
• Validation Period: 2021-01-15 to 2021-01-30 (360 timesteps)
• Temporal Resolution: 10-minute intervals
• Spatial Coverage: Global cubed-sphere grid

Distributed Training Setup

Multi-Node Configuration:

• Nodes: 2 (NASA Cabeus HPC)
• GPUs per Node: 4 × NVIDIA A100-SXM4-80GB
• Total GPUs: 8
• Batch Size: 64 per GPU (512 global batch size)
• Data Parallelism: DistributedDataParallel (DDP) with NCCL backend

Optimization:

• Optimizer: AdamW (lr=2e-4, weight_decay=1e-4)
• Scheduler: CosineAnnealingLR (T_max=50 epochs, η_min=1e-6)
• Gradient Clipping: max_norm=1.0
• Mixed Precision: FP16 with GradScaler
• Epochs: 50

Data Loading:

• Smart Batch Sampling: Custom sampler for efficient multi-timestep batching
• Distributed Sampling: Each GPU processes unique data subset
• Workers: 0 (single-threaded due to xarray/netCDF4 multiprocessing constraints)

Computational Details

• Patches per Timestep: 216 (36×6 decomposition)
• Training Batches per Epoch: ~25,920 (8,640 timesteps × 216 patches / 512 batch size / 8 GPUs)
• Model Parameters: 44.19M
• Peak GPU Memory: ~70 GB per GPU
• Estimated Training Time: ~5-8 hours per epoch on 8×A100

---

Scientific Rationale

Problem Motivation

High-resolution climate model simulations (e.g., km-scale) are computationally prohibitive for long-term climate projections. Statistical downscaling using deep learning offers a promising alternative:

1. Computational Efficiency: Train once on expensive simulations, deploy for cheap inference
2. Physical Consistency: Learn from physics-based models rather than pure statistical relationships
3. Scalability: Apply to arbitrary future scenarios without retraining

Why This Architecture?

ConvNeXt-based Progressive Upsampling addresses key challenges in climate super-resolution:

1. Large Upsampling Factor (8×):

   • Progressive stages prevent feature collapse
   • Gradual refinement improves training stability
   • Each stage learns at appropriate spatial scale

2. Global Grid Continuity:

   • Reflection padding preserves boundary conditions
   • Critical for cubed-sphere geometry
   • Prevents artifacts at face edges

3. Feature Preservation:

   • ConvNeXt blocks maintain large receptive fields (7×7 depthwise)
   • Residual connections prevent gradient degradation
   • Layer scaling enables deep networks (18 total blocks)

4. Physical Realism:

   • Multi-objective loss balances accuracy and sharpness
   • Spectral loss ensures realistic fine-scale structures
   • Gradient loss prevents over-smoothing

Comparison to Alternatives

Approach	Pros	Cons	Why Not Used
Bicubic Interpolation	Fast, simple	Over-smooths, no learned patterns	Lacks physical realism
Single-Step Upsampling	Fewer parameters	Checkerboard artifacts at 8×	Quality degradation
Transformer-based	Strong for global context	High memory cost, slower	Resource constraints
GAN-based	Sharp outputs	Training instability, mode collapse	Production reliability

Expected Outcomes

1. Spatial Resolution: 8× finer grid spacing
2. Feature Preservation: Sharp gradients in temperature/precipitation fronts
3. Physical Consistency: Realistic spectral energy distribution
4. Computational Gain: >100× speedup vs. high-res simulation

---

Usage

Single-Node Training

python A65_DYAMOND_v7.py

Multi-Node Distributed Training

qsub A65_DYAMOND_multinode.j

Key Files

• A65_DYAMOND_v7_multinode.py: Main training script with DDP support
• A65_DYAMOND_multinode.j: PBS job submission script for HPC
• README.md: This documentation

---

Citation

If you use this code or methodology in your research, please cite:

@article{a65_dyamond_v7,
  title={Deep Learning Super-Resolution for Climate Model Output: A ConvNeXt-based Approach},
  author={[Your Name]},
  journal={[Journal Name]},
  year={2025},
  note={DYAMOND Winter Dataset}
}

---

References

1. DYAMOND: Stevens et al. (2019). DYAMOND: The DYnamics of the Atmospheric general circulation Modeled On Non-hydrostatic Domains. JAMES.
2. ConvNeXt: Liu et al. (2022). A ConvNet for the 2020s. CVPR.
3. PixelShuffle: Shi et al. (2016). Real-Time Single Image and Video Super-Resolution. CVPR.
4. Climate Downscaling: Stengel et al. (2020). Adversarial super-resolution of climatological wind and solar data. PNAS.

---

Acknowledgments

• NASA High-End Computing (HEC) Program for computational resources
• DYAMOND initiative for high-resolution climate model data
• PyTorch Distributed Data Parallel for efficient multi-GPU training