NaiNUQ - Sea Ice Emulator

A PyTorch-based UNet model for rapid sea-ice forecasting in the Arctic. This model emulates high-fidelity sea ice simulations from NANUQ, enabling fast predictions of sea ice properties including volume, concentration, and velocity and snow volume.

Model Description

Model Architecture: UNet with optional Partial Convolution layers Input Resolution: 128×128 grid cells
Prediction Variables: Sea ice volume (SIT), concentration (SIC), velocity (SIU/SIV), snow volume (SNT) Forecast Horizon: Configurable (typically 10-120 timesteps) Temporal Resolution: 1, 6, 12, or 24 hours

Intended Use

• Rapid ensemble sea ice forecasting
• Climate model emulation
• Arctic weather prediction
• Research on sea ice dynamics
• Educational applications in polar science

Model Inputs

The model accepts multi-channel input tensors representing:

• Sea ice state variables (5 channels for sit, sic, siu, siv, snt)
• Atmospheric forcing (7 channels)
• Ocean forcing (optional, 5 channels)
• Subsurface fields (optional, 2 channels) instead of ocean surface currents

Total input channels: 17 or 12 (if no ocean inputs)

Input shape: (batch_size, in_channels, 128, 128)

Model Outputs

Predictions of sea ice variables at each forecast timestep:

• SIT (Sea Ice Thickness): meters
• SIC (Sea Ice Concentration): fraction (0-1)
• SIU/SIV (Sea Ice Velocity): m/s
• SNT (Snow Thickness): meters
• Ocean variables (optional): various physical units

Output shape: (batch_size, out_channels, 128, 128)

Training Data

The model was trained on:

• Domain: Arctic (NANUK1 grid configuration)
• Temporal coverage: Multiple years of coupled sea ice-ocean simulations
• Data format: TFRecord dataset (compatible conversion scripts provided)
• Normalization: Applied using per-variable mean and standard deviation

Quick start

Installation

pip install torch torchvision numpy

# Clone the repository
git clone https://github.com/nanuqhub/nainuq.git
cd nainuq

Basic Usage

import torch
from layers.full_UNet import UNetModel

# Load the model
model = UNetModel(
    in_channels=17,      # 7 atmosphere + 5 ocean + 5 sea ice
    out_channels=5,      # sit, sic, siu, siv, snt
    base_features=32
)

# Load weights
checkpoint = torch.load("checkpoint_epoch_100.pt", weights_only=True)
model.load_state_dict(checkpoint["state_dict"])
model.eval()

# Prepare input (batch_size=1, 17 channels, 128x128 grid)
x = torch.randn(1, 17, 128, 128)

# Make prediction
with torch.no_grad():
    output = model(x)  # Shape: (1, 5, 128, 128)

Advanced Usage with Inference Class

from inference.test_utils import Test
import argparse

# Configure inference
args = argparse.Namespace(
    sea_ice_variables=['sit', 'sic', 'siu', 'siv', 'snt'],
    use_ocean_as_forcings=True,
    ocean_under=True,
    k=120,              # forecast horizon (timesteps)
    timestep=1,         # temporal resolution (hours)
    n_cycle=100,        # number of forecast cycles
    frequency=24,       # offset between cycles
    post_processing=True,
    save_pred=True,
    noise=0.0,
    noise_init=True,
    ocean=True
)

# Initialize test class
test = Test(
    args=args,
    model=model,
    use_ocean_as_forcings=True,
    N_ocean=5,
    N_under=2,
    N_inputs=17,
    N_outputs=5,
    frequency=24,
    post_processing=True,
    ocean=True,
    season="all",
    k=120,
    N_cycle=100,
    save_pred=True,
    timestep=1,
    path_to_save="./results",
    path_to_data="./data",
    noise=0.0,
    noise_init=True
)

# Run inference
fs, fs_pers, bias, predictions, truth = test.test_model()

Configuration

The model configuration is stored in config.json:

{
  "in_channels": 17,
  "out_channels": 5,
  "base_features": 32,
  "lr": 1e-4,
  "weight_decay": 1e-3,
  "lambda_": 100.0
}

Performance Metrics

Evaluation Metrics

• RMSE (Root Mean Square Error): Quantifies prediction accuracy
• Bias: Systematic model error
• Forecast Skill (FS): Anomaly correlation compared to climatology
• Persistence Forecast Skill (FS_pers): Skill of persistence baseline

Typical Performance

Results vary by:

• Forecast horizon (skill decreases with time)
• Variable (some are more predictable than others)
• Season (ice season vs. melt season)
• Configuration (with/without ocean forcing)

For detailed metrics, see test results in the repository.

Post-Processing

Optional physical post-processing enforces:

• Minimum ice thickness (h_min = 1e-5 m)
• Minimum ice concentration (c_min = 1e-4)
• Physical constraints on output fields

Enable with --post_processing True

Limitations

• Fixed grid size (128×128)
• Trained on Arctic domain only
• Performance degrades for very long forecasts (>2 weeks)
• Requires normalized input data
• Ocean variables should be derived from coupled simulations

Hardware Requirements

• GPU (Recommended): NVIDIA GPU with 6GB+ VRAM (e.g., RTX 2070)
• CPU: ~30-60 seconds per forecast on modern CPU
• Memory: ~2GB RAM for single predictions

Known Issues

• Mask loading path is hardcoded (requires updating in test_utils.py)
• TFRecord data format requires specific preprocessing
• Ocean data access paths need configuration

Citation

If you use this model, please cite:

License

MIT - See LICENSE file for details

Contributing

Contributions are welcome! Please:

1. Fork the repository
2. Create a feature branch
3. Submit a pull request with clear descriptions

Support & Issues

For questions, bugs, or feature requests, please open an issue on GitHub: https://github.com/nanuqhub/nainuq/issues

Acknowledgments

• Arctic domain configuration based on NANUK1 grid
• Training data from coupled sea ice-ocean simulations
• UNet architecture adapted for multi-scale feature extraction

References

• U-Net: Convolutional Networks for Biomedical Image Segmentation (Ronneberger et al., 2015)
• Partial Convolutions (Liu et al., 2018)

---

Last Updated: 2024
Model Status: Active
Maintained By: Nanuq Hub