Upload folder using huggingface_hub

.gitattributes

@@ -1,14 +1,8 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
+*.hdf5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
 *.npy filter=lfs diff=lfs merge=lfs -text
@@ -16,20 +10,11 @@
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+weight/* filter=lfs diff=lfs merge=lfs -text
+model/data/era5_tl31_19590102T00.nc filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,204 @@
+---
+frameworks: JAX
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecasting
+- Climate Simulation
+- Hybrid Physics-ML
+- ERA5
+- NeuralGCM
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">NeuralGCM</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+NeuralGCM (Neural General Circulation Models) is an open-source hybrid machine-learning and physics-based atmospheric model developed by Google Research for weather forecasting and climate simulation.
+
+Paper: Neural General Circulation Models for Weather and Climate
+
+https://arxiv.org/abs/2311.07222
+
+# Model Description
+
+NeuralGCM is built around a differentiable atmospheric dynamical core. Neural networks represent unresolved physical processes, the encoder, and the decoder, improving forecast efficiency while retaining physical constraints.
+
+| Profile | Resolution | Type | Bundled official checkpoint |
+| :--- | :---: | :--- | :--- |
+| `weather_forecast` | 0.7 degrees (`512 x 256`) | Deterministic weather forecasting for approximately 2 to 15 days | `weight/models_v1_deterministic_0_7_deg.pkl` |
+| `climate_scale` | 1.4 degrees (`256 x 128`) | Deterministic climate-scale simulation | `weight/models_v1_deterministic_1_4_deg.pkl` |
+| `forecast_2_8_deg` | 2.8 degrees (`128 x 64`) | Deterministic weather forecasting | `weight/models_v1_deterministic_2_8_deg.pkl` |
+| `stochastic_1_4_deg` | 1.4 degrees (`256 x 128`) | Stochastic weather forecasting | `weight/models_v1_stochastic_1_4_deg.pkl` |
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather forecasting | Train the 0.7-degree model on ERA5 data for short- to medium-range weather forecasting. |
+| Climate-scale simulation | Train the 1.4-degree model on ERA5 data for longer atmospheric simulations. |
+| Low-resolution experiments | Use the 2.8-degree data profile for lower-cost weather forecasting experiments. |
+| Local quick validation | Generate HDF5 data with the required channel protocol using `scripts/fake_data.py` and validate the data, model, and checkpoint workflows. |
+| ModelScope / OneCode execution | Download the standalone model package, install the OneScience and JAX dependencies, and run the scripts directly. |
+| Multi-device training | Run synchronous data-parallel training on multiple local accelerators. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import checks and small-scale connectivity validation, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version compatible with the current cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/NeuralGCM --local-dir ./NeuralGCM
+cd NeuralGCM
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and conda first.
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported.
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate conda first.
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# Installation with uv is also supported.
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The OneScience community provides an ERA5 data slice for training. Download it with the following command and confirm that the data path in `conf/config.yaml` is correct:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+### Generate Synthetic Data
+
+```bash
+python scripts/fake_data.py
+```
+
+The script creates yearly HDF5 files under `data/data/`, writes synthetic static fields to `data/static.nc`, and saves channel, time-window, and grid metadata to `data/metadata/dataset_card.json`. The synthetic fields use approximate physical units but are intended only for shape, loading, regridding, and numerical-stability checks.
+
+### Training
+
+Single device:
+
+```bash
+# 0.7-degree deterministic short- to medium-range weather forecasting
+python scripts/train_weather_forecast.py
+# 1.4-degree deterministic climate-scale simulation
+python scripts/train_climate_scale.py
+# 2.8-degree deterministic low-resolution weather forecasting
+python scripts/train_forecast_2_8_deg.py
+# 1.4-degree stochastic weather forecasting
+python scripts/train_stochastic_1_4_deg.py
+```
+
+Multiple devices:
+
+```bash
+# 0.7-degree deterministic short- to medium-range weather forecasting
+python scripts/train_weather_forecast.py --devices 8
+# 1.4-degree deterministic climate-scale simulation
+python scripts/train_climate_scale.py --devices 8
+# 2.8-degree deterministic low-resolution weather forecasting
+python scripts/train_forecast_2_8_deg.py --devices 8
+# 1.4-degree stochastic weather forecasting
+python scripts/train_stochastic_1_4_deg.py --devices 8
+```
+
+### Fine-tuning
+
+Fine-tuning can start from either a checkpoint produced by local training or the bundled official checkpoint for the selected profile.
+
+```bash
+# Use the bundled official checkpoint for each profile.
+python scripts/train_weather_forecast.py --finetune weight/models_v1_deterministic_0_7_deg.pkl
+python scripts/train_climate_scale.py --finetune weight/models_v1_deterministic_1_4_deg.pkl
+python scripts/train_forecast_2_8_deg.py --finetune weight/models_v1_deterministic_2_8_deg.pkl
+python scripts/train_stochastic_1_4_deg.py --finetune weight/models_v1_stochastic_1_4_deg.pkl
+
+# Alternatively, provide a local checkpoint explicitly.
+python scripts/train_weather_forecast.py --finetune ./data/checkpoint/model_bak.pkl
+```
+
+For multi-device fine-tuning, add `--devices` to the corresponding command.
+
+### Pre-trained Weights
+
+This project includes the following official pre-trained checkpoints:
+
+| Local file | Official release path |
+| :--- | :--- |
+| `weight/models_v1_deterministic_0_7_deg.pkl` | `gs://neuralgcm/models/v1/deterministic_0_7_deg.pkl` |
+| `weight/models_v1_deterministic_1_4_deg.pkl` | `gs://neuralgcm/models/v1/deterministic_1_4_deg.pkl` |
+| `weight/models_v1_deterministic_2_8_deg.pkl` | `gs://neuralgcm/models/v1/deterministic_2_8_deg.pkl` |
+| `weight/models_v1_stochastic_1_4_deg.pkl` | `gs://neuralgcm/models/v1/stochastic_1_4_deg.pkl` |
+
+### Inference
+
+```bash
+# 0.7-degree deterministic short- to medium-range weather forecasting
+python scripts/inference.py --mode weather_forecast --checkpoint weight/models_v1_deterministic_0_7_deg.pkl
+# 1.4-degree deterministic climate-scale simulation
+python scripts/inference.py --mode climate_scale --checkpoint weight/models_v1_deterministic_1_4_deg.pkl
+# 2.8-degree deterministic low-resolution weather forecasting
+python scripts/inference.py --mode forecast_2_8_deg --checkpoint weight/models_v1_deterministic_2_8_deg.pkl
+# 1.4-degree stochastic weather forecasting
+python scripts/inference.py --mode stochastic_1_4_deg --checkpoint weight/models_v1_stochastic_1_4_deg.pkl
+```
+
+Without an explicit `--checkpoint`, inference first checks `./data/checkpoint/model_bak.pkl`. The default output is `results/predictions.nc`, containing pressure-level variables with their official names and rollout time coordinates.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is a reproduction of the original NeuralGCM paper.
+- The repository code is provided under the Apache License 2.0.
+- The trained model weights released by Google, including the four checkpoints in this directory, are licensed under the Creative Commons Attribution-ShareAlike 4.0 International license (CC BY-SA 4.0). Redistribution or adaptation of the weights must preserve attribution and use the same license as required by those terms.

conf/config.yaml

@@ -0,0 +1,527 @@
+project:
+  name: neuralgcm_develop
+  task: earth_system_forecasting
+  seed: 20260904
+
+paths:
+  project_root: .
+  # Upstream source is supplied by the external neuralgcm package.
+  official_source_dir: null
+  virtual_era5_dir: data
+  checkpoint_dir: data/checkpoint
+  result_dir: results
+  metadata_dir: metadata
+
+model:
+  # Native NeuralGCM pressure-level input contract.
+  variant: weather_forecast
+  grid_degrees: 0.7
+  # Gaussian grid shape in [longitude, latitude] order (the model API reports
+  # the same grid as (latitude, longitude) when printing sizes).
+  grid_shape: [512, 256]
+  profiles:
+    weather_forecast:
+      description: "未来2至15天天气预报"
+      official_reference: models_v1_deterministic_0_7_deg.pkl
+      grid_degrees: 0.7
+      grid_shape: [512, 256]
+    climate_scale:
+      description: "气候尺度模拟"
+      official_reference: models_v1_deterministic_1_4_deg.pkl
+      grid_degrees: 1.4
+      grid_shape: [256, 128]
+    forecast_2_8_deg:
+      description: "2.8度天气预报"
+      official_reference: models_v1_deterministic_2_8_deg.pkl
+      grid_degrees: 2.8
+      grid_shape: [128, 64]
+    stochastic_1_4_deg:
+      description: "1.4度随机预报"
+      official_reference: models_v1_stochastic_1_4_deg.pkl
+      grid_degrees: 1.4
+      grid_shape: [256, 128]
+  pressure_levels_hpa: [1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125, 150, 175, 200, 225, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 775, 800, 825, 850, 875, 900, 925, 950, 975, 1000]
+  input_variables: [geopotential, specific_humidity, temperature, u_component_of_wind, v_component_of_wind]
+  optional_input_variables: [specific_cloud_ice_water_content, specific_cloud_liquid_water_content]
+  forcing_variables: [sea_ice_cover, sea_surface_temperature]
+  official_checkpoint: null
+  load_pretrained: false
+
+data:
+  dataset_class: onescience.datapipes.climate.ERA5Dataset
+  data_dir: data
+  # Optional auxiliary static fields generated by fake_data.py or supplied by
+  # a real ERA5 preprocessing job.  Dynamic channels remain in data/*.h5.
+  static_file: data/static.nc
+  # Exact Gaussian-grid static fields extracted from the four official
+  # checkpoints by scripts/prepare_static_data.py. These take precedence over
+  # the source-grid synthetic fallback above.
+  static_files:
+    weather_forecast: data/static/weather_forecast.nc
+    climate_scale: data/static/climate_scale.nc
+    forecast_2_8_deg: data/static/forecast_2_8_deg.nc
+    stochastic_1_4_deg: data/static/stochastic_1_4_deg.nc
+  field_key: fields
+  time_step_hours: 6
+  input_steps: 1
+  # Official training consumes a time trajectory. Increase for production
+  # rollouts; 1 is retained for the minimal data validation command.
+  output_steps: 1
+  normalize: false
+  batch_size: 1
+  num_workers: 0
+  train_years: [1999]
+  val_years: [2000]
+  test_years: [2001]
+  virtual:
+    # Memory-conscious default: one initial frame + eight future 6-hour
+    # frames. Use --forecast-steps 60 for the full official 15-day horizon.
+    timesteps_per_year: 9
+    forecast_steps: 8
+    forecast_horizon_days: 2
+    height: 721
+    width: 1440
+    seed: 20260904
+  # Exact flattened fields order used by fake_data.py and ERA5Dataset.
+  channel_order: &channel_order
+    - geopotential_1
+    - geopotential_2
+    - geopotential_3
+    - geopotential_5
+    - geopotential_7
+    - geopotential_10
+    - geopotential_20
+    - geopotential_30
+    - geopotential_50
+    - geopotential_70
+    - geopotential_100
+    - geopotential_125
+    - geopotential_150
+    - geopotential_175
+    - geopotential_200
+    - geopotential_225
+    - geopotential_250
+    - geopotential_300
+    - geopotential_350
+    - geopotential_400
+    - geopotential_450
+    - geopotential_500
+    - geopotential_550
+    - geopotential_600
+    - geopotential_650
+    - geopotential_700
+    - geopotential_750
+    - geopotential_775
+    - geopotential_800
+    - geopotential_825
+    - geopotential_850
+    - geopotential_875
+    - geopotential_900
+    - geopotential_925
+    - geopotential_950
+    - geopotential_975
+    - geopotential_1000
+    - specific_humidity_1
+    - specific_humidity_2
+    - specific_humidity_3
+    - specific_humidity_5
+    - specific_humidity_7
+    - specific_humidity_10
+    - specific_humidity_20
+    - specific_humidity_30
+    - specific_humidity_50
+    - specific_humidity_70
+    - specific_humidity_100
+    - specific_humidity_125
+    - specific_humidity_150
+    - specific_humidity_175
+    - specific_humidity_200
+    - specific_humidity_225
+    - specific_humidity_250
+    - specific_humidity_300
+    - specific_humidity_350
+    - specific_humidity_400
+    - specific_humidity_450
+    - specific_humidity_500
+    - specific_humidity_550
+    - specific_humidity_600
+    - specific_humidity_650
+    - specific_humidity_700
+    - specific_humidity_750
+    - specific_humidity_775
+    - specific_humidity_800
+    - specific_humidity_825
+    - specific_humidity_850
+    - specific_humidity_875
+    - specific_humidity_900
+    - specific_humidity_925
+    - specific_humidity_950
+    - specific_humidity_975
+    - specific_humidity_1000
+    - temperature_1
+    - temperature_2
+    - temperature_3
+    - temperature_5
+    - temperature_7
+    - temperature_10
+    - temperature_20
+    - temperature_30
+    - temperature_50
+    - temperature_70
+    - temperature_100
+    - temperature_125
+    - temperature_150
+    - temperature_175
+    - temperature_200
+    - temperature_225
+    - temperature_250
+    - temperature_300
+    - temperature_350
+    - temperature_400
+    - temperature_450
+    - temperature_500
+    - temperature_550
+    - temperature_600
+    - temperature_650
+    - temperature_700
+    - temperature_750
+    - temperature_775
+    - temperature_800
+    - temperature_825
+    - temperature_850
+    - temperature_875
+    - temperature_900
+    - temperature_925
+    - temperature_950
+    - temperature_975
+    - temperature_1000
+    - u_component_of_wind_1
+    - u_component_of_wind_2
+    - u_component_of_wind_3
+    - u_component_of_wind_5
+    - u_component_of_wind_7
+    - u_component_of_wind_10
+    - u_component_of_wind_20
+    - u_component_of_wind_30
+    - u_component_of_wind_50
+    - u_component_of_wind_70
+    - u_component_of_wind_100
+    - u_component_of_wind_125
+    - u_component_of_wind_150
+    - u_component_of_wind_175
+    - u_component_of_wind_200
+    - u_component_of_wind_225
+    - u_component_of_wind_250
+    - u_component_of_wind_300
+    - u_component_of_wind_350
+    - u_component_of_wind_400
+    - u_component_of_wind_450
+    - u_component_of_wind_500
+    - u_component_of_wind_550
+    - u_component_of_wind_600
+    - u_component_of_wind_650
+    - u_component_of_wind_700
+    - u_component_of_wind_750
+    - u_component_of_wind_775
+    - u_component_of_wind_800
+    - u_component_of_wind_825
+    - u_component_of_wind_850
+    - u_component_of_wind_875
+    - u_component_of_wind_900
+    - u_component_of_wind_925
+    - u_component_of_wind_950
+    - u_component_of_wind_975
+    - u_component_of_wind_1000
+    - v_component_of_wind_1
+    - v_component_of_wind_2
+    - v_component_of_wind_3
+    - v_component_of_wind_5
+    - v_component_of_wind_7
+    - v_component_of_wind_10
+    - v_component_of_wind_20
+    - v_component_of_wind_30
+    - v_component_of_wind_50
+    - v_component_of_wind_70
+    - v_component_of_wind_100
+    - v_component_of_wind_125
+    - v_component_of_wind_150
+    - v_component_of_wind_175
+    - v_component_of_wind_200
+    - v_component_of_wind_225
+    - v_component_of_wind_250
+    - v_component_of_wind_300
+    - v_component_of_wind_350
+    - v_component_of_wind_400
+    - v_component_of_wind_450
+    - v_component_of_wind_500
+    - v_component_of_wind_550
+    - v_component_of_wind_600
+    - v_component_of_wind_650
+    - v_component_of_wind_700
+    - v_component_of_wind_750
+    - v_component_of_wind_775
+    - v_component_of_wind_800
+    - v_component_of_wind_825
+    - v_component_of_wind_850
+    - v_component_of_wind_875
+    - v_component_of_wind_900
+    - v_component_of_wind_925
+    - v_component_of_wind_950
+    - v_component_of_wind_975
+    - v_component_of_wind_1000
+    - specific_cloud_ice_water_content_1
+    - specific_cloud_ice_water_content_2
+    - specific_cloud_ice_water_content_3
+    - specific_cloud_ice_water_content_5
+    - specific_cloud_ice_water_content_7
+    - specific_cloud_ice_water_content_10
+    - specific_cloud_ice_water_content_20
+    - specific_cloud_ice_water_content_30
+    - specific_cloud_ice_water_content_50
+    - specific_cloud_ice_water_content_70
+    - specific_cloud_ice_water_content_100
+    - specific_cloud_ice_water_content_125
+    - specific_cloud_ice_water_content_150
+    - specific_cloud_ice_water_content_175
+    - specific_cloud_ice_water_content_200
+    - specific_cloud_ice_water_content_225
+    - specific_cloud_ice_water_content_250
+    - specific_cloud_ice_water_content_300
+    - specific_cloud_ice_water_content_350
+    - specific_cloud_ice_water_content_400
+    - specific_cloud_ice_water_content_450
+    - specific_cloud_ice_water_content_500
+    - specific_cloud_ice_water_content_550
+    - specific_cloud_ice_water_content_600
+    - specific_cloud_ice_water_content_650
+    - specific_cloud_ice_water_content_700
+    - specific_cloud_ice_water_content_750
+    - specific_cloud_ice_water_content_775
+    - specific_cloud_ice_water_content_800
+    - specific_cloud_ice_water_content_825
+    - specific_cloud_ice_water_content_850
+    - specific_cloud_ice_water_content_875
+    - specific_cloud_ice_water_content_900
+    - specific_cloud_ice_water_content_925
+    - specific_cloud_ice_water_content_950
+    - specific_cloud_ice_water_content_975
+    - specific_cloud_ice_water_content_1000
+    - specific_cloud_liquid_water_content_1
+    - specific_cloud_liquid_water_content_2
+    - specific_cloud_liquid_water_content_3
+    - specific_cloud_liquid_water_content_5
+    - specific_cloud_liquid_water_content_7
+    - specific_cloud_liquid_water_content_10
+    - specific_cloud_liquid_water_content_20
+    - specific_cloud_liquid_water_content_30
+    - specific_cloud_liquid_water_content_50
+    - specific_cloud_liquid_water_content_70
+    - specific_cloud_liquid_water_content_100
+    - specific_cloud_liquid_water_content_125
+    - specific_cloud_liquid_water_content_150
+    - specific_cloud_liquid_water_content_175
+    - specific_cloud_liquid_water_content_200
+    - specific_cloud_liquid_water_content_225
+    - specific_cloud_liquid_water_content_250
+    - specific_cloud_liquid_water_content_300
+    - specific_cloud_liquid_water_content_350
+    - specific_cloud_liquid_water_content_400
+    - specific_cloud_liquid_water_content_450
+    - specific_cloud_liquid_water_content_500
+    - specific_cloud_liquid_water_content_550
+    - specific_cloud_liquid_water_content_600
+    - specific_cloud_liquid_water_content_650
+    - specific_cloud_liquid_water_content_700
+    - specific_cloud_liquid_water_content_750
+    - specific_cloud_liquid_water_content_775
+    - specific_cloud_liquid_water_content_800
+    - specific_cloud_liquid_water_content_825
+    - specific_cloud_liquid_water_content_850
+    - specific_cloud_liquid_water_content_875
+    - specific_cloud_liquid_water_content_900
+    - specific_cloud_liquid_water_content_925
+    - specific_cloud_liquid_water_content_950
+    - specific_cloud_liquid_water_content_975
+    - specific_cloud_liquid_water_content_1000
+    - sea_ice_cover
+    - sea_surface_temperature
+
+training:
+  mode: weather_forecast
+  max_steps: 3
+  trajectory_length: 2
+  # Global batch size. For --devices N it is rounded up to a multiple of N;
+  # each replica then receives distinct trajectories.
+  samples_per_step: 1
+  # Number of local JAX devices for optional synchronous data parallelism.
+  devices: 1
+  shuffle: true
+  drop_last: true
+  # OneScience ERA5Dataset samples are prefetched on host threads while the
+  # current DCU step runs. Keep the queue shallow for full 721x1440 fields.
+  data_num_workers: 2
+  prefetch_batches: 1
+  # Full params/EMA/optimizer/reader state is always saved on clean exit. Set a
+  # positive interval for periodic resumable checkpoints during long runs.
+  checkpoint_interval: 0
+  learning_rate: 0.0001
+  optimizer:
+    name: adam
+    schedule: constant
+    b1: 0.9
+    b2: 0.95
+    eps: 1.0e-6
+    # Optional piecewise constant schedule. Empty boundaries use base LR.
+    rates: []
+    boundaries: []
+  # Public Experiment tracks an EMA for evaluation/checkpointing. Set to 0 to
+  # disable; otherwise this is the effective average window in optimizer steps.
+  ema_num_steps: 1000
+  rollout_schedule: []
+  # Public NeuralGCM uses transformed trajectory losses. The private job loss
+  # bindings and complete normalization tables are unavailable, so every
+  # published coefficient and every auditable fallback remain explicit here.
+  gradient_clip_norm: 1.0
+  loss:
+    backend: official
+    # Supplementary G.4 deterministic objective coefficients:
+    # 20*data MSE + 0.1*data spectrum MSE + 1*model MSE
+    # + 0.1*model spectrum MSE + 2*batch spectral bias MSE.
+    data_weight: 20.0
+    data_spectrum_weight: 0.1
+    model_weight: 1.0
+    model_spectrum_weight: 0.1
+    bias_weight: 2.0
+    accuracy_time_scale_hours: 24.0
+    spectral_time_scale_hours: 40.0
+    spectral_cutoff_by_mode:
+      weather_forecast: 120
+      climate_scale: 80
+      forecast_2_8_deg: 42
+    # Optional exact PerVariableRescaling weights. Each value multiplies the
+    # squared error. When null, factor/scale below multiplies the error.
+    variable_weights: null
+    # The paper uses ERA5 24-hour difference standard deviations, but does not
+    # publish the complete numerical tables. These auditable fallbacks keep
+    # physical variables balanced; replace them with statistics calculated
+    # from the exact ERA5 training vintage for a precision reproduction.
+    time_rescaling: legacy
+    spectral_weight: 0.0
+    variable_scales:
+      z: 10000.0
+      t: 30.0
+      u: 30.0
+      v: 30.0
+      specific_humidity: 0.01
+      specific_cloud_ice_water_content: 1.0e-5
+      specific_cloud_liquid_water_content: 2.0e-5
+      divergence: 0.1
+      vorticity: 0.1
+      log_surface_pressure: 0.1
+      default: 1.0
+    # Additional balancing factors stated explicitly in Supplementary G.3.
+    variable_factors:
+      z: 2.0
+      specific_humidity: 0.66
+      log_surface_pressure: 5.0
+      specific_cloud_ice_water_content: 0.05
+      specific_cloud_liquid_water_content: 0.05
+      default: 1.0
+    # Order 12 is exact. Absolute half-power cutoffs below are digitized from
+    # Supplementary Fig. 8 because the underlying numeric table was not
+    # released. Interpolation is performed at the configured output times.
+    predictability_filter:
+      enabled: true
+      order: 12
+      lead_hours: [0, 6, 12, 24, 36, 48, 60, 72]
+      cutoffs:
+        temperature: [80, 120, 120, 95, 45, 35, 30, 25]
+        wind: [80, 120, 115, 82, 48, 36, 29, 24]
+        moisture: [80, 120, 110, 52, 34, 28, 24, 21]
+        divergence: [80, 120, 105, 43, 24, 19, 16, 14]
+        default: [80, 120, 115, 82, 48, 36, 29, 24]
+    # Optional multiplicative weights for pressure levels, ordered as the
+    # configured ERA5 pressure-level list. Empty means uniform weighting.
+    level_weights: []
+  # Long-run reproduction settings inferred from the public paper description
+  # and released training pseudocode. The paper's private job bindings are not
+  # available, so these are explicit project settings rather than exact claims.
+  # They are enabled only by --paper-defaults; CLI values remain highest priority.
+  profiles:
+    weather_forecast:
+      max_steps: 25000
+      learning_rate: 0.001
+      optimizer: &paper_optimizer
+        schedule: neuralgcm
+        warmup_steps: 2000
+        decay_start: 15000
+        decay_steps: 10000
+        decay_rate: 0.5
+      rollout_schedule:
+        - {trajectory_length: 2, until_step: 0}      # 6 h
+        - {trajectory_length: 3, until_step: 500}    # 12 h
+        - {trajectory_length: 4, until_step: 2000}   # 18 h
+        - {trajectory_length: 5, until_step: 4500}   # 24 h
+        - {trajectory_length: 7, until_step: 8000}   # 36 h
+        - {trajectory_length: 9, until_step: 12500}  # 48 h
+        - {trajectory_length: 11, until_step: 18000} # 60 h
+    climate_scale:
+      max_steps: 26000
+      learning_rate: 0.002
+      optimizer: *paper_optimizer
+      rollout_schedule: &coarse_rollout_schedule
+        - {trajectory_length: 3, until_step: 0}       # 12 h
+        - {trajectory_length: 5, until_step: 2000}    # 24 h
+        - {trajectory_length: 7, until_step: 5656}    # 36 h
+        - {trajectory_length: 9, until_step: 10392}   # 48 h
+        - {trajectory_length: 11, until_step: 16000}  # 60 h
+        - {trajectory_length: 13, until_step: 22360}  # 72 h
+    forecast_2_8_deg:
+      max_steps: 38000
+      learning_rate: 0.002
+      optimizer: *paper_optimizer
+      rollout_schedule: *coarse_rollout_schedule
+    stochastic_1_4_deg:
+      max_steps: 43000
+      learning_rate: 0.001
+      ensemble_size: 2
+      optimizer: *paper_optimizer
+      rollout_schedule:
+        - {trajectory_length: 2, until_step: 0}       # 6 h
+        - {trajectory_length: 3, until_step: 500}     # 12 h
+        - {trajectory_length: 4, until_step: 2000}    # 18 h
+        - {trajectory_length: 5, until_step: 4500}    # 24 h
+        - {trajectory_length: 7, until_step: 8000}    # 36 h
+        - {trajectory_length: 9, until_step: 12500}   # 48 h
+        - {trajectory_length: 11, until_step: 18000}  # 60 h
+        - {trajectory_length: 13, until_step: 24500}  # 72 h
+        - {trajectory_length: 17, until_step: 32000}  # 96 h
+        - {trajectory_length: 21, until_step: 40500}  # 120 h
+      loss:
+        backend: crps
+        variable_weights: null
+        variable_scale: 1.0
+        nodal_time_scale_hours: 24.0
+        spectral_time_scale_hours: 40.0
+        spectral_max_wavenumber: 80
+  checkpoint: null
+  gin_config: null
+  train_dataset: null
+  eval_dataset: null
+
+inference:
+  mode: weather_forecast
+  # Used by stochastic profiles; deterministic checkpoints ignore the key.
+  seed: 20260904
+  # Memory-conscious default: eight 6-hour outputs reach forecast day 2.
+  # Set --steps 60 to exercise the model's full 15-day capability.
+  prediction_steps: 8
+  output_interval_hours: 6
+  # An official-format local model_bak.pkl takes precedence when present;
+  # otherwise inference selects the profile's bundled checkpoint.
+  checkpoint: data/checkpoint/model_bak.pkl
+  output: results/predictions.nc
+
+runtime:
+  platform: auto
+  dcu_device: 0

config.json

@@ -0,0 +1,147 @@
+{
+  "model_name": "NeuralGCM",
+  "model_type": "neuralgcm",
+  "architectures": [
+    "PressureLevelModel"
+  ],
+  "framework": "JAX with Haiku and Gin",
+  "domain": "atmosphere-and-climate",
+  "task": "global-weather-forecasting-and-climate-simulation",
+  "implementation": {
+    "entry_point": "model/NeuralGCM.py",
+    "scope": "vendored NeuralGCM pressure-level inference and training facade with differentiable dynamics and neural parameterizations"
+  },
+  "architecture": {
+    "family": "hybrid differentiable general circulation model",
+    "dynamical_core": "spectral atmospheric dynamics",
+    "learned_components": [
+      "encoder",
+      "decoder",
+      "subgrid physical parameterizations"
+    ],
+    "input_format": "xarray pressure-level fields",
+    "grid_shape_order": [
+      "longitude",
+      "latitude"
+    ],
+    "time_step_hours": 6,
+    "pressure_level_count": 37,
+    "dynamic_variable_count": 7,
+    "forcing_variable_count": 2,
+    "flattened_input_channels": 261,
+    "profiles": {
+      "weather_forecast": {
+        "type": "deterministic weather forecast",
+        "resolution_degrees": 0.7,
+        "grid_shape": [
+          512,
+          256
+        ],
+        "checkpoint": "weight/models_v1_deterministic_0_7_deg.pkl"
+      },
+      "climate_scale": {
+        "type": "deterministic climate simulation",
+        "resolution_degrees": 1.4,
+        "grid_shape": [
+          256,
+          128
+        ],
+        "checkpoint": "weight/models_v1_deterministic_1_4_deg.pkl"
+      },
+      "forecast_2_8_deg": {
+        "type": "deterministic weather forecast",
+        "resolution_degrees": 2.8,
+        "grid_shape": [
+          128,
+          64
+        ],
+        "checkpoint": "weight/models_v1_deterministic_2_8_deg.pkl"
+      },
+      "stochastic_1_4_deg": {
+        "type": "stochastic weather forecast",
+        "resolution_degrees": 1.4,
+        "grid_shape": [
+          256,
+          128
+        ],
+        "checkpoint": "weight/models_v1_stochastic_1_4_deg.pkl"
+      }
+    }
+  },
+  "data": {
+    "dataset": "ERA5",
+    "source_grid_shape": [
+      721,
+      1440
+    ],
+    "time_step_hours": 6,
+    "input_steps": 1,
+    "pressure_levels_hpa": [
+      1,
+      2,
+      3,
+      5,
+      7,
+      10,
+      20,
+      30,
+      50,
+      70,
+      100,
+      125,
+      150,
+      175,
+      200,
+      225,
+      250,
+      300,
+      350,
+      400,
+      450,
+      500,
+      550,
+      600,
+      650,
+      700,
+      750,
+      775,
+      800,
+      825,
+      850,
+      875,
+      900,
+      925,
+      950,
+      975,
+      1000
+    ],
+    "input_variables": [
+      "geopotential",
+      "specific_humidity",
+      "temperature",
+      "u_component_of_wind",
+      "v_component_of_wind"
+    ],
+    "optional_input_variables": [
+      "specific_cloud_ice_water_content",
+      "specific_cloud_liquid_water_content"
+    ],
+    "forcing_variables": [
+      "sea_ice_cover",
+      "sea_surface_temperature"
+    ],
+    "protocol": "era5_37_pressure_levels_261_channel"
+  },
+  "weights": {
+    "license": "CC BY-SA 4.0",
+    "source": "gs://neuralgcm/models/v1/"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "configuration.json",
+    "model/NeuralGCM.py",
+    "model/legacy",
+    "model/reference_code",
+    "scripts/common.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "JAX",
+  "task": "weather_and_climate_simulation",
+  "model": "NeuralGCM",
+  "input_format": "xarray_pressure_level",
+  "protocol": "era5_37_pressure_levels_261_channel",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/NeuralGCM.py

@@ -0,0 +1,230 @@
+"""Official NeuralGCM implementation facade.
+
+The upstream legacy model, encoders, decoders, dynamical core and reference
+training utilities are vendored directly under this project's ``model``
+namespace (``model/legacy`` and ``model/reference_code``). This file is the
+single project-facing entry point; no external ``neuralgcm`` source directory
+is required at runtime.
+"""
+from __future__ import annotations
+
+import pickle
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+PROFILE_GIN = {
+    "weather_forecast": "deterministic_0_7_deg.gin",
+    "climate_scale": "deterministic_1_4_deg.gin",
+    "forecast_2_8_deg": "deterministic_2_8_deg.gin",
+    "stochastic_1_4_deg": "stochastic_1_4_deg.gin",
+}
+
+MODE_ALIASES = {
+    "forecast": "weather_forecast",
+    "weather_forecast": "weather_forecast",
+    "climate": "climate_scale",
+    "climate_scale": "climate_scale",
+    "forecast_2_8_deg": "forecast_2_8_deg",
+    "stochastic_1_4_deg": "stochastic_1_4_deg",
+}
+
+
+class OfficialNeuralGCMUnavailable(RuntimeError):
+    """Raised when the official runtime package is not available."""
+
+
+class CheckpointFormatError(ValueError):
+    """Raised when a file is not an official NeuralGCM checkpoint."""
+
+
+def checkpoint_mode(payload: object) -> str | None:
+    """Infer the project profile declared by an official-format checkpoint."""
+    if not isinstance(payload, dict):
+        return None
+    if payload.get("mode"):
+        value = str(payload["mode"])
+        return MODE_ALIASES.get(value, value)
+    text = str(payload.get("model_config_str", ""))
+    if "GridTL255" in text:
+        return "weather_forecast"
+    if "GridTL63" in text:
+        return "forecast_2_8_deg"
+    if "GridTL127" in text:
+        return "stochastic_1_4_deg" if "FIELD_SUBSET" in text else "climate_scale"
+    return None
+
+
+def validate_checkpoint_mode(payload: object, mode: str, path: str | Path) -> None:
+    """Reject a checkpoint whose grid/profile differs from the requested mode."""
+    stored_mode = checkpoint_mode(payload)
+    if stored_mode and stored_mode != mode:
+        raise ValueError(
+            f"Checkpoint {path} is for mode={stored_mode!r}, but mode={mode!r} "
+            "was requested. Select the matching mode or checkpoint."
+        )
+
+
+def parameter_summary(params: Any) -> dict[str, Any]:
+    """Return reproducible parameter count, storage size and dtype statistics."""
+    import jax
+
+    leaves = jax.tree_util.tree_leaves(params)
+    array_leaves = [leaf for leaf in leaves if hasattr(leaf, "shape") and hasattr(leaf, "dtype")]
+    count = sum(int(np.prod(leaf.shape, dtype=np.int64)) for leaf in array_leaves)
+    nbytes = sum(
+        int(np.prod(leaf.shape, dtype=np.int64)) * np.dtype(leaf.dtype).itemsize
+        for leaf in array_leaves
+    )
+    dtype_counts: dict[str, int] = {}
+    for leaf in array_leaves:
+        dtype = str(np.dtype(leaf.dtype))
+        dtype_counts[dtype] = dtype_counts.get(dtype, 0) + int(
+            np.prod(leaf.shape, dtype=np.int64)
+        )
+    return {
+        "count": count,
+        "nbytes": nbytes,
+        "leaves": len(array_leaves),
+        "dtypes": dtype_counts,
+    }
+
+
+def format_parameter_summary(params: Any) -> str:
+    """Format a compact ``params.count``-style model summary."""
+    summary = parameter_summary(params)
+    dtype_text = ",".join(
+        f"{dtype}:{count:,}" for dtype, count in sorted(summary["dtypes"].items())
+    )
+    return (
+        f"params.count={summary['count']:,} "
+        f"params.bytes={summary['nbytes']:,} "
+        f"params.mib={summary['nbytes'] / 2**20:.2f} "
+        f"params.leaves={summary['leaves']} dtypes={dtype_text}"
+    )
+
+
+def load_checkpoint(path: str | Path):
+    """Load an official checkpoint through the vendored PressureLevelModel."""
+    try:
+        from model.legacy.api import PressureLevelModel
+    except Exception as exc:  # pragma: no cover - runtime-dependent
+        raise OfficialNeuralGCMUnavailable(
+            "Unable to import the vendored NeuralGCM implementation. Check "
+            "JAX, Haiku, Gin and Dinosaur dependencies in develop_base."
+        ) from exc
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(path)
+    with path.open("rb") as handle:
+        checkpoint = pickle.load(handle)
+    required = {"model_config_str", "aux_ds_dict", "params"}
+    if not isinstance(checkpoint, dict) or not required.issubset(checkpoint):
+        keys = sorted(checkpoint) if isinstance(checkpoint, dict) else type(checkpoint).__name__
+        raise CheckpointFormatError(
+            f"{path} is not an official checkpoint; expected keys "
+            f"{sorted(required)}, got {keys}"
+        )
+    return PressureLevelModel.from_checkpoint(checkpoint)
+
+
+def official_runtime_available() -> bool:
+    try:
+        from model.legacy.api import PressureLevelModel  # noqa: F401
+    except Exception:
+        return False
+    return True
+
+
+def build_from_scratch(dataset, mode: str):
+    """Build the public WhirlModel used for random parameter initialization.
+
+    Parameter initialization itself needs a concrete trajectory and is performed
+    by ``scripts/train.py`` through the returned model's Haiku rollout function.
+    This compatibility facade deliberately does not import the unreleased Google
+    experiment runner.
+    """
+    return build_training_model(dataset, mode)
+
+
+def build_training_model(dataset, mode: str):
+    """Build an official ``WhirlModel`` from the fused Gin profile."""
+    if mode not in PROFILE_GIN:
+        raise ValueError(f"Unknown NeuralGCM mode {mode!r}")
+    import gin
+    from model.legacy import model_builder
+
+    config_path = Path(__file__).resolve().parent / "reference_code" / "paper_configs" / PROFILE_GIN[mode]
+    gin_text = config_path.read_text(encoding="utf-8")
+    # The released Gin profiles use ``orography_data_path = None`` and rely on
+    # the official xarray auxiliary-dataset escape hatch for static fields.
+    # ``get_whirl_model`` normally obtains this from dataset metadata; supply it
+    # explicitly for synthetic/OneScience datasets that have no metadata attrs.
+    from dinosaur import xarray_utils
+    try:
+        aux_features = xarray_utils.aux_features_from_xarray(dataset)
+    except (KeyError, AttributeError):
+        aux_features = {}
+    aux_features[xarray_utils.XARRAY_DS_KEY] = dataset
+    dataset = dataset.copy()
+    dataset.attrs = dict(dataset.attrs)
+    dataset.attrs[xarray_utils.XR_AUX_FEATURES_LIST_KEY] = ",".join(
+        key for key in aux_features if key != xarray_utils.XARRAY_DS_KEY
+    )
+    # get_whirl_model reads serializable aux variables from attrs. Injecting the
+    # xarray dataset directly is handled below through a temporary wrapper.
+    original = model_builder.xarray_utils.aux_features_from_xarray
+    model_builder.xarray_utils.aux_features_from_xarray = lambda _: aux_features
+    try:
+        model = model_builder.get_whirl_model(dataset, gin_text)
+    finally:
+        model_builder.xarray_utils.aux_features_from_xarray = original
+    # The profile's xarray conversion callbacks are configured through Gin;
+    # get_whirl_model returns the fully bound model object.
+    return model, gin_text
+
+
+def make_rollout_functions(
+    whirl_model, trajectory_length: int, *, inner_steps: int = 1
+):
+    """Return Haiku init/apply functions using the official rollout helpers."""
+    import haiku as hk
+    from model.legacy import model_utils
+
+    @hk.transform
+    def rollout_fn(target, forcing):
+        model = whirl_model.model_cls()
+        trajectory_fn = model_utils.trajectory_with_inputs_and_forcing(
+            model, num_init_frames=1, start_with_input=True
+        )
+        _, predicted = trajectory_fn(
+            target,
+            forcing,
+            outer_steps=trajectory_length,
+            inner_steps=inner_steps,
+        )
+        return model_utils.compute_prediction_and_target_representations(
+            predicted, target, forcing, model
+        )
+
+    return rollout_fn
+
+
+def save_official_checkpoint(path: str | Path, params: Any, dataset, model_config_str: str, *, metadata: dict[str, Any] | None = None):
+    """Write a checkpoint consumable by ``PressureLevelModel.from_checkpoint``."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "model_config_str": model_config_str,
+        "aux_ds_dict": dataset.to_dict(),
+        "params": params,
+    }
+    if metadata:
+        payload.update(metadata)
+    with path.open("wb") as handle:
+        pickle.dump(payload, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    return path
+
+
+NeuralGCMAdapter = load_checkpoint

model/NeuralGCM_demo.py

@@ -0,0 +1,60 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import importlib.resources
+import pickle
+
+from dinosaur import coordinate_systems
+from dinosaur import horizontal_interpolation
+from dinosaur import spherical_harmonic
+import model as neuralgcm
+import numpy as np
+import xarray
+
+
+def _horizontal_regrid(
+    regridder: horizontal_interpolation.Regridder, dataset: xarray.Dataset
+) -> xarray.Dataset:
+  """Horizontally regrid an xarray Dataset."""
+  # TODO(shoyer): consider moving to public API
+  regridded = xarray.apply_ufunc(
+      regridder,
+      dataset,
+      input_core_dims=[['longitude', 'latitude']],
+      output_core_dims=[['longitude', 'latitude']],
+      exclude_dims={'longitude', 'latitude'},
+      vectorize=True,  # loops over level, for lower memory usage
+  )
+  regridded.coords['longitude'] = np.rad2deg(regridder.target_grid.longitudes)
+  regridded.coords['latitude'] = np.rad2deg(regridder.target_grid.latitudes)
+  return regridded
+
+
+def load_checkpoint_tl63_stochastic():
+  """Load a checkpoint for a toy TL63 stochastic model."""
+  package = importlib.resources.files(neuralgcm)
+  file = package.joinpath('data/tl63_stochastic_mini.pkl')
+  return pickle.loads(file.read_bytes())
+
+
+def load_data(coords: coordinate_systems.CoordinateSystem) -> xarray.Dataset:
+  """Load demo data for the given coordinate system."""
+  if coords.vertical.layers != 37:
+    raise ValueError('can only load demo data for 37 pressure levels')
+  package = importlib.resources.files(neuralgcm)
+  with package.joinpath('data/era5_tl31_19590102T00.nc').open('rb') as f:
+    ds = xarray.load_dataset(f).expand_dims('time')
+  regridder = horizontal_interpolation.ConservativeRegridder(
+      spherical_harmonic.Grid.TL31(), coords.horizontal
+  )
+  return _horizontal_regrid(regridder, ds)

model/data/era5_tl31_19590102T00.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f66e795af9f564a2b6e0d861b9a51e74ce831a82675b47ff957be709554a5e
+size 2141788

model/data/tl63_stochastic_mini.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de0f8e8cb31ecca59469104440808a3ea8e31418d5a8a48220e3c1fa64303baa
+size 1030438

model/legacy/api.py

@@ -0,0 +1,601 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Public API for NeuralGCM models."""
+from __future__ import annotations
+
+from collections import abc
+import datetime
+import functools
+from typing import Any, Callable
+
+from dinosaur import coordinate_systems
+from dinosaur import scales
+from dinosaur import time_integration
+from dinosaur import typing
+from dinosaur import xarray_utils
+import jax
+from jax import tree_util
+import jax.numpy as jnp
+from model.legacy import gin_utils
+from model.legacy import model_builder
+from model.legacy import physics_specifications
+import numpy as np
+import pandas as pd
+import xarray
+
+
+ArrayLike = float | np.ndarray | jax.Array
+Params = dict[str, dict[str, ArrayLike]]
+TimedeltaLike = str | np.timedelta64 | pd.Timestamp | datetime.timedelta
+Numeric = float | np.ndarray | jax.Array | xarray.DataArray
+
+# TODO(shoyer): make these types more precise
+Inputs = dict[str, ArrayLike]
+Forcings = dict[str, ArrayLike]
+TemporalForcings = dict[str, ArrayLike]
+Outputs = dict[str, jax.Array]
+BatchedOutputs = dict[str, jax.Array]
+State = Any
+
+
+def _sim_time_from_state(state: State) -> jax.Array:
+  """Extract sim_time from model state."""
+  # TODO(shoyer): eliminate whichever of these two cases is no longer needed!
+  # TODO(shoyer): consider renaming `sim_time` to `time`?
+  if isinstance(state, typing.ModelState):
+    sim_time = getattr(state.state, 'sim_time', None)
+  else:
+    sim_time = getattr(state, 'sim_time', None)
+  return sim_time  # pyrefly: ignore[bad-return]
+
+
+def _calculate_sub_steps(
+    timestep: np.timedelta64, duration: TimedeltaLike
+) -> int:
+  """Calculate the number of time-steps required to simulate a time interval."""
+  duration = pd.Timedelta(duration)
+  time_step_ratio = duration / timestep
+  if abs(time_step_ratio - round(time_step_ratio)) > 1e-6:
+    raise ValueError(
+        f'non-integral time-step ratio: {duration=} is not a multiple of '
+        f'the internal model timestep {timestep}'
+    )
+  return round(time_step_ratio)
+
+
+def _prepend_dummy_time_axis(state: typing.Pytree) -> typing.Pytree:
+  return tree_util.tree_map(lambda x: jnp.expand_dims(x, axis=0), state)
+
+
+def _static_gin_config(method):
+  """Decorator to add static gin config to a method."""
+
+  @functools.wraps(method)
+  def _method(self, *args, **kwargs):
+    with gin_utils.specific_config(self.gin_config):
+      return method(self, *args, **kwargs)
+
+  return _method
+
+
+def _check_variables(
+    dataset: xarray.Dataset,
+    desired_level_variables: abc.Sequence[str] = (),
+    desired_surface_variables: abc.Sequence[str] = (),
+):
+  """Checks that a dataset has the desired variables."""
+  T, Z, X, Y = ('time', 'level', 'longitude', 'latitude')  # pylint: disable=invalid-name
+
+  for k in desired_level_variables:
+    if k not in dataset.data_vars:
+      raise ValueError(f'expected variable {k} not found')
+    dims = dataset[k].dims
+    if not (set(dims) == {Z, X, Y} or set(dims) == {T, Z, X, Y}):
+      raise ValueError(
+          f'expected variable {k} to have dims {(Z, X, Y)} or {(T, Z, X, Y)},'
+          f' but got {dims}'
+      )
+
+  for k in desired_surface_variables:
+    if k not in dataset.data_vars:
+      raise ValueError(f'expected variable {k} not found')
+    dims = dataset[k].dims
+    if not (set(dims) == {X, Y} or set(dims) == {T, X, Y}):
+      raise ValueError(
+          f'expected variable {k} to have dims {(X, Y)} or {(T, X, Y)},'
+          f' but got {dims}'
+      )
+
+
+def _check_coords(
+    actual_coords: coordinate_systems.CoordinateSystem,
+    desired_coords: coordinate_systems.CoordinateSystem,
+) -> None:
+  """Checks that a dataset has the desired coordinates."""
+  if not np.allclose(
+      actual := actual_coords.horizontal.longitudes,
+      desired := desired_coords.horizontal.longitudes,
+      atol=1e-3,
+  ):
+    raise ValueError(f'longitude coordinate mismatch: {actual=}, {desired=}')
+
+  if not np.allclose(
+      actual := actual_coords.horizontal.latitudes,
+      desired := desired_coords.horizontal.latitudes,
+      atol=1e-3,
+  ):
+    raise ValueError(f'latitude coordinate mismatch: {actual=}, {desired=}')
+
+  if actual_coords.vertical is not None and not np.allclose(
+      actual := actual_coords.vertical.centers,
+      desired := desired_coords.vertical.centers,
+      atol=1e-3,
+  ):
+    raise ValueError(
+        f'pressure level coordinate mismatch: {actual=}, {desired=}'
+    )
+
+
+def _rename_if_found(
+    dataset: xarray.Dataset, names: dict[str, str]
+) -> xarray.Dataset:
+  return dataset.rename({k: v for k, v in names.items() if k in dataset})
+
+
+_ABBREVIATED_NAMES = {
+    'u_component_of_wind': 'u',
+    'v_component_of_wind': 'v',
+    'geopotential': 'z',
+    'temperature': 't',
+    'longitude': 'lon',
+    'latitude': 'lat',
+}
+_FULL_NAMES = {v: k for k, v in _ABBREVIATED_NAMES.items()}
+
+
+def _expand_tracers(inputs: dict) -> dict:
+  inputs = inputs.copy()
+  inputs.update(inputs.pop('tracers'))
+  assert not inputs['diagnostics']
+  del inputs['diagnostics']
+  return inputs
+
+
+@tree_util.register_pytree_node_class
+class PressureLevelModel:
+  """Inference-only API for models that predict dense data on pressure levels.
+
+  These models are trained on ECMWF ERA5 data on pressure-levels as stored in
+  the Copernicus Data Store.
+
+  This class encapsulates the details of defining models (e.g., with Haiku) and
+  hence should remain stable even for future NeuralGCM models.
+  """
+
+  def __init__(
+      self,
+      structure: model_builder.WhirlModel,
+      params: Params,
+      gin_config: str,
+  ):
+    self._structure = structure
+    self._params = params
+    self.gin_config = gin_config
+
+    self._tracer_variables = [
+        'specific_humidity',
+    ]
+    self._input_variables = [
+        'geopotential',
+        'specific_humidity',
+        'temperature',
+        'u_component_of_wind',
+        'v_component_of_wind',
+    ]
+    # Some old model versions do not use cloud variables.
+    # TODO(shoyer): remove this once all integration tests are updated.
+    cloud_variables = [
+        'specific_cloud_ice_water_content',
+        'specific_cloud_liquid_water_content',
+    ]
+    for variable in cloud_variables:
+      if variable in self.gin_config:
+        self._tracer_variables.append(variable)
+        self._input_variables.append(variable)
+
+    self._forcing_variables = [
+        'sea_ice_cover',
+        'sea_surface_temperature',
+    ]
+
+  def __repr__(self):
+    return (
+        f'{self.__class__.__name__}(structure={self._structure},'
+        f' params={self._params})'
+    )
+
+  @property
+  def params(self) -> Params:
+    return self._params
+
+  def tree_flatten(self):
+    leaves, params_def = tree_util.tree_flatten(self.params)
+    return (leaves, (params_def, self._structure, self.gin_config))
+
+  @classmethod
+  def tree_unflatten(cls, aux_data, leaves):
+    params_def, structure, gin_config = aux_data
+    params = tree_util.tree_unflatten(params_def, leaves)
+    return cls(structure, params, gin_config)
+
+  @property
+  def input_variables(self) -> list[str]:
+    """List of variable names required in `inputs` by this model."""
+    return list(self._input_variables)
+
+  @property
+  def forcing_variables(self) -> list[str]:
+    """List of variable names required in `forcings` by this model."""
+    return list(self._forcing_variables)
+
+  @property
+  def timestep(self) -> np.timedelta64:
+    """Spacing between internal model timesteps."""
+    to_timedelta = (
+        self._structure.specs.physics_specs.dimensionalize_timedelta64
+    )
+    return to_timedelta(self._structure.specs.dt)
+
+  @property
+  def data_coords(self) -> coordinate_systems.CoordinateSystem:
+    """Coordinate system for input and output data."""
+    return self._structure.data_coords
+
+  @property
+  def model_coords(self) -> coordinate_systems.CoordinateSystem:
+    """Coordinate system for internal model state."""
+    return self._structure.coords
+
+  def _check_coords(self, dataset: xarray.Dataset):
+    dataset_coords = model_builder.coordinate_system_from_dataset(dataset)
+    _check_coords(dataset_coords, self.data_coords)
+
+  def _dataset_with_sim_time(self, dataset: xarray.Dataset) -> xarray.Dataset:
+    ref_datetime = self._structure.specs.aux_features['reference_datetime']
+    return xarray_utils.ds_with_sim_time(
+        dataset,
+        self._structure.specs.physics_specs,
+        reference_datetime=ref_datetime,
+    )
+
+  def _to_abbreviated_names_and_tracers(self, inputs: dict) -> dict:
+    inputs = {_ABBREVIATED_NAMES.get(k, k): v for k, v in inputs.items()}
+    inputs['tracers'] = {
+        k: inputs.pop(k) for k in self._tracer_variables if k in inputs
+    }
+    inputs['diagnostics'] = {}
+    return inputs
+
+  def _from_abbreviated_names_and_tracers(self, outputs: dict) -> dict:
+    outputs = {_FULL_NAMES.get(k, k): v for k, v in outputs.items()}
+    outputs |= outputs.pop('tracers')
+    outputs |= outputs.pop('diagnostics')
+    return outputs
+
+  def to_nondim_units(self, value: Numeric, units: str) -> Numeric:
+    """Scale a value to the model's internal non-dimensional units."""
+    scale_ = self._structure.specs.physics_specs.scale
+    units_ = scales.parse_units(units)
+    return scale_.nondimensionalize(value * units_)
+
+  def from_nondim_units(self, value: Numeric, units: str) -> Numeric:
+    """Scale a value from the model's internal non-dimensional units."""
+    scale_ = self._structure.specs.physics_specs.scale
+    units_ = scales.parse_units(units)
+    return scale_.dimensionalize(value, units_).magnitude
+
+  def datetime64_to_sim_time(self, datetime64: np.ndarray) -> np.ndarray:
+    """Converts a datetime64 array to sim_time."""
+    ref_datetime = self._structure.specs.aux_features['reference_datetime']
+    return xarray_utils.datetime64_to_nondim_time(
+        datetime64,
+        self._structure.specs.physics_specs,
+        reference_datetime=ref_datetime,
+    )
+
+  def sim_time_to_datetime64(self, sim_time: np.ndarray) -> np.ndarray:
+    """Converts a sim_time array to datetime64."""
+    ref_datetime = self._structure.specs.aux_features['reference_datetime']
+    return xarray_utils.nondim_time_to_datetime64(
+        sim_time,
+        self._structure.specs.physics_specs,
+        reference_datetime=ref_datetime,
+    )
+
+  def _data_from_xarray(
+      self, dataset: xarray.Dataset, variables: list[str]
+  ) -> dict[str, np.ndarray]:
+    self._check_coords(dataset)
+    dataset = dataset[variables]
+    dataset = self._dataset_with_sim_time(dataset)
+    dataset = _rename_if_found(dataset, {'longitude': 'lon', 'latitude': 'lat'})
+    return xarray_utils.xarray_to_data_dict(dataset)
+
+  def inputs_from_xarray(
+      self, dataset: xarray.Dataset
+  ) -> dict[str, np.ndarray]:
+    """Extract inputs from an xarray.Dataset."""
+    _check_variables(dataset, desired_level_variables=self._input_variables)
+    return self._data_from_xarray(dataset, self._input_variables)
+
+  def forcings_from_xarray(
+      self, dataset: xarray.Dataset
+  ) -> dict[str, np.ndarray]:
+    """Extract forcings from an xarray.Dataset."""
+    _check_variables(dataset, desired_surface_variables=self._forcing_variables)
+    return self._data_from_xarray(dataset, self._forcing_variables)
+
+  def data_from_xarray(
+      self, dataset: xarray.Dataset
+  ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+    """Extracts data and forcings from an xarray.Dataset."""
+    inputs = self.inputs_from_xarray(dataset)
+    forcings = self.forcings_from_xarray(dataset)
+    return (inputs, forcings)
+
+  def data_to_xarray(
+      self,
+      data: dict[str, ArrayLike],
+      times: np.ndarray | None,
+      decoded: bool = True,
+  ) -> xarray.Dataset:
+    """Converts decoded model predictions to xarray.Dataset format.
+
+    Args:
+      data: dict of arrays with shapes matching input/outputs or encoded model
+        state for this model, i.e., with shape `([time,] level, longitude,
+        latitude)`, where `[time,]` indicates an optional leading time
+        dimension.
+      times: either `None` indicating no leading time dimension on any
+        variables, or a coordinate array of times with shape `(time,)`.
+      decoded: if `True`, use `self.data_coords` to determine the output
+        coordinates; otherwise use `self.model_coords`.
+
+    Returns:
+      An xarray.Dataset with appropriate coordinates and dimensions.
+    """
+    coords = self.data_coords if decoded else self.model_coords
+    dataset = xarray_utils.data_to_xarray(data, coords=coords, times=times)
+    dataset = _rename_if_found(dataset, {'lon': 'longitude', 'lat': 'latitude'})
+    return dataset
+
+  def _squeeze_level_from_forcings(self, forcings: Forcings) -> Forcings:
+    # Due to a bug in xarray_to_dynamic_covariate_data, we were accidentally
+    # not inserting a level dimension in forcings.
+    forcings = dict(forcings)
+    for k in self._forcing_variables:
+      if k in forcings:
+        assert isinstance(forcings[k], (np.ndarray, jax.Array))
+        forcings[k] = forcings[k].squeeze(axis=-3)  # pyrefly: ignore[missing-attribute]
+    return forcings
+
+  @jax.jit
+  @_static_gin_config
+  def encode(
+      self,
+      inputs: Inputs,
+      forcings: Forcings,
+      rng_key: typing.PRNGKeyArray | None = None,
+  ) -> State:
+    """Encode from pressure-level inputs & forcings to model state.
+
+    Args:
+      inputs: input data on pressure-levels, as a dict where each entry is an
+        array with shape `[level, longitude, latitude]` matching `data_coords`.
+      forcings: forcing data on pressure-levels, as a dict where each entry is
+        an array with shape `[level, longitude, latitude]` matching
+        `data_coords`. Single level data (e.g., sea surface temperature) should
+        have a `level` dimension of size 1.
+      rng_key: optional JAX RNG key to use for encoding the state. Required if
+        using stochastic models, otherwise ignored.
+
+    Returns:
+      Dynamical core state on sigma levels, where all arrays have dimensions
+      `[level, zonal_wavenumber, total_wavenumber]` matching `model_coords`.
+    """
+    sim_time = inputs['sim_time']
+    inputs = self._to_abbreviated_names_and_tracers(inputs)
+    inputs = _prepend_dummy_time_axis(inputs)
+    forcings = self._squeeze_level_from_forcings(forcings)
+    forcings = _prepend_dummy_time_axis(forcings)
+    f = self._structure.forcing_fn(self.params, None, forcings, sim_time)
+    return self._structure.encode_fn(self.params, rng_key, inputs, f)
+
+  @jax.jit
+  @_static_gin_config
+  def advance(self, state: State, forcings: Forcings) -> State:
+    """Advance model state one timestep forward.
+
+    Args:
+      state: dynamical core state on sigma levels, where all arrays have
+        dimensions `[level, zonal_wavenumber, total_wavenumber]` matching
+        `model_coords`
+      forcings: forcing data on pressure-levels, as a dict where each entry is
+        an array with shape `[level, longitude, latitude]` matching
+        `data_coords`. Single level data (e.g., sea surface temperature) should
+        have a `level` dimension of size 1.
+
+    Returns:
+      State advanced one time-step forward.
+    """
+    sim_time = _sim_time_from_state(state)
+    forcings = self._squeeze_level_from_forcings(forcings)
+    forcings = _prepend_dummy_time_axis(forcings)
+    f = self._structure.forcing_fn(self.params, None, forcings, sim_time)
+    state = self._structure.advance_fn(self.params, None, state, f)
+    return state
+
+  @jax.jit
+  @_static_gin_config
+  def decode(self, state: State, forcings: Forcings) -> Outputs:
+    """Decode from model state to pressure-level outputs.
+
+    Args:
+      state: dynamical core state on sigma levels, where all arrays have
+        dimensions `[level, zonal_wavenumber, total_wavenumber]` matching
+        `model_coords`.
+      forcings: forcing data on pressure-levels, as a dict where each entry is
+        an array with shape `[level, longitude, latitude]` matching
+        `data_coords`. Single level data (e.g., sea surface temperature) should
+        have a `level` dimension of size 1.
+
+    Returns:
+      Outputs on pressure-levels, as a dict where each entry is an array with
+      shape `[level, longitude, latitude]` matching `data_coords`.
+    """
+    sim_time = _sim_time_from_state(state)
+    forcings = self._squeeze_level_from_forcings(forcings)
+    forcings = _prepend_dummy_time_axis(forcings)
+    f = self._structure.forcing_fn(self.params, None, forcings, sim_time)
+    outputs = self._structure.decode_fn(self.params, None, state, f)
+    outputs = self._from_abbreviated_names_and_tracers(outputs)
+    return outputs
+
+  @functools.partial(
+      jax.jit,
+      static_argnames=[
+          'steps',
+          'timedelta',
+          'start_with_input',
+          'post_process_fn',
+      ],
+  )
+  @_static_gin_config
+  def unroll(
+      self,
+      state: State,
+      forcings: TemporalForcings,
+      *,
+      steps: int,
+      timedelta: TimedeltaLike | None = None,
+      start_with_input: bool = False,
+      post_process_fn: Callable[[State], Any] | None = None,
+  ) -> tuple[State, BatchedOutputs]:
+    """Unroll predictions over many time-steps.
+
+    Usage:
+
+      advanced_state, outputs = model.unroll(state, forcings, steps=N)
+
+    where ``advanced_state`` is the advanced model state after ``N`` steps and
+    ``outputs`` is a trajectory of decoded states on pressure-levels with a
+    leading dimension of size ``N``.
+
+    Args:
+      state: initial model state.
+      forcings: forcing data over the time-period spanned by the desired output
+        trajectory. Should include a leading time-axis, but times can be at any
+        desired granularity (e.g., it should be fine to supply daily forcing
+        data, even if producing hourly outputs). The nearest forcing in time
+        will be used for each internal ``advance()`` and ``decode()`` call.
+      steps: number of time-steps to take.
+      timedelta: size of each time-step to take, which must be a multiple of the
+        internal model timestep. By default uses the internal model timestep.
+      start_with_input: if ``True``, outputs are at times ``[0, ..., (steps - 1)
+        * timestep]`` relative to the initial time; if ``False``, outputs are at
+        times ``[timestep, ..., steps * timestep]``.
+      post_process_fn: optional function to apply to each advanced state and
+        current forcings to create outputs like ``post_process_fn(state,
+        forcings)``, where ``forcings`` does not include a time axis. By
+        default, uses ``model.decode``.
+
+    Returns:
+      A tuple of the advanced state at time ``steps * timestamp``, and outputs
+      with a leading ``time`` axis at the time-steps specified by ``steps``,
+      ``timedelta`` and ``start_with_input``.
+    """
+    if timedelta is None:
+      timedelta = self.timestep
+
+    def get_nearest_forcings(sim_time):
+      times = forcings['sim_time']
+      assert isinstance(times, jax.Array)
+      approx_index = jnp.interp(sim_time, times, jnp.arange(times.size))
+      index = jnp.round(approx_index).astype(jnp.int32)
+      return jax.tree.map(lambda x: x[index, ...], forcings)
+
+    def with_nearest_forcings(func):
+      def wrapped(state):
+        sim_time = _sim_time_from_state(state)
+        forcings = get_nearest_forcings(sim_time)
+        return func(state, forcings)
+
+      return wrapped
+
+    if post_process_fn is None:
+      post_process_fn = self.decode
+
+    inner_steps = _calculate_sub_steps(self.timestep, timedelta)
+    trajectory_func = time_integration.trajectory_from_step(
+        with_nearest_forcings(self.advance),
+        outer_steps=steps,
+        inner_steps=inner_steps,
+        start_with_input=start_with_input,
+        post_process_fn=with_nearest_forcings(post_process_fn),
+    )
+    state, outputs = trajectory_func(state)
+    return state, outputs
+
+  @classmethod
+  def from_checkpoint(cls, checkpoint: Any) -> PressureLevelModel:
+    """Creates a PressureLevelModel from a checkpoint.
+
+    Args:
+      checkpoint: dictionary with keys "model_config_str", "aux_ds_dict" and
+        "params" that specifies model gin configuration, supplemental xarray
+        dataset with model-specific static features, and model parameters.
+
+    Returns:
+      Instance of a `PressureLevelModel` with weights and configuration
+      specified by the checkpoint.
+    """
+    # Hard code radius=1.0 to enable breaking changes in Dinosaur.
+    model_config_str = (
+        checkpoint['model_config_str'].replace(
+            'GridWithWavenumbers.radius = None',
+            'GridWithWavenumbers.radius = 1.0',
+        )
+        + '\n\n'
+        + '\n'.join([
+            'GridTL63.radius = 1.0',
+            'GridTL127.radius = 1.0',
+            'GridTL255.radius = 1.0',
+        ])
+    )
+    with gin_utils.specific_config(model_config_str):
+      physics_specs = physics_specifications.get_physics_specs()
+      aux_ds = xarray.Dataset.from_dict(checkpoint['aux_ds_dict'])
+      data_coords = model_builder.coordinate_system_from_dataset(aux_ds)
+      model_specs = model_builder.get_model_specs(
+          data_coords, physics_specs, {xarray_utils.XARRAY_DS_KEY: aux_ds}
+      )
+      whirl_model = model_builder.WhirlModel(
+          coords=model_specs.coords,
+          dt=model_specs.dt,
+          physics_specs=model_specs.physics_specs,
+          aux_features=model_specs.aux_features,
+          input_coords=data_coords,
+          output_coords=data_coords,
+      )
+      return cls(whirl_model, checkpoint['params'], model_config_str)

model/legacy/correctors.py

@@ -0,0 +1,177 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules that predict refinement or updates of time-advanced states."""
+
+import dataclasses
+from typing import Any, Callable, Optional
+from dinosaur import coordinate_systems
+from dinosaur import time_integration
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+from model.legacy import equations
+from model.legacy import features
+from model.legacy import filters
+from model.legacy import integrators
+from model.legacy import mappings
+
+Pytree = typing.Pytree
+PyTreeState = typing.PyTreeState
+Forcing = typing.Forcing
+
+CorrectorFn = typing.CorrectorFn
+CorrectorModule = typing.CorrectorModule
+EquationModule = equations.EquationModule
+FeaturesModule = features.FeaturesModule
+MappingModule = mappings.MappingModule
+StepModule = typing.StepModule
+StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
+TimeIntegrator = integrators.TimeIntegrator
+TransformModule = typing.TransformModule
+
+
+@gin.register
+class PredictorEulerCorrector(hk.Module):
+  """Corrector that takes Euler step ontop of a predictor step."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      predictor_module: StepModule,
+      filter_module: StepFilterModule = filters.NoFilter,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.dt = dt
+    self.step_fn = predictor_module(coords, dt, physics_specs, aux_features)
+    self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
+
+  def __call__(
+      self,
+      state: typing.PyTreeState,
+      tendencies: typing.PyTreeState,
+      forcing: Optional[Forcing] = None,
+  ) -> typing.PyTreeState:
+    state = self.step_fn(state, forcing)
+    euler_add_fn = lambda x, y: x + self.dt * y if y is not None else x
+    result = jax.tree_util.tree_map(euler_add_fn, state, tendencies)
+    return self.filter_fn(state, result)
+
+
+@gin.register
+class DycoreWithPhysicsCorrector(hk.Module):
+  """Corrector that runs dycore with physics tendencies added to explicit terms.
+
+  This corrector treats predicted physics tendencies constant at each time
+  interval and includes them to all substeps of the dycore step. To achieve this
+  the dycore in this module is specified by the governing equation, rather than
+  an `EquationStep`.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      dycore_equation_module: EquationModule = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      dycore_substeps: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      time_integrator: TimeIntegrator = integrators.imex_rk_sil3,
+      filter_module: StepFilterModule = filters.NoFilter,
+      checkpoint_explicit_terms: bool = True,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    dycore_equation = dycore_equation_module(
+        coords, dt, physics_specs, aux_features)
+    if checkpoint_explicit_terms:
+      dycore_equation = time_integration.ImplicitExplicitODE.from_functions(
+          hk.remat(dycore_equation.explicit_terms),
+          dycore_equation.implicit_terms,
+          dycore_equation.implicit_inverse)  # pyrefly: ignore[bad-argument-type]
+    self.coords = coords
+    self.dycore_equation = dycore_equation
+    self.dycore_substeps = dycore_substeps
+    self.inner_dt = dt / dycore_substeps
+    self.dt = dt
+    self.time_integrator = time_integrator
+    self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
+
+  def __call__(
+      self,
+      state: typing.PyTreeState,
+      tendencies: typing.PyTreeState,
+      forcing: Optional[Forcing] = None,
+  ) -> typing.PyTreeState:
+    state, tendencies = self.coords.with_dycore_sharding((state, tendencies))
+    physics_parametrization_eq = time_integration.ExplicitODE.from_functions(
+        lambda state: tendencies)
+    all_equations = (self.dycore_equation, physics_parametrization_eq)
+    equation = time_integration.compose_equations(all_equations)
+    step_fn = self.time_integrator(equation, self.inner_dt)
+    # TODO(dkochkov) make step_with_filters work with single filter.
+    step_fn = time_integration.step_with_filters(step_fn, [self.filter_fn])
+    step_fn = time_integration.repeated(step_fn, self.dycore_substeps, hk.scan)
+    state = time_integration.maybe_fix_sim_time_roundoff(
+        step_fn(state), self.dt
+    )
+    state = self.coords.with_dycore_sharding(state)
+    return state
+
+
+@gin.register
+class CustomCoordsCorrector(hk.Module):
+  """Corrector module that uses gin-configured coordinates instead of coords.
+
+  This class currently supports model states in spectral representation. It
+  could be easily extended to nodal-state models by converting to modal space
+  prior to spectral interpolation and back after the timestep if performed.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      corrector_module: CorrectorModule = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      custom_coords: coordinate_systems.CoordinateSystem = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    custom_coords = dataclasses.replace(
+        custom_coords, spmd_mesh=coords.spmd_mesh
+    )
+    self.corrector_fn = corrector_module(
+        custom_coords, dt, physics_specs, aux_features)
+    self.to_custom_coords_fn = coordinate_systems.get_spectral_interpolate_fn(
+        coords, custom_coords)
+    self.from_custom_coords_fn = coordinate_systems.get_spectral_interpolate_fn(
+        custom_coords, coords)
+
+  def __call__(
+      self,
+      state: typing.PyTreeState,
+      tendencies: typing.PyTreeState,
+      forcing: Optional[Forcing] = None,
+  ) -> typing.PyTreeState:
+    state = self.to_custom_coords_fn(state)
+    tendencies = self.to_custom_coords_fn(tendencies)
+    # TODO(dkochkov) Consider adding forcing interpolated to custom coords.
+    custom_out = self.corrector_fn(state, tendencies, None)
+    return self.from_custom_coords_fn(custom_out)

model/legacy/decoders.py

@@ -0,0 +1,749 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines `decoder` modules that map model state to output data format."""
+
+import functools
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar
+import zlib
+
+from dinosaur import coordinate_systems
+from dinosaur import primitive_equations
+from dinosaur import pytree_utils
+from dinosaur import scales
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+from dinosaur import vertical_interpolation
+from dinosaur import weatherbench_utils
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import diagnostics
+from model.legacy import features
+from model.legacy import filters
+from model.legacy import mappings
+from model.legacy import orographies
+from model.legacy import perturbations
+from model.legacy import stochastic
+from model.legacy import transforms
+import numpy as np
+
+
+# long lines are better than splitting argument definitions onto two lines
+# pylint: disable=line-too-long
+
+# We ♥ λ's
+# pylint: disable=g-long-lambda
+
+DataState = typing.DataState
+DiagnosticModule = diagnostics.DiagnosticModule
+FeaturesModule = features.FeaturesModule
+FilterModule = Callable[..., typing.PyTreeFilterFn]
+Forcing = typing.Forcing
+MappingModule = mappings.MappingModule
+PyTreeState = typing.PyTreeState
+ModelState = typing.ModelState
+TransformModule = typing.TransformModule
+OrographyModule = orographies.OrographyModule
+PerturbationModule = perturbations.PerturbationModule
+RandomnessModule = stochastic.RandomnessModule
+
+
+@gin.register
+class DecoderIdentityTransform(hk.Module):
+  """Transformation that returns inputs without modification."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del coords, dt, physics_specs, aux_features, output_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    return inputs
+
+
+@gin.register
+class DecoderFilterTransform(hk.Module):
+  """Transformation that returns truncated and filtered modal inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      filter_module: FilterModule = filters.DataNoFilter,
+      return_nodal: bool = True,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.output_coords = output_coords
+    self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
+    self.return_nodal = return_nodal
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    modal_inputs = coordinate_systems.maybe_to_modal(inputs, self.output_coords)
+    filtered_inputs = self.filter_fn(modal_inputs)
+    if self.return_nodal:
+      return self.output_coords.horizontal.to_nodal(filtered_inputs)
+    return filtered_inputs
+
+
+@gin.register
+class OutputModalToModalTransform(hk.Module):
+  """Transformation that truncates modal state to output coords."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.output_coords = output_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    downsample_fn = coordinate_systems.get_spectral_downsample_fn(
+        self.coords, self.output_coords
+    )
+    return downsample_fn(inputs)
+
+
+@gin.register
+class OutputModalToNodalTransform(hk.Module):
+  """Transformation that converts modal state to nodal representation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.output_coords = output_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    to_nodal_fn = self.output_coords.horizontal.to_nodal
+    downsample_fn = coordinate_systems.get_spectral_downsample_fn(
+        self.coords, self.output_coords
+    )
+    return jax.tree_util.tree_map(
+        lambda x: to_nodal_fn(downsample_fn(x)), inputs
+    )
+
+
+@gin.register
+class OutputNodalToModalTransform(hk.Module):
+  """Transformation that converts nodal state to modal representation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.output_coords = output_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    return self.output_coords.horizontal.to_modal(inputs)
+
+
+@gin.register
+class ModalOutputLearnedAdaptorTransform(hk.Module):
+  """Transformation using a tower to adapt modal outputs to the data domain."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      output_transform_module: TransformModule,
+      name: Optional[str] = None,
+  ):
+    del output_coords  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.modal_to_nodal_features_fn = modal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, coords
+    )
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    """Applies transform to modal inputs, returns modal outputs."""
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    prediction_shapes = jax.tree_util.tree_map(self.get_nodal_shape_fn, inputs)
+    # if `inputs` contain `sim_time` - remove it from corrections.
+    sim_time_shape = prediction_shapes.pop('sim_time', None)
+    net = self.nodal_mapping_module(prediction_shapes)
+    nodal_input_features = self.modal_to_nodal_features_fn(inputs, None)
+    nodal_corrections = self.output_transform_fn(net(nodal_input_features))
+    corrections = self.coords.horizontal.to_modal(nodal_corrections)
+    if sim_time_shape is not None:
+      corrections['sim_time'] = 0.0
+    outputs = jax.tree_util.tree_map(lambda x, y: x + y, inputs, corrections)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+class NodalOutputLearnedAdaptorTransform(hk.Module):
+  """Transformation using a tower to adapt nodal outputs to the data domain."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      nodal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      output_transform_module: TransformModule,
+      name: Optional[str] = None,
+  ):
+    del output_coords  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.nodal_to_nodal_features_fn = nodal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, coords
+    )
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    """Applies transform to nodal inputs, returns nodal outputs."""
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    prediction_shapes = jax.tree_util.tree_map(self.get_nodal_shape_fn, inputs)
+    # if `inputs` contain `sim_time` - remove it from corrections.
+    sim_time_shape = prediction_shapes.pop('sim_time', None)
+    net = self.nodal_mapping_module(prediction_shapes)
+    input_features = self.nodal_to_nodal_features_fn(inputs, None)
+    corrections = self.output_transform_fn(net(input_features))
+    if sim_time_shape is not None:
+      corrections['sim_time'] = 0.0
+    outputs = jax.tree_util.tree_map(lambda x, y: x + y, inputs, corrections)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+class DecoderCombinedTransform(hk.Module):
+  """Module that applies multiple transformations sequentially."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      output_coords: coordinate_systems.CoordinateSystem,
+      transforms: Tuple[TransformModule, ...],  # pylint: disable=redefined-outer-name
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.transform_fns = [
+        module(coords, dt, physics_specs, aux_features, output_coords)
+        for module in transforms
+    ]
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    for transform_fn in self.transform_fns:
+      inputs = transform_fn(inputs)
+    return inputs
+
+
+@gin.register
+class IdentityDecoder(hk.Module):
+  """Decoder that returns model state unaltered."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features, output_coords
+    super().__init__(name=name)
+
+  def __call__(self, x: ModelState, forcing: Forcing) -> DataState:
+    del forcing
+    return x.state
+
+
+@gin.register
+class StateToDictDecoder(hk.Module):
+  """Decoder that returns a dict representation of a model state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+
+  def __call__(self, x: ModelState, forcing: Forcing) -> DataState:
+    del forcing
+    state_dict, _ = pytree_utils.as_dict(x.state)
+    return self.transform_fn(state_dict)
+
+
+@gin.register
+class LeapfrogSliceDecoder(hk.Module):
+  """Decoder that returns one slice out of a leapfrog pair."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      slice_id: int = 0,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_id = slice_id
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+
+  def __call__(self, x: ModelState, forcing: Forcing) -> DataState:
+    del forcing
+    return self.transform_fn(x.state[self.slice_id])
+
+
+@gin.register
+class LeapfrogSliceDictDecoder(hk.Module):
+  """Decoder that returns one slice out of a leapfrog pair as dictionary."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      slice_id: int = 0,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_id = slice_id
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+
+  def __call__(self, x: ModelState, forcing: Forcing) -> DataState:
+    del forcing
+    state_dict, _ = pytree_utils.as_dict(x.state[self.slice_id])
+    return self.transform_fn(state_dict)
+
+
+@gin.configurable
+class PrimitiveToWeatherbenchDecoder(hk.Module):
+  """Decoder that converts `StateWithTime` to  `weatherbench.State`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    ref_temps = aux_features[xarray_utils.REF_TEMP_KEY]
+    self.ref_temps = ref_temps[..., np.newaxis, np.newaxis]
+    self.output_coords = output_coords
+    self.coords = coords
+    self.physics_specs = physics_specs
+    self.velocity_fn = functools.partial(
+        spherical_harmonic.vor_div_to_uv_nodal,
+        output_coords.horizontal,
+    )
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features
+    )
+    orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    self.nodal_orography = coords.horizontal.to_nodal(orography)
+    self.geopotential_fn = functools.partial(
+        primitive_equations.get_geopotential_with_moisture,
+        nodal_orography=self.nodal_orography,
+        coordinates=coords.vertical,
+        gravity_acceleration=physics_specs.gravity_acceleration,
+        ideal_gas_constant=physics_specs.ideal_gas_constant,
+        water_vapor_gas_constant=physics_specs.water_vapor_gas_constant,
+    )
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+
+  def primitive_to_weatherbench(
+      self,
+      inputs: primitive_equations.StateWithTime,
+  ) -> weatherbench_utils.State:
+    """Converts pe_state to weatherbench state on pressure levels."""
+    # output state is computed on output_coords.
+    to_nodal_fn = self.output_coords.horizontal.to_nodal
+    u, v = self.velocity_fn(  # returned in nodal space.
+        vorticity=inputs.vorticity, divergence=inputs.divergence
+    )
+    t = self.ref_temps + to_nodal_fn(inputs.temperature_variation)
+    tracers = to_nodal_fn(inputs.tracers)
+    z = self.geopotential_fn(t, tracers['specific_humidity'])
+    surface_pressure = jnp.exp(to_nodal_fn(inputs.log_surface_pressure))
+    u, v, t, z, tracers, surface_pressure = (
+        self.coords.dycore_to_physics_sharding(
+            (u, v, t, z, tracers, surface_pressure)
+        )
+    )
+    interpolate_with_linear_extrap_fn = (
+        vertical_interpolation.vectorize_vertical_interpolation(
+            vertical_interpolation.linear_interp_with_linear_extrap
+        )
+    )
+    interpolate_with_constant_extrap_fn = (
+        vertical_interpolation.vectorize_vertical_interpolation(
+            vertical_interpolation.vertical_interpolation
+        )
+    )
+    regrid_with_linear_fn = functools.partial(
+        vertical_interpolation.interp_sigma_to_pressure,
+        pressure_coords=self.output_coords.vertical,
+        sigma_coords=self.coords.vertical,
+        surface_pressure=surface_pressure,
+        interpolate_fn=interpolate_with_linear_extrap_fn,
+    )
+    regrid_with_constant_fn = functools.partial(
+        vertical_interpolation.interp_sigma_to_pressure,
+        pressure_coords=self.output_coords.vertical,
+        sigma_coords=self.coords.vertical,
+        surface_pressure=surface_pressure,
+        interpolate_fn=interpolate_with_constant_extrap_fn,
+    )
+    # closes regridding options based on http://shortn/_X09ZAU1jsx.
+    # use constant extrapolation for `u, v, tracers`.
+    # use linear extrapolation for `z, t`.
+    return weatherbench_utils.State(
+        u=regrid_with_constant_fn(u),  # pyrefly: ignore[unexpected-keyword]
+        v=regrid_with_constant_fn(v),  # pyrefly: ignore[unexpected-keyword]
+        t=regrid_with_linear_fn(t),  # pyrefly: ignore[unexpected-keyword]
+        z=regrid_with_linear_fn(z),  # pyrefly: ignore[unexpected-keyword]
+        sim_time=inputs.sim_time,  # pyrefly: ignore[unexpected-keyword]
+        tracers=regrid_with_constant_fn(tracers),  # pyrefly: ignore[unexpected-keyword]
+    )
+
+  def __call__(
+      self, inputs: ModelState, forcing: Forcing
+  ) -> DataState:
+    del forcing
+    wb_on_sigma = self.primitive_to_weatherbench(inputs.state)
+    return self.transform_fn(wb_on_sigma.asdict())  # pyrefly: ignore[missing-attribute]
+
+
+_DECODER_SALT = zlib.crc32(b'decoder')  # arbitrary uint32 value
+
+
+def _decoder_prng_key(
+    randomness: typing.RandomnessState,
+) -> typing.PRNGKeyArray | None:
+  """Get a PRNG Key suitable for decoder randomness."""
+  if randomness.prng_key is None:
+    return None
+  salt = jnp.uint32(_DECODER_SALT) + jnp.uint32(randomness.prng_step)
+  return jax.random.fold_in(randomness.prng_key, salt)
+
+
+@gin.register
+class LearnedPrimitiveToWeatherbenchDecoder(PrimitiveToWeatherbenchDecoder):
+  """Similar to `PrimitiveToWeatherbenchDecoder` with learned interpolation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      modal_to_nodal_model_features_module: FeaturesModule,
+      modal_to_nodal_data_features_module: FeaturesModule,
+      correction_transform_module: TransformModule,
+      nodal_mapping_module: MappingModule,
+      prediction_mask: typing.Pytree,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      diagnostics_module: DiagnosticModule = diagnostics.NoDiagnostics,
+      name: Optional[str] = None,
+  ):
+    super().__init__(
+        coords=coords,
+        dt=dt,
+        physics_specs=physics_specs,
+        aux_features=aux_features,
+        output_coords=output_coords,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        name=name,
+    )  # don't pass the transform, as we apply it at the end.
+    self.prediction_mask = prediction_mask
+    # features are computed on both coordinate systems.
+    self.model_features_fn = modal_to_nodal_model_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.data_features_fn = modal_to_nodal_data_features_module(
+        output_coords, dt, physics_specs, aux_features
+    )
+    self.corrections_transform_fn = correction_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    # corrections are computed in real space on output coordinates.
+    self.nodal_mapping_module = nodal_mapping_module
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, output_coords
+    )
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+    self.randomness_fn = randomness_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.perturbation_fn = perturbation_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.diagnostic_fn = diagnostics_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self, inputs: ModelState, forcing: Forcing
+  ) -> DataState:
+    randomness = self.randomness_fn.unconditional_sample(
+        _decoder_prng_key(inputs.randomness)
+    )
+    prognostics = self.perturbation_fn(
+        inputs=self.coords.with_dycore_sharding(inputs.state),
+        state=None,
+        randomness=self.coords.with_dycore_sharding(randomness.nodal_value),
+    )
+    inputs.state = prognostics  # compute diagnostics from the perturbed state.
+
+    # TODO(dkochkov) Could we pass physics_tendencies here?
+    # TODO(janniyuval) Consider using evaporation diagnostics for training.
+    decoder_diagnostics = self.diagnostic_fn(inputs, None)
+    wb_on_pressure_dict = self.primitive_to_weatherbench(prognostics).asdict()  # pyrefly: ignore[missing-attribute]
+    wb_on_pressure_modal = coordinate_systems.maybe_to_modal(
+        self.coords.physics_to_dycore_sharding(wb_on_pressure_dict),
+        self.output_coords,
+    )
+    wb_on_pressure_dict['diagnostics'] = decoder_diagnostics
+    prediction_mask = pytree_utils.replace_with_matching_or_default(
+        wb_on_pressure_dict, self.prediction_mask, default=False)
+    prediction_shapes = jax.tree_util.tree_map(
+        lambda x, y: self.get_nodal_shape_fn(x) if y else None,
+        wb_on_pressure_dict,
+        prediction_mask,
+    )
+    net = self.nodal_mapping_module(prediction_shapes)
+    model_features = self.model_features_fn(
+        prognostics.asdict(), forcing=forcing,
+        randomness=randomness.nodal_value
+    )
+    data_features = self.data_features_fn(wb_on_pressure_modal, forcing=forcing)
+    data_features = transforms.add_prefix(data_features, 'data_')
+    model_features = transforms.add_prefix(model_features, 'model_')
+    all_features = self.coords.dycore_to_physics_sharding(
+        data_features | model_features
+    )
+
+    nodal_outputs = self.corrections_transform_fn(net(all_features))
+    add_fn = lambda x, y: x + y if y is not None else x
+    wb_on_pressure_dict = jax.tree_util.tree_map(
+        add_fn, wb_on_pressure_dict, nodal_outputs
+    )
+    return self.transform_fn(wb_on_pressure_dict)
+
+
+@gin.register
+class DimensionalPrimitiveToWeatherbenchDecoder(PrimitiveToWeatherbenchDecoder):
+  """Same as PrimitiveToWeatherbenchDecoder, but with dimensional output."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    nondim_pressure_centers = physics_specs.nondimensionalize(
+        output_coords.vertical.centers * scales.units.millibar
+    )
+    nondim_output_coords = coordinate_systems.CoordinateSystem(
+        output_coords.horizontal,
+        vertical_interpolation.PressureCoordinates(nondim_pressure_centers),
+        spmd_mesh=output_coords.spmd_mesh,
+    )
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        output_coords=nondim_output_coords,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        transform_module=transform_module,
+        name=name,
+    )
+    self.redimensionalize_fn = transforms.RedimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        output_coords=output_coords,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+
+  def __call__(
+      self, inputs: ModelState, forcing: Forcing
+  ) -> DataState:
+    return self.redimensionalize_fn(super().__call__(inputs, forcing))
+
+
+@gin.configurable
+class DimensionalLearnedPrimitiveToWeatherbenchDecoder(
+    LearnedPrimitiveToWeatherbenchDecoder
+):
+  """Same as LearnedPrimitiveToWeatherbenchDecoder, but with dimensional output."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      output_coords: coordinate_systems.CoordinateSystem,
+      modal_to_nodal_model_features_module: FeaturesModule,
+      modal_to_nodal_data_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      correction_transform_module: TransformModule,
+      prediction_mask: typing.Pytree,
+      inputs_to_units_mapping: Dict[str, str],
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = DecoderIdentityTransform,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      diagnostics_module: DiagnosticModule = diagnostics.NoDiagnostics,
+      name: Optional[str] = None,
+  ):
+    nondim_pressure_centers = physics_specs.nondimensionalize(
+        output_coords.vertical.centers * scales.units.millibar
+    )
+    nondim_output_coords = coordinate_systems.CoordinateSystem(
+        output_coords.horizontal,
+        vertical_interpolation.PressureCoordinates(nondim_pressure_centers),
+        spmd_mesh=output_coords.spmd_mesh,
+    )
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        output_coords=nondim_output_coords,
+        modal_to_nodal_model_features_module=(
+            modal_to_nodal_model_features_module
+        ),
+        modal_to_nodal_data_features_module=modal_to_nodal_data_features_module,
+        nodal_mapping_module=nodal_mapping_module,
+        correction_transform_module=correction_transform_module,
+        prediction_mask=prediction_mask,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        transform_module=transform_module,
+        randomness_module=randomness_module,
+        perturbation_module=perturbation_module,
+        diagnostics_module=diagnostics_module,
+        name=name,
+    )
+    self.redimensionalize_fn = transforms.RedimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        output_coords=output_coords,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+
+  def __call__(
+      self, inputs: ModelState, forcing: Forcing
+  ) -> DataState:
+    return self.redimensionalize_fn(super().__call__(inputs, forcing))

model/legacy/diagnostics.py

@@ -0,0 +1,420 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines `diagnostic` modules that compute diagnostic predictions."""
+
+from collections import abc
+from typing import Any, Callable, Optional, Protocol
+
+from dinosaur import coordinate_systems
+from dinosaur import scales
+from dinosaur import sigma_coordinates
+from dinosaur import typing
+
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+TransformModule = typing.TransformModule
+
+PRECIPITATION = 'precipitation'
+EVAPORATION = 'evaporation'
+
+
+class DiagnosticFn(Protocol):
+  """Implements initialization and computation of model diagnostic fields."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+  ):
+    del coords, dt, physics_specs, aux_features
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> dict[str, jax.Array]:
+    """Computes diagnostic field from `model_state` and `physics_tendencies`."""
+    ...
+
+
+DiagnosticModule = Callable[..., DiagnosticFn]
+
+
+@gin.register
+class NoDiagnostics:
+  """Diagnostic module that computes no diagnostics."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+  ):
+    del coords, dt, physics_specs, aux_features
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> dict[str, jax.Array]:
+    return {}
+
+
+@gin.register
+class CombinedDiagnostics:
+  """Computes a combination of multiple diagnostics."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+      diagnostic_modules: abc.Sequence[DiagnosticModule] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+  ):
+    self.diagnostic_fns = [
+        module(coords, dt, physics_specs, aux_features)
+        for module in diagnostic_modules
+    ]
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> dict[str, jax.Array]:
+    diagnostics = {}
+    for fn in self.diagnostic_fns:
+      new_diagnostics = fn(model_state, physics_tendencies, forcing)
+      if any(k in diagnostics for k in new_diagnostics):
+        raise ValueError(
+            f'{new_diagnostics.keys()} overlaps with {diagnostics.keys()}'
+        )
+      diagnostics.update(new_diagnostics)
+    return diagnostics
+
+
+@gin.register
+class PrecipitationMinusEvaporationDiagnostics:
+  """Computes `P-E` by integrating physics_tendencies.
+
+  Depending on the `method` computes either precipitation minus evaporation
+  rate, which in ERA5 has units `kg m**-2 s**-1` or time-accumulated value
+  in `kg m**-2` if `method == cumulative`.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+      moisture_species: tuple[str, ...] = (
+          'specific_humidity',
+          'specific_cloud_ice_water_content',
+          'specific_cloud_liquid_water_content',
+      ),
+      method: str = 'rate',
+  ):
+    del aux_features
+    self.coords = coords
+    self.dt = dt
+    self.physics_specs = physics_specs
+    self.moisture_species = moisture_species
+    self.method = method
+    self.to_nodal_fn = coords.horizontal.to_nodal
+
+  def _compute_evaporation_minus_precipitation(
+      self, model_state: typing.ModelState, physics_tendencies: typing.Pytree
+  ) -> typing.Array:
+    """Computes evaporation minus precipitation."""
+    lsp = model_state.state.log_surface_pressure
+    p_surface = jnp.squeeze(jnp.exp(self.to_nodal_fn(lsp)), axis=0)
+    moisture_tendencies = [
+        v
+        for tracer, v in physics_tendencies.tracers.items()
+        if tracer in self.moisture_species
+    ]
+    moisture_tendencies = sum(self.to_nodal_fn(moisture_tendencies))
+    scale = p_surface / self.physics_specs.g
+    e_minus_p = scale * sigma_coordinates.sigma_integral(
+        moisture_tendencies, self.coords.vertical, keepdims=False
+    )
+    return e_minus_p
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> typing.Pytree:
+    """Computes precipitation minus evaporation."""
+    del forcing  # unused
+    e_minus_p = self._compute_evaporation_minus_precipitation(
+        model_state, physics_tendencies
+    )
+    if self.method == 'rate':
+      return {'P_minus_E_rate': -e_minus_p}
+    elif self.method == 'cumulative':
+      # TODO(dkochkov) Address possible precision loss due to small deltas.
+      surface_nodal_shape = self.coords.horizontal.nodal_shape
+      previous = model_state.diagnostics.get(
+          'P_minus_E_cumulative',
+          jnp.zeros(surface_nodal_shape))
+      return {'P_minus_E_cumulative': previous - (e_minus_p * self.dt)}
+    else:
+      raise ValueError(f'Unknown {self.method=}, must be `rate`/`cumulative`')
+
+
+@gin.register
+class PrecipitableWaterDiagnostics:
+  """Computes cumulative preciptable water in the state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+      moisture_species: tuple[str, ...] = (
+          'specific_humidity',
+          'specific_cloud_ice_water_content',
+          'specific_cloud_liquid_water_content',
+      ),
+  ):
+    del dt, aux_features
+    self.coords = coords
+    self.physics_specs = physics_specs
+    self.moisture_species = moisture_species
+    self.to_nodal_fn = coords.horizontal.to_nodal
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> typing.Pytree:
+    """Computes preciptable water."""
+    del physics_tendencies, forcing  # unused
+    lsp = model_state.state.log_surface_pressure
+    p_surface = jnp.squeeze(jnp.exp(self.to_nodal_fn(lsp)), axis=0)
+    moisture_tracers = [
+        v
+        for tracer, v in model_state.tracers.items()  # pyrefly: ignore[missing-attribute]
+        if tracer in self.moisture_species
+    ]
+    moisture = sum(self.to_nodal_fn(moisture_tracers))
+    water_density = self.physics_specs.nondimensionalize(scales.WATER_DENSITY)
+    scale = p_surface / (self.physics_specs.g * water_density)
+    water = scale * sigma_coordinates.sigma_integral(
+        moisture, self.coords.vertical, keepdims=False
+    )
+    return {'precipitable_water': water}
+
+
+@gin.register
+class NodalModelDiagnosticsDecoder:
+  """Diagnostics decoder that returns elements from model_state.diagnostics."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+  ):
+    del dt, aux_features
+    self.coords = coords
+    self.physics_specs = physics_specs
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> typing.Pytree:
+    """Computes precipitation minus evaporation."""
+    del physics_tendencies, forcing  # unused.
+    nodal_diagnostics = coordinate_systems.maybe_to_nodal(
+        model_state.diagnostics, self.coords
+    )
+    return nodal_diagnostics
+
+
+# TODO(janniyuval) add a decoder that can add some Gaussian noise to evap/precip
+@gin.register
+class PrecipitationDiagnosticsConstrained(
+    hk.Module, PrecipitationMinusEvaporationDiagnostics
+):
+  """Predict evaporation and computes cumulative precipitation.
+
+  Calculation is based on calculating `P-E` by integrating physics_tendencies.
+  Depending on the `method` computes either precipitation
+  rate, (which in ERA5 has units `kg m**-2 s**-1`) or time-accumulated value
+  in `Length` units (GPCP uses mm/day) if `method == cumulative`.
+  Evaporation has the units of `kg m**-2 s**-1` in ERA5.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+      embedding_module: typing.EmbeddingModule,
+      moisture_species: tuple[str, ...] = (
+          'specific_humidity',
+          'specific_cloud_ice_water_content',
+          'specific_cloud_liquid_water_content',
+      ),
+      is_precipitation: bool = True,
+      method_precipitation: str = 'cumulative',
+      method_evaporation: str = 'rate',
+      name: Optional[str] = None,
+      field_name: str = 'total_precipitation',
+  ):
+    # del aux_features
+    super().__init__(name=name)
+    self.coords = coords
+    self.dt = dt
+    self.physics_specs = physics_specs
+    self.moisture_species = moisture_species
+    self.method_precipitation = method_precipitation
+    self.method_evaporation = method_evaporation
+    self.to_nodal_fn = coords.horizontal.to_nodal
+    self.is_precipitation = is_precipitation
+    if self.is_precipitation:
+      predicted_name = PRECIPITATION
+      diagnosed_name = EVAPORATION
+    else:
+      predicted_name = EVAPORATION
+      diagnosed_name = PRECIPITATION
+
+    self.predicted_name = predicted_name
+    self.diagnosed_name = diagnosed_name
+
+    output_shapes = {
+        f'{predicted_name}': np.asarray(coords.surface_nodal_shape)
+    }
+
+    self.embedding_fn = embedding_module(
+        coords, dt, physics_specs, aux_features, output_shapes=output_shapes
+    )
+    self.water_density = self.physics_specs.nondimensionalize(
+        scales.WATER_DENSITY
+    )
+    self.field_name = field_name
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> typing.Pytree:
+    """Computes precipitation minus evaporation."""
+    e_minus_p = self._compute_evaporation_minus_precipitation(
+        model_state, physics_tendencies
+    )
+    water_budget = self.embedding_fn(
+        model_state.state,
+        model_state.memory,
+        model_state.diagnostics,
+        model_state.randomness,
+        forcing,
+    )
+    water_budget[self.diagnosed_name] = (
+        -e_minus_p - water_budget[self.predicted_name]
+    )
+
+    # Note: In ERA5 mean_evaporation_rate (kg m**-2 s**-1)
+    # is negative for evaporation.
+    # In GPCP precipitation is positive (mm/day).
+    # Here e_minus_p is positive for evaporation.
+    output_dict = {}
+    surface_nodal_shape = self.coords.horizontal.nodal_shape
+    if self.method_precipitation == 'rate':  # units: length/time
+      output_dict[PRECIPITATION + '_rate'] = (
+          water_budget[PRECIPITATION]
+      ) / self.water_density
+    elif self.method_precipitation == 'cumulative':  # units: length
+      previous = model_state.diagnostics.get(
+          self.field_name, jnp.zeros(surface_nodal_shape)
+      )
+      # TODO(janniyuval) remove precipitation_cumulative_mean once no models
+      # use it.
+      assert self.field_name in [
+          'total_precipitation',
+          'precipitation_cumulative_mean',
+      ], self.field_name
+      output_dict[self.field_name] = previous + (
+          (water_budget[PRECIPITATION] / self.water_density) * self.dt
+      )
+    else:
+      raise ValueError(
+          f'Precipitation method is {self.method_precipitation=}, but it must'
+          ' be `rate`/`cumulative`'
+      )
+    if self.method_evaporation == 'rate':  # units: mass length**-2 time**-1
+      output_dict[EVAPORATION] = water_budget[EVAPORATION]
+    elif self.method_evaporation == 'cumulative':  # units: length
+      previous_evap = model_state.diagnostics.get(
+          EVAPORATION + '_cumulative', jnp.zeros(surface_nodal_shape)
+      )
+      output_dict[EVAPORATION + '_cumulative'] = (
+          previous_evap
+          + (water_budget[EVAPORATION] / self.water_density) * self.dt
+      )
+    else:
+      raise ValueError(
+          f'Evaporation method is {self.method_evaporation=},  but it must be'
+          ' `rate`/`cumulative`'
+      )
+    return output_dict
+
+
+@gin.register
+class SurfacePressureDiagnostics:
+  """Getting the surface pressure of the state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: dict[str, Any],
+  ):
+    del dt, aux_features, physics_specs
+    self.to_nodal_fn = coords.horizontal.to_nodal
+
+  def __call__(
+      self,
+      model_state: typing.ModelState,
+      physics_tendencies: typing.Pytree,
+      forcing: typing.Forcing | None = None,
+  ) -> typing.Pytree:
+    """Computes surface pressure."""
+    del physics_tendencies, forcing  # unused
+    lsp = model_state.state.log_surface_pressure
+    surface_pressure = jnp.squeeze(jnp.exp(self.to_nodal_fn(lsp)), axis=0)
+    return {'surface_pressure': surface_pressure}

model/legacy/embeddings.py

@@ -0,0 +1,380 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules that predict an embedding from the model state."""
+from typing import Any, Optional
+from dinosaur import coordinate_systems
+from dinosaur import pytree_utils
+from dinosaur import scales
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import features
+from model.legacy import mappings
+from model.legacy import transforms
+
+EmbeddingFn = typing.EmbeddingFn
+EmbeddingModule = typing.EmbeddingModule
+Forcing = typing.Forcing
+TransformModule = typing.TransformModule
+
+units = scales.units
+
+
+@gin.register
+class ModalToNodalEmbedding(hk.Module):
+  """Embedding that expects modal state input and returns nodal output."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      output_shapes: typing.Pytree,
+      modal_to_nodal_features_module: features.FeaturesModule,
+      nodal_mapping_module: mappings.MappingModule,
+      output_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.output_shapes = output_shapes
+    self.modal_to_nodal_features_fn = modal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      state: typing.Pytree,
+      memory: Optional[typing.Pytree] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.Pytree] = None,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> typing.Pytree:
+    """Returns the embedding output on nodal locations."""
+    net = self.nodal_mapping_module(self.output_shapes)
+    # Need to check if dict when embedding is not within the parameterization
+    # (e.g., for diagnostic NN)
+    state, _ = pytree_utils.as_dict(state)
+    nodal_inputs = self.modal_to_nodal_features_fn(
+        state, memory, diagnostics, randomness, forcing
+    )
+    nodal_outputs = self.output_transform_fn(net(nodal_inputs))
+    return nodal_outputs
+
+
+# TODO(pnorgaard) Refactor default embeddings to separate object
+@gin.register
+class NodalSurfaceModelEmbedding(hk.Module):
+  """Embedding to represent a nodal space surface model."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      output_shapes: typing.Pytree,
+      static_vars_ds_path: str,
+      land_embedding: Optional[EmbeddingModule] = None,
+      sea_embedding: Optional[EmbeddingModule] = None,
+      sea_ice_embedding: Optional[EmbeddingModule] = None,
+      snow_embedding: Optional[EmbeddingModule] = None,
+      output_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.output_shapes = output_shapes
+
+    # Basic surface embedding settings
+    self.feature_axis = -3
+    param_init = hk.initializers.TruncatedNormal()
+    output_size = sum([x[self.feature_axis]
+                       for x in jax.tree_util.tree_leaves(output_shapes)])
+    param_shape = (output_size, 1, 1)  # uniform across lon, lat
+    surface_nodal_shape = self.coords.surface_nodal_shape
+
+    if land_embedding is not None:
+      self.land_embedding_fn = land_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.land_parameters = hk.get_parameter(
+          'land_params', param_shape,
+          jnp.float32, init=param_init)
+      def land_embedding_fn(state, memory, randomness, forcing):
+        del state, memory, randomness, forcing  # unused
+        outputs = self.land_parameters * jnp.ones(surface_nodal_shape)
+        return pytree_utils.unpack_to_pytree(
+            outputs, self.output_shapes, self.feature_axis
+        )
+      self.land_embedding_fn = land_embedding_fn
+
+    if sea_embedding is not None:
+      self.sea_embedding_fn = sea_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.sea_parameters = hk.get_parameter(
+          'sea_params', param_shape,
+          jnp.float32, init=param_init)
+      def sea_embedding_fn(state, memory, randomness, forcing):
+        del state, memory, randomness, forcing  # unused
+        outputs = self.sea_parameters * jnp.ones(surface_nodal_shape)
+        return pytree_utils.unpack_to_pytree(
+            outputs, self.output_shapes, self.feature_axis
+        )
+      self.sea_embedding_fn = sea_embedding_fn
+
+    if sea_ice_embedding is not None:
+      self.sea_ice_embedding_fn = sea_ice_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.sea_ice_parameters = hk.get_parameter(
+          'sea_ice_params', param_shape,
+          jnp.float32, init=param_init)
+      def sea_ice_embedding_fn(state, memory, randomness, forcing):
+        del state, memory, randomness, forcing  # unused
+        outputs = self.sea_ice_parameters * jnp.ones(surface_nodal_shape)
+        return pytree_utils.unpack_to_pytree(
+            outputs, self.output_shapes, self.feature_axis
+        )
+      self.sea_ice_embedding_fn = sea_ice_embedding_fn
+
+    if snow_embedding is not None:
+      self.snow_embedding_fn = snow_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.snow_parameters = hk.get_parameter(
+          'snow_params', param_shape,
+          jnp.float32, init=param_init)
+      def snow_embedding_fn(state, memory, randomness, forcing):
+        del state, memory, randomness, forcing  # unused
+        outputs = self.snow_parameters * jnp.ones(surface_nodal_shape)
+        return pytree_utils.unpack_to_pytree(
+            outputs, self.output_shapes, self.feature_axis
+        )
+      self.snow_embedding_fn = snow_embedding_fn
+
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+    ds = xarray_utils.ds_from_path_or_aux(static_vars_ds_path, aux_features)
+    self.land_sea_mask = xarray_utils.nodal_land_sea_mask_from_ds(ds)
+
+    # snow data is provided as depth (in meters). It is converted to snow_cover
+    # by choosing a threshold such that snow_cover = 0 below that value and
+    # snow cover = 1 above that value.
+    self.snow_cover_threshold = physics_specs.nondimensionalize(1 * units.meter)  # pyrefly: ignore[unsupported-operation]
+
+  def __call__(
+      self,
+      state: typing.Pytree,
+      memory: Optional[typing.Pytree] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.Pytree] = None,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> typing.Pytree:
+    """Returns the embedding output on nodal locations."""
+    land_outputs = self.land_embedding_fn(
+        state, memory, diagnostics, randomness, forcing)  # pyrefly: ignore[bad-argument-count]
+    sea_outputs = self.sea_embedding_fn(
+        state, memory, diagnostics, randomness, forcing)  # pyrefly: ignore[bad-argument-count]
+    sea_ice_outputs = self.sea_ice_embedding_fn(
+        state, memory, diagnostics, randomness, forcing  # pyrefly: ignore[bad-argument-count]
+    )
+    snow_outputs = self.snow_embedding_fn(
+        state, memory, diagnostics, randomness, forcing)  # pyrefly: ignore[bad-argument-count]
+
+    # prepare masks with fractional values in [0, 1]
+    land_fraction = self.land_sea_mask
+    sea_fraction = 1 - land_fraction
+    sea_ice_fraction = forcing[xarray_utils.SEA_ICE_COVER]  # pyrefly: ignore[unsupported-operation]
+    snow_fraction = forcing[xarray_utils.SNOW_DEPTH] > self.snow_cover_threshold  # pyrefly: ignore[unsupported-operation]
+
+    # weight and combine outputs
+    snow_weight = snow_fraction * land_fraction  # snow covered land
+    land_weight = (1 - snow_fraction) * land_fraction  # land without snow
+    sea_ice_weight = sea_ice_fraction * sea_fraction  # ice covered sea
+    sea_weight = (1 - sea_ice_fraction) * sea_fraction  # sea without ice
+
+    def tree_scale(a, x):
+      # Multiply leaves of `x` by `a`.
+      return jax.tree_util.tree_map(lambda y: a * y, x)
+
+    surface_outputs = jax.tree_util.tree_map(
+        lambda a, b, c, d: a + b + c + d,
+        tree_scale(land_weight, land_outputs),
+        tree_scale(sea_weight, sea_outputs),
+        tree_scale(sea_ice_weight, sea_ice_outputs),
+        tree_scale(snow_weight, snow_outputs),
+    )
+
+    return self.output_transform_fn(surface_outputs)
+
+
+@gin.register
+class NodalLandSeaIceEmbedding(hk.Module):
+  """Embedding to represent a nodal land/sea/sea-ice surface."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      output_shapes: typing.Pytree,
+      static_vars_ds_path: str,
+      land_embedding: Optional[EmbeddingModule] = None,
+      sea_embedding: Optional[EmbeddingModule] = None,
+      sea_ice_embedding: Optional[EmbeddingModule] = None,
+      output_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.output_shapes = output_shapes
+
+    # Basic surface embedding settings
+    self.feature_axis = -3
+    surface_nodal_shape = self.coords.surface_nodal_shape
+    param_init = hk.initializers.TruncatedNormal()
+    output_size = sum([x[self.feature_axis]
+                       for x in jax.tree_util.tree_leaves(output_shapes)])
+    uniform_param_shape = (output_size, 1, 1)  # uniform across lon, lat
+    # Alternative for lon,lat dependent parameters, e.g. for land model
+    # spatial_params_shape = (output_size, surface_nodal_shape[-2:])
+
+    def get_parameters_fn(
+        shape: tuple[int, int, int],
+        name: str = ''):
+      parameters = hk.get_parameter(
+          name + '_params', shape, jnp.float32, init=param_init
+      )
+      def parameters_fn(state, memory, diagnostics, randomness, forcing):
+        del state, memory, diagnostics, randomness, forcing  # unused
+        outputs = parameters * jnp.ones(surface_nodal_shape)
+        return pytree_utils.unpack_to_pytree(
+            outputs, output_shapes, self.feature_axis,
+        )
+      return parameters_fn
+
+    if land_embedding is not None:
+      self.land_embedding_fn = land_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.land_embedding_fn = get_parameters_fn(uniform_param_shape, 'land')
+
+    if sea_embedding is not None:
+      self.sea_embedding_fn = sea_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.sea_embedding_fn = get_parameters_fn(uniform_param_shape, 'sea')
+
+    if sea_ice_embedding is not None:
+      self.sea_ice_embedding_fn = sea_ice_embedding(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          output_shapes=output_shapes,
+      )
+    else:
+      self.sea_ice_embedding_fn = get_parameters_fn(
+          uniform_param_shape, 'sea_ice'
+      )
+
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    ds = xarray_utils.ds_from_path_or_aux(static_vars_ds_path, aux_features)
+    self.land_sea_mask = xarray_utils.nodal_land_sea_mask_from_ds(ds)
+
+  def __call__(
+      self,
+      state: typing.Pytree,
+      memory: Optional[typing.Pytree] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.Pytree] = None,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> typing.Pytree:
+    """Returns the embedding output on nodal locations."""
+    # get outputs from each model
+    land_outputs = self.land_embedding_fn(
+        state, memory, diagnostics, randomness, forcing)
+    sea_outputs = self.sea_embedding_fn(
+        state, memory, diagnostics, randomness, forcing)
+    sea_ice_outputs = self.sea_ice_embedding_fn(
+        state, memory, diagnostics, randomness, forcing
+    )
+
+    # prepare masks with fractional values in [0, 1]
+    land_fraction = self.land_sea_mask
+    sea_fraction = 1 - land_fraction
+    sea_ice_fraction = forcing[xarray_utils.SEA_ICE_COVER]  # pyrefly: ignore[unsupported-operation]
+
+    # weight and combine outputs
+    land_weight = land_fraction
+    sea_ice_weight = sea_ice_fraction * sea_fraction  # ice covered sea
+    sea_weight = (1 - sea_ice_fraction) * sea_fraction  # sea without ice
+
+    def tree_scale(a, x):
+      # Multiply leaves of `x` by `a`.
+      return jax.tree_util.tree_map(lambda y: a * y, x)
+
+    surface_outputs = jax.tree_util.tree_map(
+        lambda a, b, c: a + b + c,
+        tree_scale(land_weight, land_outputs),
+        tree_scale(sea_weight, sea_outputs),
+        tree_scale(sea_ice_weight, sea_ice_outputs),
+    )
+
+    return self.output_transform_fn(surface_outputs)

model/legacy/encoders.py

@@ -0,0 +1,874 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines `encoder` modules that map input trajectories to model states.
+
+All encoder modules return the encoder-specific model state that represents the
+state of the system at the latest time provided in the input trajectory.
+The inputs are expected to consist of arrays with `time` as a leading axis.
+"""
+
+# TODO(dkochkov) make all encoders take in trajectories and return ModelState.
+
+import functools
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
+
+from dinosaur import coordinate_systems
+from dinosaur import primitive_equations
+from dinosaur import pytree_utils
+from dinosaur import scales
+from dinosaur import shallow_water
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+from dinosaur import vertical_interpolation
+from dinosaur import weatherbench_utils
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import features
+from model.legacy import mappings
+from model.legacy import orographies
+from model.legacy import perturbations
+from model.legacy import stochastic
+from model.legacy import transforms
+import numpy as np
+
+
+Array = Union[np.ndarray, jnp.ndarray]
+DataState = typing.DataState
+FeaturesModule = features.FeaturesModule
+FilterModule = Callable[..., typing.PyTreeFilterFn]
+Forcing = typing.Forcing
+MappingModule = mappings.MappingModule
+PyTreeState = typing.PyTreeState
+ModelState = typing.ModelState
+TransformModule = typing.TransformModule
+OrographyModule = orographies.OrographyModule
+PerturbationModule = perturbations.PerturbationModule
+RandomnessModule = stochastic.RandomnessModule
+
+# We ♥ λ's
+# pylint: disable=g-long-lambda
+
+
+@gin.register
+class EncoderIdentityTransform(hk.Module):
+  """Transformation that returns inputs without modification."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del coords, dt, physics_specs, aux_features, input_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    return inputs
+
+
+@gin.register
+class EncoderFilterTransform(hk.Module):
+  """Transformation that returns truncated and filtered modal inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      filter_modules: Sequence[FilterModule] = tuple(),
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.filter_fns = [
+        module(coords, dt, physics_specs, aux_features)
+        for module in filter_modules
+    ]
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    for filter_fn in self.filter_fns:
+      inputs = filter_fn(inputs)
+    return inputs
+
+
+@gin.register
+class InputClipTransform(hk.Module):
+  """Filter that clips highest total wavenumber the input state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      wavenumbers_to_clip: int = 1,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.input_coords = input_coords
+    self.wavenumbers_to_clip = wavenumbers_to_clip
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    return self.input_coords.horizontal.clip_wavenumbers(
+        inputs, self.wavenumbers_to_clip
+    )
+
+
+@gin.register
+class InputNodalToModalTransform(hk.Module):
+  """Transformation that converts nodal inputs to modal representation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.input_coords = input_coords
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    to_modal_fn = self.input_coords.horizontal.to_modal
+    downsample_fn = coordinate_systems.get_spectral_downsample_fn(
+        self.input_coords, self.coords, expect_same_vertical=False
+    )
+    return jax.tree_util.tree_map(
+        lambda x: downsample_fn(to_modal_fn(x)), inputs
+    )
+
+
+@gin.register
+class ModalInputLearnedAdaptorTransform(hk.Module):
+  """Transformation using a tower to adapt modal inputs to the model domain."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      output_transform_module: TransformModule,
+      name: Optional[str] = None,
+  ):
+    del input_coords  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.modal_to_nodal_features_fn = modal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, coords
+    )
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    """Applies transform to modal inputs, returns modal outputs."""
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    prediction_shapes = jax.tree_util.tree_map(self.get_nodal_shape_fn, inputs)
+    # if `inputs` contain `sim_time` - remove it from corrections.
+    sim_time_shape = prediction_shapes.pop('sim_time', None)
+    net = self.nodal_mapping_module(prediction_shapes)
+    nodal_input_features = self.modal_to_nodal_features_fn(inputs, None)
+    nodal_corrections = self.output_transform_fn(net(nodal_input_features))
+    corrections = self.coords.horizontal.to_modal(nodal_corrections)
+    if sim_time_shape is not None:
+      corrections['sim_time'] = 0.0
+    outputs = jax.tree_util.tree_map(lambda x, y: x + y, inputs, corrections)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+class NodalInputLearnedAdaptorTransform(hk.Module):
+  """Transformation using a tower to adapt nodal inputs to the model domain."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      nodal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      output_transform_module: TransformModule,
+      name: Optional[str] = None,
+  ):
+    del input_coords  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.nodal_to_nodal_features_fn = nodal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = output_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, coords
+    )
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    """Applies transform to nodal inputs, returns nodal outputs."""
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    prediction_shapes = jax.tree_util.tree_map(self.get_nodal_shape_fn, inputs)
+    # if `inputs` contain `sim_time` - remove it from corrections.
+    sim_time_shape = prediction_shapes.pop('sim_time', None)
+    net = self.nodal_mapping_module(prediction_shapes)
+    input_features = self.nodal_to_nodal_features_fn(inputs, None)
+    corrections = self.output_transform_fn(net(input_features))
+    if sim_time_shape is not None:
+      corrections['sim_time'] = 0.0
+    outputs = jax.tree_util.tree_map(lambda x, y: x + y, inputs, corrections)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+class EncoderCombinedTransform(hk.Module):
+  """Module that applies multiple transformations sequentially."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      input_coords: coordinate_systems.CoordinateSystem,
+      transforms: Tuple[TransformModule, ...] = tuple(),  # pylint: disable=redefined-outer-name
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.transform_fns = [
+        module(coords, dt, physics_specs, aux_features, input_coords)
+        for module in transforms
+    ]
+
+  def __call__(self, inputs: PyTreeState) -> PyTreeState:
+    for transform_fn in self.transform_fns:
+      inputs = transform_fn(inputs)
+    return inputs
+
+
+@gin.register
+class ShallowWaterStateEncoder(hk.Module):
+  """Encoder that extracts shallow_water.State pair from inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=-1)
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def __call__(
+      self, inputs: DataState, forcing: Forcing
+  ) -> shallow_water.State:
+    del forcing
+    state = self.transform_fn(shallow_water.State(**self.slice_fn(inputs)))
+    return ModelState(state)  # pyrefly: ignore[bad-argument-count, bad-return]
+
+
+@gin.register
+class ShallowWaterLeapfrogEncoder(hk.Module):
+  """Encoder that extracts shallow_water.State pair from inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=slice(-2, None))
+    self.time_axis = time_axis
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def __call__(
+      self, inputs: DataState, forcing: Forcing
+  ) -> Tuple[shallow_water.State, ...]:
+    del forcing
+    last_two_frames = pytree_utils.split_axis(
+        self.slice_fn(inputs), self.time_axis
+    )
+    state = self.transform_fn(
+        tuple(shallow_water.State(**items) for items in last_two_frames)
+    )
+    return ModelState(state)  # pyrefly: ignore[bad-argument-count, bad-return]
+
+
+@gin.register
+class PrimitiveEquationStateEncoder(hk.Module):
+  """Encoder that extracts primitive_equations.State from inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=-1)
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def __call__(
+      self, inputs: DataState, forcing: Forcing
+  ) -> primitive_equations.State:
+    del forcing
+    state = self.transform_fn(
+        primitive_equations.State(**self.slice_fn(inputs)))
+    return ModelState(state)  # pyrefly: ignore[bad-argument-count, bad-return]
+
+
+@gin.register
+class PrimitiveEquationLeapfrogEncoder(hk.Module):
+  """Encoder that extracts primitive_equations.State pair from inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=slice(-2, None))
+    self.time_axis = time_axis
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def __call__(
+      self, inputs: DataState, forcing: Forcing
+  ) -> Tuple[primitive_equations.State, ...]:
+    del forcing
+    last_two_frames = pytree_utils.split_axis(
+        self.slice_fn(inputs), self.time_axis
+    )
+    state = self.transform_fn(
+        tuple(primitive_equations.State(**items) for items in last_two_frames)
+    )
+    return ModelState(state)  # pyrefly: ignore[bad-argument-count, bad-return]
+
+
+@gin.register
+class PrimitiveEquationStateWithTimeEncoder(hk.Module):
+  """Encoder that extracts primitive_equations.StateWithTime from inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=-1)
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def __call__(
+      self, inputs: DataState, forcing: Forcing
+  ) -> ModelState:
+    del forcing
+    sliced_inputs = self.slice_fn(inputs)
+    state = self.transform_fn(
+        primitive_equations.StateWithTime(**sliced_inputs))
+    return ModelState(state)  # pyrefly: ignore[bad-argument-count]
+
+
+@gin.register
+class WeatherbenchToPrimitiveEncoder(hk.Module):
+  """Encoder that extracts primitive_equations.StateWithTime from WB inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    ref_temps = aux_features[xarray_utils.REF_TEMP_KEY]
+    self.ref_temps = ref_temps[..., np.newaxis, np.newaxis]
+    self.coords = coords
+    self.input_coords = input_coords
+    self.slice_fn = functools.partial(
+        pytree_utils.slice_along_axis, axis=time_axis, idx=-1)
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features
+    )
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    self.surface_pressure_fn = functools.partial(
+        vertical_interpolation.get_surface_pressure,
+        input_coords.vertical,
+        orography=input_coords.horizontal.to_nodal(modal_orography),
+        gravity_acceleration=physics_specs.gravity_acceleration,
+    )
+    self.curl_and_div_fn = functools.partial(
+        spherical_harmonic.uv_nodal_to_vor_div_modal,
+        input_coords.horizontal,
+    )
+    self.modal_interpolate_fn = coordinate_systems.get_spectral_interpolate_fn(
+        input_coords, coords, expect_same_vertical=False
+    )
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+
+  def weatherbench_to_primitive(
+      self,
+      wb_state_nodal: weatherbench_utils.State,
+  ) -> ModelState:
+    """Converts wb_state on pressure coordinates to primitive on sigma."""
+    # Note: the returned values have mixed nodal/modal representations.
+    surface_pressure = self.surface_pressure_fn(wb_state_nodal.z)
+    interpolate_fn = vertical_interpolation.vectorize_vertical_interpolation(
+        vertical_interpolation.vertical_interpolation
+    )
+    regrid_fn = functools.partial(
+        vertical_interpolation.interp_pressure_to_sigma,
+        pressure_coords=self.input_coords.vertical,
+        sigma_coords=self.coords.vertical,
+        surface_pressure=surface_pressure,
+        interpolate_fn=interpolate_fn,
+    )
+    wb_state_on_sigma = regrid_fn(wb_state_nodal)
+    u, v = self.coords.physics_to_dycore_sharding(
+        (wb_state_on_sigma.u, wb_state_on_sigma.v)
+    )
+    vorticity, divergence = self.coords.dycore_to_physics_sharding(
+        self.curl_and_div_fn(u, v)
+    )
+    pe_state_on_sigma = primitive_equations.StateWithTime(
+        divergence=divergence,  # pyrefly: ignore[unexpected-keyword]
+        vorticity=vorticity,  # pyrefly: ignore[unexpected-keyword]
+        temperature_variation=(wb_state_on_sigma.t - self.ref_temps),  # pyrefly: ignore[unexpected-keyword]
+        log_surface_pressure=jnp.log(surface_pressure),  # pyrefly: ignore[unexpected-keyword]
+        sim_time=wb_state_on_sigma.sim_time,  # pyrefly: ignore[unexpected-keyword]
+        tracers=wb_state_on_sigma.tracers,  # pyrefly: ignore[unexpected-keyword]
+    )
+    return pe_state_on_sigma  # pyrefly: ignore[bad-return]
+
+  def __call__(
+      self,
+      inputs: DataState,
+      forcing: Forcing,
+  ) -> ModelState:
+    del forcing
+    wb_state = weatherbench_utils.State(**self.slice_fn(inputs))
+    wb_state = coordinate_systems.maybe_to_nodal(wb_state, self.input_coords)
+    pe_state = self.weatherbench_to_primitive(wb_state)
+    pe_state = coordinate_systems.maybe_to_modal(pe_state, self.input_coords)
+    pe_state = self.modal_interpolate_fn(pe_state)
+    return ModelState(state=self.transform_fn(pe_state))  # pyrefly: ignore[unexpected-keyword]
+
+
+@gin.register
+class LearnedWeatherbenchToPrimitiveEncoder(WeatherbenchToPrimitiveEncoder):
+  """Same as `WeatherbenchToPrimitiveEncoder`, but with learned corrections."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      modal_to_nodal_data_features_module: FeaturesModule,
+      modal_to_nodal_model_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      correction_transform_module: TransformModule,
+      prediction_mask: typing.Pytree,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      name: Optional[str] = None,
+  ):
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords=input_coords,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        name=name,
+    )
+    self.prediction_mask = prediction_mask
+    # data features are computed in real space on input coordinates.
+    self.data_features_fn = modal_to_nodal_data_features_module(
+        input_coords, dt, physics_specs, aux_features
+    )
+    self.model_features_fn = modal_to_nodal_model_features_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.nodal_mapping_module = nodal_mapping_module
+    self.output_transform_fn = correction_transform_module(
+        input_coords, dt, physics_specs, aux_features
+    )
+    self.get_nodal_shape_fn = lambda x: coordinate_systems.get_nodal_shapes(
+        x, coords
+    )
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+    self.randomness_fn = randomness_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.perturbation_fn = perturbation_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: DataState,
+      forcing: Forcing,
+  ) -> ModelState:
+    randomness = self.randomness_fn.unconditional_sample(
+        hk.maybe_next_rng_key()
+    )
+    wb_state = self.coords.with_physics_sharding(
+        weatherbench_utils.State(**self.slice_fn(inputs))
+    )
+    wb_state_nodal = self.coords.with_physics_sharding(
+        coordinate_systems.maybe_to_nodal(wb_state, self.input_coords)
+    )
+    wb_state_modal = self.coords.with_physics_sharding(
+        coordinate_systems.maybe_to_modal(wb_state, self.input_coords)
+    )
+    pe_state = self.coords.physics_to_dycore_sharding(
+        self.weatherbench_to_primitive(wb_state_nodal)
+    )
+    # Computing corrections to the primitive_equations state.
+    pe_state_modal = coordinate_systems.maybe_to_modal(
+        pe_state, self.input_coords
+    )
+    # we need to interpolate `pe_state_modal` to self.coords to compute
+    # features in model space. In most cases this is no-op as grids match.
+    pe_state_modal = self.modal_interpolate_fn(pe_state_modal)
+    pe_state_nodal = coordinate_systems.maybe_to_nodal(
+        pe_state_modal, self.coords
+    )
+    prediction_shapes = jax.tree_util.tree_map(
+        lambda x, y: self.get_nodal_shape_fn(x) if y else None,
+        pe_state_nodal.asdict(),
+        self.prediction_mask,
+    )
+    prediction_shapes = primitive_equations.StateWithTime(**prediction_shapes)
+    net = self.nodal_mapping_module(prediction_shapes)
+    # we need modal values to compute features for ML corrections.
+    data_features = self.data_features_fn(
+        wb_state_modal.asdict(), forcing=forcing,
+    )
+    model_features = self.model_features_fn(
+        pe_state_modal.asdict(), forcing=forcing,
+        randomness=randomness.nodal_value,
+    )
+    data_features = transforms.add_prefix(data_features, 'data_')
+    model_features = transforms.add_prefix(model_features, 'model_')
+
+    all_features = self.coords.with_physics_sharding(
+        data_features | model_features
+    )
+
+    nodal_corrections = self.coords.with_physics_sharding(
+        self.output_transform_fn(net(all_features))
+    )
+
+    perturbed_correction = self.perturbation_fn(
+        state=None,  # Unused
+        inputs=nodal_corrections,
+        randomness=randomness.nodal_value,
+    )
+
+    add_fn = lambda x, y: x + y if y is not None else x
+    corrected_pe_state = self.coords.physics_to_dycore_sharding(
+        jax.tree_util.tree_map(
+            add_fn,
+            coordinate_systems.maybe_to_modal(pe_state_nodal, self.coords),
+            coordinate_systems.maybe_to_modal(
+                perturbed_correction, self.coords
+            ),
+        )
+    )
+    return ModelState(state=self.transform_fn(corrected_pe_state))  # pyrefly: ignore[unexpected-keyword]
+
+
+@gin.register
+class DimensionalWeatherbenchToPrimitiveEncoder(WeatherbenchToPrimitiveEncoder):
+  """Same as WeatherbenchToPrimitiveEncoder, but with dimensional inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      name: Optional[str] = None,
+  ):
+    nondim_pressure_centers = physics_specs.nondimensionalize(
+        input_coords.vertical.centers * scales.units.millibar
+    )
+    nondim_input_coords = coordinate_systems.CoordinateSystem(
+        input_coords.horizontal,
+        vertical_interpolation.PressureCoordinates(nondim_pressure_centers),
+        spmd_mesh=input_coords.spmd_mesh,
+    )
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords=nondim_input_coords,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        transform_module=transform_module,
+        name=name,
+    )
+    self.nondim_transform_fn = transforms.NondimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        nondim_input_coords,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+
+  def __call__(
+      self,
+      inputs: DataState,
+      forcing: Forcing,
+  ) -> primitive_equations.StateWithTime:
+    nondim_inputs = self.nondim_transform_fn(inputs)
+    return super().__call__(nondim_inputs, forcing)  # pyrefly: ignore[bad-return]
+
+
+@gin.register
+class DimensionalLearnedWeatherbenchToPrimitiveEncoder(
+    LearnedWeatherbenchToPrimitiveEncoder
+):
+  """Same as LearnedWeatherbenchToPrimitiveEncoder, but with dimensional inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      modal_to_nodal_data_features_module: FeaturesModule,
+      modal_to_nodal_model_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      correction_transform_module: TransformModule,
+      prediction_mask: typing.Pytree,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      name: Optional[str] = None,
+  ):
+    nondim_pressure_centers = physics_specs.nondimensionalize(
+        input_coords.vertical.centers * scales.units.millibar
+    )
+    nondim_input_coords = coordinate_systems.CoordinateSystem(
+        input_coords.horizontal,
+        vertical_interpolation.PressureCoordinates(nondim_pressure_centers),
+        spmd_mesh=input_coords.spmd_mesh,
+    )
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords=nondim_input_coords,
+        modal_to_nodal_data_features_module=modal_to_nodal_data_features_module,
+        modal_to_nodal_model_features_module=(
+            modal_to_nodal_model_features_module
+        ),
+        nodal_mapping_module=nodal_mapping_module,
+        correction_transform_module=correction_transform_module,
+        prediction_mask=prediction_mask,
+        time_axis=time_axis,
+        orography_module=orography_module,
+        transform_module=transform_module,
+        randomness_module=randomness_module,
+        perturbation_module=perturbation_module,
+        name=name,
+    )
+    self.nondim_transform_fn = transforms.NondimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        nondim_input_coords,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+
+  def __call__(
+      self,
+      inputs: DataState,
+      forcing: Forcing,
+  ) -> ModelState:
+    nondim_inputs = self.nondim_transform_fn(inputs)
+    return super().__call__(nondim_inputs, forcing)
+
+
+@gin.register
+class DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder(hk.Module):
+  """Same as DimensionalLearnedWeatherbenchToPrimitiveEncoder, but with memory.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      modal_to_nodal_data_features_module: FeaturesModule,
+      modal_to_nodal_model_features_module: FeaturesModule,
+      nodal_mapping_module: MappingModule,
+      correction_transform_module: TransformModule,
+      prediction_mask: typing.Pytree,
+      time_axis: int = 0,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      transform_module: TransformModule = EncoderIdentityTransform,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      name: Optional[str] = None,
+  ):
+    nondim_pressure_centers = physics_specs.nondimensionalize(
+        input_coords.vertical.centers * scales.units.millibar)
+    nondim_input_coords = coordinate_systems.CoordinateSystem(
+        input_coords.horizontal,
+        vertical_interpolation.PressureCoordinates(nondim_pressure_centers),
+        spmd_mesh=input_coords.spmd_mesh,
+    )
+    super().__init__(name=name)
+    make_encoder_fn = functools.partial(
+        LearnedWeatherbenchToPrimitiveEncoder,
+        coords=coords, dt=dt,
+        physics_specs=physics_specs, aux_features=aux_features,
+        input_coords=nondim_input_coords,
+        modal_to_nodal_data_features_module=
+        modal_to_nodal_data_features_module,
+        modal_to_nodal_model_features_module=
+        modal_to_nodal_model_features_module,
+        nodal_mapping_module=nodal_mapping_module,
+        correction_transform_module=correction_transform_module,
+        prediction_mask=prediction_mask, time_axis=time_axis,
+        orography_module=orography_module,
+        transform_module=transform_module,
+        name=name
+    )
+
+    # Memory will be deterministic. State may be random.
+    self.memory_encoder = make_encoder_fn(
+        randomness_module=stochastic.NoRandomField,
+        perturbation_module=perturbations.NoPerturbation,
+    )
+    self.state_encoder = make_encoder_fn(
+        randomness_module=randomness_module,
+        perturbation_module=perturbation_module,
+    )
+
+    self.nondim_transform_fn = transforms.NondimensionalizeTransform(
+        coords, dt, physics_specs, aux_features, nondim_input_coords,
+        inputs_to_units_mapping=inputs_to_units_mapping)
+
+  def __call__(
+      self,
+      inputs: DataState,
+      forcing: Forcing,
+  ) -> ModelState:
+    nondim_inputs = self.nondim_transform_fn(inputs)
+    memory = self.memory_encoder(nondim_inputs, forcing=forcing)
+    model_state = self.state_encoder(nondim_inputs, forcing=forcing)
+    return ModelState(
+        state=model_state.state,  # pyrefly: ignore[unexpected-keyword]
+        memory=memory.state,  # pyrefly: ignore[unexpected-keyword]
+        randomness=model_state.randomness,  # pyrefly: ignore[unexpected-keyword]
+    )

model/legacy/equations.py

@@ -0,0 +1,402 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ML modules for equation-based models."""
+
+from typing import Any, Callable, Optional, Sequence, Union
+from dinosaur import coordinate_systems
+from dinosaur import held_suarez
+from dinosaur import primitive_equations
+from dinosaur import pytree_utils
+from dinosaur import scales
+from dinosaur import shallow_water
+from dinosaur import sigma_coordinates
+from dinosaur import time_integration
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import features
+from model.legacy import mappings
+from model.legacy import orographies
+from model.legacy import parameterizations
+
+units = scales.units
+SCALE = scales.DEFAULT_SCALE
+QuantityOrStr = Union[str, scales.Quantity]
+EquationModule = Callable[..., time_integration.ImplicitExplicitODE]
+TransformModule = typing.TransformModule
+FeaturesModule = features.FeaturesModule
+OrographyModule = orographies.OrographyModule
+MappingModule = mappings.MappingModule
+StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
+
+REF_TEMP_KEY = xarray_utils.REF_TEMP_KEY
+REF_POTENTIAL_KEY = xarray_utils.REF_POTENTIAL_KEY
+OROGRAPHY = xarray_utils.OROGRAPHY
+
+
+@gin.register
+class ShallowWaterEquations(shallow_water.ShallowWaterEquations):
+  """Equation module for shallow water system."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: shallow_water.ShallowWaterSpecs,
+      aux_features: typing.AuxFeatures,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      name: Optional[str] = None,
+  ):
+    reference_potential = aux_features.get(REF_POTENTIAL_KEY, None)
+    if reference_potential is None:
+      raise ValueError(f'must supply {REF_POTENTIAL_KEY} in `aux_features`.')
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features)
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        orography=modal_orography,
+        reference_potential=reference_potential,
+    )
+
+
+@gin.register
+class PrimitiveEquations(primitive_equations.PrimitiveEquations):
+  """Equation module for primitive equations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: primitive_equations.PrimitiveEquationsSpecs,
+      aux_features: typing.AuxFeatures,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      vertical_advection: Callable[..., jax.Array] = (
+          sigma_coordinates.centered_vertical_advection
+      ),
+      include_vertical_advection: bool = True,
+      name: Optional[str] = None,
+  ):
+    ref_temperatures = aux_features.get(REF_TEMP_KEY, None)
+    if ref_temperatures is None:
+      raise ValueError(f'must supply {REF_TEMP_KEY} in `aux_features`.')
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features)
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        reference_temperature=ref_temperatures,
+        orography=modal_orography,
+        vertical_advection=vertical_advection,
+        include_vertical_advection=include_vertical_advection,
+    )
+
+
+@gin.register
+class PrimitiveEquationsWithTime(
+    primitive_equations.PrimitiveEquationsWithTime
+):
+  """Equation module for primitive equations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: primitive_equations.PrimitiveEquationsSpecs,
+      aux_features: typing.AuxFeatures,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      vertical_advection: Callable[..., jax.Array] = (
+          sigma_coordinates.centered_vertical_advection
+      ),
+      include_vertical_advection: bool = True,
+      name: Optional[str] = None,
+  ):
+    ref_temperatures = aux_features.get(REF_TEMP_KEY, None)
+    if ref_temperatures is None:
+      raise ValueError(f'must supply {REF_TEMP_KEY} in `aux_features`.')
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features)
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        reference_temperature=ref_temperatures,
+        orography=modal_orography,
+        vertical_advection=vertical_advection,
+        include_vertical_advection=include_vertical_advection,
+    )
+
+
+@gin.register
+class MoistPrimitiveEquations(
+    primitive_equations.MoistPrimitiveEquations
+):
+  """Equation module for moist primitive equations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: primitive_equations.PrimitiveEquationsSpecs,
+      aux_features: typing.AuxFeatures,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      vertical_advection: Callable[..., jax.Array] = (
+          sigma_coordinates.centered_vertical_advection
+      ),
+      include_vertical_advection: bool = True,
+      name: Optional[str] = None,
+  ):
+    ref_temperatures = aux_features.get(REF_TEMP_KEY, None)
+    if ref_temperatures is None:
+      raise ValueError(f'must supply {REF_TEMP_KEY} in `aux_features`.')
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features)
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        reference_temperature=ref_temperatures,
+        orography=modal_orography,
+        vertical_advection=vertical_advection,
+        include_vertical_advection=include_vertical_advection,
+    )
+
+
+@gin.register
+class MoistPrimitiveEquationsWithCloudMoisture(
+    primitive_equations.MoistPrimitiveEquationsWithCloudMoisture
+):
+  """Equation module for moist primitive equations with clouds."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: primitive_equations.PrimitiveEquationsSpecs,
+      aux_features: typing.AuxFeatures,
+      orography_module: OrographyModule = orographies.ClippedOrography,
+      vertical_advection: Callable[..., jax.Array] = (
+          sigma_coordinates.centered_vertical_advection
+      ),
+      include_vertical_advection: bool = True,
+      name: Optional[str] = None,
+  ):
+    ref_temperatures = aux_features.get(REF_TEMP_KEY, None)
+    if ref_temperatures is None:
+      raise ValueError(f'must supply {REF_TEMP_KEY} in `aux_features`.')
+    modal_orography_init_fn = orography_module(
+        coords, dt, physics_specs, aux_features)
+    modal_orography = modal_orography_init_fn()  # pytype: disable=not-callable  # jax-ndarray
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        reference_temperature=ref_temperatures,
+        orography=modal_orography,
+        vertical_advection=vertical_advection,
+        include_vertical_advection=include_vertical_advection,
+    )
+
+
+@gin.register
+class MoistPrimitiveEquationsWithCloudMoisutre(
+    MoistPrimitiveEquationsWithCloudMoisture
+):
+  """Temporary alias with mis-spelled name."""
+
+
+@gin.register
+class HeldSuarezEquations(held_suarez.HeldSuarezForcing):
+  """Equation module for Held-Suarez forcing equations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: primitive_equations.PrimitiveEquationsSpecs,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    ref_temperatures = aux_features.get(REF_TEMP_KEY, None)
+    if ref_temperatures is None:
+      raise ValueError(f'must supply {REF_TEMP_KEY} in `aux_features`.')
+    super().__init__(
+        coords=coords,
+        physics_specs=physics_specs,  # pyrefly: ignore[bad-argument-type]
+        reference_temperature=ref_temperatures)
+
+
+# TODO(dkochkov) Test if vertical diffusion works well with euler integrator.
+
+
+@gin.register
+class VerticalDiffusion(time_integration.ExplicitODE):
+  """Equation module that adds explicit diffusion along vertical direction."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      timescale: QuantityOrStr = gin.REQUIRED,
+  ):
+    self.coords = coords
+    timescale = dt / physics_specs.nondimensionalize(scales.Quantity(timescale))
+    timescales = coords.vertical.boundaries * timescale  # pyrefly: ignore[missing-attribute]
+    self.level_weighted_timescales = timescales[:, jnp.newaxis, jnp.newaxis]
+
+  def explicit_terms(self, state: typing.PyTreeState) -> typing.PyTreeState:
+    def vertical_diffusion_fn(x: typing.Array) -> typing.Array:
+      # TODO(dkochkov) Consider using sigma_coordinates.centered_difference.
+      x_grad = x[1:, ...] - x[:-1, ...]
+      # padding with zero values for vertical fluxes.
+      pad_width = ((1, 1), (0, 0), (0, 0))
+      x_grad = jnp.pad(x_grad, pad_width)
+      fluxes = self.level_weighted_timescales * x_grad
+      # TODO(dkochkov) Consider using sigma_coordinates.centered_difference.
+      return fluxes[1:, ...] - fluxes[:-1, ...]
+
+    nodal_state = self.coords.horizontal.to_nodal(state)
+    nodal_tendency = pytree_utils.tree_map_where(
+        condition_fn=lambda x: jnp.asarray(x).shape == self.coords.nodal_shape,  # pyrefly: ignore[bad-argument-type]
+        f=vertical_diffusion_fn,
+        g=jnp.zeros_like,
+        x=nodal_state)
+    modal_tendency = self.coords.horizontal.to_modal(nodal_tendency)
+    return self.coords.horizontal.clip_wavenumbers(modal_tendency)
+
+
+@gin.register
+class NoDynamics(time_integration.ImplicitExplicitODE):
+  """The constant ODE, ∂u/∂t = 0."""
+
+  def __init__(self, *args, **kwargs):
+    del args, kwargs
+
+  def explicit_terms(self, x: typing.PyTreeState) -> typing.PyTreeState:
+    return 0 * x  # pyrefly: ignore[bad-return, unsupported-operation]
+
+  def implicit_terms(self, x: typing.PyTreeState) -> typing.PyTreeState:
+    return 0 * x  # pyrefly: ignore[bad-return, unsupported-operation]
+
+  def implicit_inverse(
+      self, x: typing.PyTreeState, time_step: float
+  ) -> typing.PyTreeState:
+    return x
+
+
+@gin.register
+def composed_equations_module(
+    coords: coordinate_systems.CoordinateSystem,
+    dt: float,
+    physics_specs: Any,
+    aux_features: typing.AuxFeatures,
+    equation_modules: Sequence[EquationModule],
+) -> time_integration.ImplicitExplicitODE:
+  """Returns an equation module that represents a composition of equations."""
+  equations = tuple(eq(coords, dt, physics_specs, aux_features)
+                    for eq in equation_modules)
+  return time_integration.compose_equations(equations)
+
+
+@gin.register
+class DirectNeuralEquations(hk.Module, time_integration.ExplicitODE):
+  """Computes explicit tendencies for the input state.
+
+  This equation module predicts tendencies directly in the nodal representation
+  and returns values transformed back to the modal space. The nodal tendencies
+  are computed by the `nodal_mapping_module` from preprocessed nodal features
+  computed by `modal_to_nodal_features_module` followed by the
+  `tendency_transform_module`.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: mappings.MappingModule,
+      tendency_transform_module: TransformModule,
+      prediction_mask: Optional[typing.Pytree] = None,
+      filter_module: Optional[StepFilterModule] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.parameterization_fn = parameterizations.DirectNeuralParameterization(
+        coords=coords,
+        dt=dt,
+        physics_specs=physics_specs,
+        aux_features=aux_features,
+        modal_to_nodal_features_module=modal_to_nodal_features_module,
+        nodal_mapping_module=nodal_mapping_module,
+        tendency_transform_module=tendency_transform_module,
+        prediction_mask=prediction_mask,
+        filter_module=filter_module,
+        name=name,
+    )
+
+  def explicit_terms(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    modal_tendencies = self.parameterization_fn(inputs, forcing=None)
+    modal_tendencies = pytree_utils.none_to_zeros(modal_tendencies, inputs)
+    return modal_tendencies
+
+
+@gin.register
+class DivCurlNeuralEquations(hk.Module, time_integration.ExplicitODE):
+  """Computes explicit tendencies using div and curl operators for `u, v` terms.
+
+  This equation module predicts tendencies of the inputs with velocity-based
+  parameterization of the `divergence` and `vorticity` components. Specifically,
+  we replace predictions of `divergence` and `vorticity` by nodal predictions
+  of `u`, and `v`, which are then differentiated using modal representation.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: mappings.MappingModule,
+      tendency_transform_module: TransformModule,
+      prediction_mask: Optional[typing.Pytree] = None,
+      filter_module: Optional[StepFilterModule] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.parameterization_fn = parameterizations.DivCurlNeuralParameterization(
+        coords=coords,
+        dt=dt,
+        physics_specs=physics_specs,
+        aux_features=aux_features,
+        modal_to_nodal_features_module=modal_to_nodal_features_module,
+        nodal_mapping_module=nodal_mapping_module,
+        tendency_transform_module=tendency_transform_module,
+        prediction_mask=prediction_mask,
+        filter_module=filter_module,
+        name=name,
+    )
+
+  def explicit_terms(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    modal_tendencies = self.parameterization_fn(inputs, forcing=None)
+    modal_tendencies = pytree_utils.none_to_zeros(modal_tendencies, inputs)
+    return modal_tendencies

model/legacy/features.py

@@ -0,0 +1,867 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules that computes relevant state features to be used by ML components."""
+
+from typing import Any, Callable, Mapping, Optional, Protocol, Sequence
+from dinosaur import coordinate_systems
+from dinosaur import primitive_equations
+from dinosaur import pytree_utils
+from dinosaur import radiation
+from dinosaur import scales
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import transforms
+import numpy as np
+
+
+Array = typing.Array
+Pytree = typing.Pytree
+TransformModule = typing.TransformModule
+KeyWithCosLatFactor = typing.KeyWithCosLatFactor
+
+
+class FeaturesFn(Protocol):
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    ...
+
+
+FeaturesModule = Callable[..., FeaturesFn]
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class PrimitiveEquationsDiagnosticState(hk.Module):
+  """Features modules that returns processed DiagnosticState for PE."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.coords = coords
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> primitive_equations.DiagnosticState:
+    del memory, diagnostics, randomness, forcing  # unused
+    if not isinstance(inputs, primitive_equations.State):
+      inputs = primitive_equations.State(**inputs)
+    d_state = primitive_equations.compute_diagnostic_state(inputs, self.coords)
+    return self.features_transform_fn(d_state.asdict())
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class VelocityAndPrognostics(hk.Module):
+  """Features module that returns prognostics + u,v and optionally gradients."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      fields_to_include: Optional[Sequence[str]] = None,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      compute_gradients_module: TransformModule = transforms.EmptyTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.coords = coords
+    self.fields_to_include = fields_to_include
+    self.compute_gradients_fn = compute_gradients_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def _extract_features(
+      self,
+      inputs: typing.Pytree,
+      prefix: str = '',
+  ) -> typing.Pytree:
+    """Returns a nodal velocity and prognostic features."""
+    # Note: all intermediate features have an explicit cos-lat factors in key.
+    # These factors are removed in the `__call__` method before returning.
+
+    # compute `u, v` if div/curl is available and `u, v` not in prognosics.
+    if set(['vorticity', 'divergence']).issubset(inputs.keys()) and not set(
+        ['u', 'v']
+    ).intersection(inputs.keys()):
+      cos_lat_u, cos_lat_v = spherical_harmonic.get_cos_lat_vector(
+          inputs['vorticity'], inputs['divergence'], self.coords.horizontal
+      )
+      modal_features = {
+          KeyWithCosLatFactor(prefix + 'u', 1): cos_lat_u,
+          KeyWithCosLatFactor(prefix + 'v', 1): cos_lat_v,
+      }
+    else:
+      modal_features = {}
+    prognostics_keys = list(inputs.keys())
+    prognostics_keys.remove('tracers')
+    prognostics_keys.remove('sim_time')
+    for k in prognostics_keys:
+      if self.fields_to_include is None or k in self.fields_to_include:
+        modal_features[KeyWithCosLatFactor(prefix + k, 0)] = inputs[k]
+
+    for k, v in inputs['tracers'].items():
+      if self.fields_to_include is None or k in self.fields_to_include:
+        modal_features[KeyWithCosLatFactor(prefix + k, 0)] = v
+    # Computing gradient features and adjusting cos_lat factors.
+    modal_features = self.coords.with_dycore_sharding(modal_features)
+    diff_operator_features = self.compute_gradients_fn(modal_features)
+    sec_lat = 1 / self.coords.horizontal.cos_lat
+    sec2_lat = self.coords.horizontal.sec2_lat
+    sec_lat_scales = {0: 1, 1: sec_lat, 2: sec2_lat}
+    # Computing all features in nodal space.
+    features = {}
+    for k, v in (diff_operator_features | modal_features).items():
+      sec_lat_scale = sec_lat_scales[k.factor_order]
+      features[k.name] = self.coords.horizontal.to_nodal(v) * sec_lat_scale
+    features = self.coords.with_dycore_sharding(features)
+    return features
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del memory, diagnostics, randomness, forcing  # unused.
+    nodal_features = self._extract_features(inputs)
+    return self.features_transform_fn(nodal_features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class MemoryVelocityAndValues(VelocityAndPrognostics):
+  """Similar to `VelocityAndPrognostics`, but operates on memory."""
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del inputs, diagnostics, randomness, forcing  # unused.
+    nodal_features = self._extract_features(memory, 'memory_')
+    return self.features_transform_fn(nodal_features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class NodalInputVelocityAndPrognostics(VelocityAndPrognostics):
+  """Features modules that returns velocities, temperature, and optionally gradients."""
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    to_modal_fn = self.coords.horizontal.to_modal
+    inputs = to_modal_fn(inputs)
+    memory = to_modal_fn(memory)
+    return super().__call__(inputs, memory, randomness, forcing)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class RadiationFeatures(hk.Module):
+  """Feature module that computes incident radiation flux."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    ref_datetime_str = str(aux_features[xarray_utils.REFERENCE_DATETIME_KEY])
+    self.solar_radiation = radiation.SolarRadiation.normalized(
+        coords=coords,
+        physics_specs=physics_specs,
+        reference_datetime=np.datetime64(ref_datetime_str),
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del memory, diagnostics, randomness, forcing  # unused.
+    features = {}
+    features['radiation'] = self.solar_radiation.radiation_flux(
+        inputs['sim_time']
+    )
+    # TODO(janniyuval) add a flag that allow to get radiation of next time step
+    # insert a feature axis.
+    features = jax.tree_util.tree_map(lambda x: jnp.expand_dims(x, 0), features)
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class OrbitalTimeFeatures(hk.Module):
+  """Feature module that computes orbital time features."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    ref_datetime_str = str(aux_features[xarray_utils.REFERENCE_DATETIME_KEY])
+    self.solar_radiation = radiation.SolarRadiation.normalized(
+        coords=coords,
+        physics_specs=physics_specs,
+        reference_datetime=np.datetime64(ref_datetime_str),
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del memory, diagnostics, randomness, forcing  # unused.
+    features = {}
+    # Cosine and sine of Earth's orbital phase around the Sun
+    orbital_time = self.solar_radiation.time_to_orbital_time(inputs['sim_time'])
+    # Convert from orbital_phase=0 on January 1st UTC to orbital_phase=0 at the
+    # approximate perihelion (when earth is closest to the sun).
+    orbital_phase = orbital_time.orbital_phase - radiation.PERIHELION
+    # All longitude, latitude locations share the same orbital phase
+    ones = jnp.ones(self.solar_radiation.coords.surface_nodal_shape)
+    features['cos_orbital_phase'] = jnp.cos(orbital_phase) * ones
+    features['sin_orbital_phase'] = jnp.sin(orbital_phase) * ones
+    # Cosine and sine of local hour angle (angle from solar noon)
+    solar_hour_angle = self.solar_radiation.solar_hour_angle(inputs['sim_time'])
+    solar_hour_angle = jnp.expand_dims(solar_hour_angle, 0)
+    features['cos_solar_hour'] = jnp.cos(solar_hour_angle)
+    features['sin_solar_hour'] = jnp.sin(solar_hour_angle)
+    # TODO(janniyuval) add a flag that allow to get radiation of next time step
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class ForcingFeatures(hk.Module):
+  """Feature module that provides forcing values as features."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      forcing_to_include: Sequence[str] = tuple(),
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.forcing_to_include = forcing_to_include
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> Pytree:
+    del inputs, memory, diagnostics, randomness
+    features = {}
+    for key in self.forcing_to_include:
+      value = forcing[key]  # pyrefly: ignore[unsupported-operation]
+      # Expect singleton "level" dimension for surface forcings
+      if value.ndim > 3:
+        raise ValueError(
+            f'Expected forcing "{key}" to have ndim <= 3, got {value.ndim}'
+        )
+      if value.ndim == 2:
+        value = jnp.expand_dims(value, axis=0)
+      if value.shape[0] != 1:
+        raise ValueError(
+            f'Expected forcing "{key}" to have leading dimension 1'
+            f'for level, got {value.shape}'
+        )
+      features[key] = value
+
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class LatitudeFeatures(hk.Module):
+  """Feature module that creates cos and sin of latitude as features."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.coords = coords
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del inputs, memory, diagnostics, randomness, forcing  # unused.
+    _, sin_lat = self.coords.horizontal.nodal_mesh
+    sin_features = sin_lat[np.newaxis, ...]
+    cos_features = jnp.cos(jnp.arcsin(sin_features))
+    features = {
+        'cos_latitude': cos_features,
+        'sin_latitude': sin_features,
+    }
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class RandomnessFeatures(hk.Module):
+  """Feature module that returns fields from `randomness` as features."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del inputs, memory, diagnostics, forcing  # unused.
+    if randomness is None:
+      random_features = {}
+    elif isinstance(randomness, dict):
+      random_features, _ = pytree_utils.flatten_dict(randomness)
+    elif isinstance(randomness, jax.Array):
+      random_features = {'randomness': randomness}
+    else:
+      raise ValueError(f'randomness has unsupported {type(randomness)=}.')
+    # random fields are 2D by construction, adding a feature/level dimension.
+    if randomness is not None:
+      ndims = set(x.ndim for x in jax.tree_util.tree_leaves(random_features))
+      if not ndims.issubset({2, 3}):
+        raise ValueError(
+            f'Random fields expected to be 2D and/or 3D. Found {ndims=}'
+        )
+
+      def make_3d(x):
+        if x.ndim == 3:
+          return x
+        if x.ndim == 2:
+          return x[np.newaxis, ...]
+
+      random_features = jax.tree_util.tree_map(make_3d, random_features)
+    return self.features_transform_fn(random_features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class OrographyFeatures(hk.Module):
+  """Feature module that computes orographic features."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    if xarray_utils.OROGRAPHY not in aux_features:
+      raise ValueError('OrographyFeatures requires orography in aux_features.')
+    self.nodal_orography = aux_features[xarray_utils.OROGRAPHY]
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del inputs, memory, diagnostics, randomness, forcing  # unused.
+    features = {
+        xarray_utils.OROGRAPHY: jnp.expand_dims(self.nodal_orography, 0),
+    }
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class OneHotAuxFeatures(hk.Module):
+  """Feature module that produces one-hot encodings from binary covariates."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      covariate_keys: Sequence[str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      convert_float_to_int: bool = False,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs  # unused.
+    super().__init__(name=name)
+    covariates = {}
+    num_classes = {}
+    for key in covariate_keys:
+      if key not in aux_features:
+        raise ValueError(f'Covariate {key} not found in aux_features.')
+      if not np.issubdtype(aux_features[key].dtype, np.integer):
+        if convert_float_to_int:
+          aux_features[key] = np.round(aux_features[key]).astype(int)
+        else:
+          raise ValueError(
+              f'Covariate {key} is expected to be integer dtype, '
+              f'but is: {aux_features[key].dtype}'
+          )
+      covariates[key] = aux_features[key]
+      num_classes[key] = np.unique(aux_features[key]).size
+    self.covariates = covariates
+    self.num_classes = num_classes
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> dict[str, jnp.ndarray]:
+    del inputs, memory, diagnostics, randomness, forcing  # unused.
+    features = {
+        k: jax.nn.one_hot(v, self.num_classes[k], axis=0)
+        for k, v in self.covariates.items()
+    }
+    return features
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class LearnedPositionalFeatures(hk.Module):
+  """Feature module with learned params at surface nodal locations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      latent_size: int,
+      scale: float = 1.0,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.scale = scale
+    self.padding = coords.horizontal.nodal_padding
+    unpadded_nodal_shape = tuple(
+        x - y for x, y in zip(coords.horizontal.nodal_shape, self.padding)
+    )
+    self.positional_features = hk.get_parameter(
+        'learned_positional_features',
+        (latent_size,) + unpadded_nodal_shape,
+        jnp.float32,
+        init=hk.initializers.Constant(0.0),
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> dict[str, jnp.ndarray]:
+    """Returns scaled parameter values at surface nodal locations."""
+    del inputs, memory, diagnostics, randomness, forcing  # unused.
+    pad_x, pad_y = self.padding
+    positional_features = self.scale * jnp.pad(
+        self.positional_features, [(0, 0), (0, pad_x), (0, pad_y)]
+    )
+    return {'learned_positional_features': positional_features}
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class EmbeddingSurfaceFeatures(hk.Module):
+  """Feature module that specifies embedding surface outputs as features.
+
+  Returns {feature_name: nn_output}
+    where nn_output.shape = (output_size, lon, lat).
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      feature_name: str,
+      output_size: int,
+      embedding_module: typing.EmbeddingModule,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    # output shapes are arrays to be pytree leaves for tree_map
+    output_shapes = {
+        feature_name: np.asarray((output_size,) + coords.horizontal.nodal_shape)
+    }
+    self.embedding_fn = embedding_module(
+        coords, dt, physics_specs, aux_features, output_shapes=output_shapes
+    )
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    features = self.embedding_fn(
+        inputs, memory, diagnostics, randomness, forcing  # pyrefly: ignore[bad-argument-type]
+    )
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class EmbeddingVolumeFeatures(hk.Module):
+  """Feature module that specifies embedding volume outputs as features.
+
+  Returns {feature_name_0: nn_output_0,
+           feature_name_1: nn_output_1,
+           ...
+          }
+    where the NN output array has shape (output_size, level, lon, lat), which is
+    unpacked over output_size such that nn_output_{i}.shape = (level, lon, lat)
+    for each i in range(output_size).
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      feature_name: str,
+      output_size: int,
+      embedding_module: typing.EmbeddingModule,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    # output shapes are arrays to be pytree leaves for tree_map
+    output_shapes = {
+        f'{feature_name}_{i}': np.asarray(coords.nodal_shape)
+        for i in range(output_size)
+    }
+    self.embedding_fn = embedding_module(
+        coords, dt, physics_specs, aux_features, output_shapes=output_shapes
+    )
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    features = self.embedding_fn(
+        inputs, memory, diagnostics, randomness, forcing  # pyrefly: ignore[bad-argument-type]
+    )
+    return self.features_transform_fn(features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class FloatDataFeatures(hk.Module):
+  """Feature module that supplies floating point covariates from data."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      covariate_data_path: str = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      covariate_keys: Sequence[str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      renaming_dict: Optional[Mapping[str, str]] = None,
+      compute_gradients_module: TransformModule = transforms.EmptyTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.covariates = {}
+    self.compute_gradients_fn = compute_gradients_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.coords = coords
+    ds = xarray_utils.ds_from_path_or_aux(covariate_data_path, aux_features)
+    if renaming_dict is not None:
+      ds = ds.rename(renaming_dict)
+    lon, lat = (ds[xarray_utils.XR_LON_NAME], ds[xarray_utils.XR_LAT_NAME])
+    xarray_utils.verify_grid_consistency(lon, lat, coords.horizontal)
+    lon_lat_order = (xarray_utils.XR_LON_NAME, xarray_utils.XR_LAT_NAME)
+    for key in covariate_keys:
+      data = ds[key].transpose(*lon_lat_order)
+      data_units = scales.parse_units(data.attrs['units'])
+      data = physics_specs.nondimensionalize(data.values * data_units)
+      if data.ndim != 3:
+        data = data[np.newaxis, ...]
+      self.covariates[key] = data
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> dict[str, jnp.ndarray]:
+    del inputs, memory, diagnostics, forcing, randomness  # unused.
+    features = {k: v for k, v in self.covariates.items()}
+    modal_features = self.coords.horizontal.to_modal(features)
+    modal_features = {  # jit should eliminate to_modal if it is not used.
+        KeyWithCosLatFactor(k, 0): v for k, v in modal_features.items()
+    }
+    modal_gradient_features = self.compute_gradients_fn(modal_features)
+    sec_lat = 1 / self.coords.horizontal.cos_lat
+    sec2_lat = self.coords.horizontal.sec2_lat
+    sec_lat_scales = {0: 1, 1: sec_lat, 2: sec2_lat}
+    for k, v in modal_gradient_features.items():
+      sec_lat_scale = sec_lat_scales[k.factor_order]
+      features[k.name] = self.coords.horizontal.to_nodal(v) * sec_lat_scale
+    return features
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class CombinedFeatures(hk.Module):
+  """Feature module that combines multiple feature modules together."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      feature_modules: Sequence[FeaturesModule] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      feature_module_names_to_exclude: Sequence[str] = tuple(),
+      features_to_exclude: Sequence[str] = tuple(),
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.feature_fns = [
+        module(coords, dt, physics_specs, aux_features)
+        for module in feature_modules
+    ]
+    self.feature_module_names_to_exclude = feature_module_names_to_exclude
+    self.features_to_exclude = features_to_exclude
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> dict[str, jnp.ndarray]:
+    all_features = {}
+    for feature_fn in self.feature_fns:
+      if type(feature_fn).__name__ not in self.feature_module_names_to_exclude:
+        features = feature_fn(inputs, memory, diagnostics, randomness, forcing)
+        for k, v in features.items():
+          if k in all_features:
+            raise ValueError(f'Encountered duplicate feature {k}')
+          all_features[k] = v
+    all_features = self.features_transform_fn(all_features)
+    for k in self.features_to_exclude:
+      all_features.pop(k, None)
+    return all_features
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class NullFeatures(hk.Module):
+  """Placeholder features module that returns an empty dict."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused
+    super().__init__(name=name)
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> dict[str, jnp.ndarray]:
+    del inputs, memory, diagnostics, randomness, forcing  # unused
+    return {}
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class PressureFeatures(hk.Module):
+  """Feature module that computes pressure."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      features_transform_module: TransformModule = transforms.IdentityTransform,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.features_transform_fn = features_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def _nodal_pressure(
+      self,
+      inputs: typing.Pytree,
+      prefix: str = '',
+  ) -> Mapping[str, Array]:
+    """Computes nodal pressure from model inputs."""
+    # Compute nodal, dimensionalized quantities
+    to_nodal_fn = self.coords.horizontal.to_nodal
+    sigma = self.coords.vertical.centers
+    surface_pressure = jnp.exp(to_nodal_fn(inputs['log_surface_pressure']))
+    pressure = surface_pressure * sigma[:, jnp.newaxis, jnp.newaxis]
+    nodal_features = {prefix + 'pressure': pressure}
+    return nodal_features
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del memory, diagnostics, randomness, forcing  # unused.
+    nodal_features = self._nodal_pressure(inputs)
+    return self.features_transform_fn(nodal_features)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class MemoryPressureFeatures(PressureFeatures):
+  """Feature module that computes pressure from memory values."""
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      memory: Optional[typing.PyTreeState] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.PyTreeState] = None,
+      forcing: Optional[typing.Pytree] = None,
+  ) -> typing.Pytree:
+    del inputs, diagnostics, randomness, forcing  # unused.
+    nodal_features = self._nodal_pressure(memory, 'memory_')
+    return self.features_transform_fn(nodal_features)

model/legacy/filters.py

@@ -0,0 +1,457 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines `filtering` that aim to improve stability of integration."""
+
+from typing import Any, Callable, Dict, Optional, Sequence, Union
+from dinosaur import coordinate_systems
+from dinosaur import filtering
+from dinosaur import pytree_utils
+from dinosaur import scales
+from dinosaur import time_integration
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import numpy as np
+
+
+QuantityOrStr = Union[str, scales.Quantity]
+
+StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
+TransformModule = typing.TransformModule
+
+
+#  =============================================================================
+#  Step filters that attenuate modal components between time steps.
+#  =============================================================================
+
+
+@gin.register
+class NoFilter(hk.Module):
+  """Filter module that performs no filtering."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    del u  # unused.
+    return u_next
+
+
+@gin.register
+class ClipFilter(hk.Module):
+  """Filter that clips highest total wavenumber in the next state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      wavenumbers_to_clip: int = 1,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.wavenumbers_to_clip = wavenumbers_to_clip
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    del u  # unused.
+    return self.coords.horizontal.clip_wavenumbers(
+        u_next, self.wavenumbers_to_clip
+    )
+
+
+@gin.register
+class ExponentialLeapfrogFilter(hk.Module):
+  """Filter that removes high frequency components from a spectral state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      tau: QuantityOrStr = '0.010938',
+      order: int = 18,
+      cutoff: float = 0,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del aux_features  # unused.
+    super().__init__(name=name)
+    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
+    self.filter_fn = time_integration.exponential_leapfrog_step_filter(
+        coords.horizontal, dt, tau, order, cutoff)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    return self.filter_fn(u, u_next)
+
+
+@gin.register
+class ExponentialFilter(hk.Module):
+  """Filter that removes high frequency components from a spectral state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      tau: QuantityOrStr = '0.010938',
+      order: int = 18,
+      cutoff: float = 0,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_step_filter` for details."""
+    del aux_features  # unused.
+    super().__init__(name=name)
+    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
+    self.filter_fn = time_integration.exponential_step_filter(
+        coords.horizontal, dt, tau, order, cutoff)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    return self.filter_fn(u, u_next)
+
+
+@gin.register
+class HorizontalDiffusionFilter(hk.Module):
+  """Filter that applies implicit diffusion operator to a spectral state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      tau: QuantityOrStr = '',
+      order: int = 1,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.horizontal_diffusion_filter` for details."""
+    del aux_features  # unused.
+    super().__init__(name=name)
+    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
+    self.filter_fn = time_integration.horizontal_diffusion_step_filter(
+        coords.horizontal, dt, tau, order)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    del u  # unused
+    return self.filter_fn(u_next)  # pytype: disable=wrong-arg-count  # always-use-return-annotations
+
+
+@gin.register
+class RobertAsselinLeapfrogFilter(hk.Module):
+  """Time smoothing filter."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      strength: float = 0.05,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.robert_asselin_leapfrog_filter` for details."""
+    del dt, coords, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.filter_fn = time_integration.robert_asselin_leapfrog_filter(strength)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    return self.filter_fn(u, u_next)
+
+
+@gin.register
+class LearnedExponentialFilter(hk.Module):
+  """Low pass filter with learned parameters."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    self.coords = coords
+    self.a_init = hk.initializers.Constant(16)
+    self.p_init = hk.initializers.Constant(18)
+    self.c_init = hk.initializers.Constant(0)
+    super().__init__(name=name)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    del u  # unused.
+    a_logit = hk.get_parameter('attenuation_logit', shape=(), init=self.a_init)
+    p_logit = hk.get_parameter('order_logit', shape=(), init=self.p_init)
+    c_logit = hk.get_parameter('threshold_logit', shape=(), init=self.c_init)
+    a = jax.nn.softplus(a_logit)
+    p = jax.nn.softplus(p_logit)
+    c = jax.nn.sigmoid(c_logit)
+    filter_fn = filtering.exponential_filter(self.coords.horizontal, a, p, c)  # pytype: disable=wrong-arg-types  # jax-nn-types
+    return filter_fn(u_next)
+
+
+@gin.register
+class SequentialStepFilter(hk.Module):
+  """Filter module that combines multiple step filters applied sequentially."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      filter_modules: Sequence[StepFilterModule],
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.filter_fns = [module(coords, dt, physics_specs, aux_features)
+                       for module in filter_modules]
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    for filter_fn in self.filter_fns:
+      u_next = filter_fn(u, u_next)
+    return u_next
+
+
+@gin.register
+class LayeredStepFilter(hk.Module):
+  """Filter decorator that uses varying time-scales at different levels."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      filter_module: Union[HorizontalDiffusionFilter, ExponentialFilter],
+      tau_vals: Union[Sequence[float], np.ndarray],
+      tau_units: QuantityOrStr,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    tau = (np.asarray(tau_vals) * scales.Quantity(tau_units))
+    tau = tau[:, np.newaxis, np.newaxis]  # add spatial axes.
+    self.filter_fn = filter_module(
+        coords, dt, physics_specs, aux_features, tau=tau)  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    return self.filter_fn(u, u_next)  # pyrefly: ignore[not-callable]
+
+
+@gin.register
+class MaskedFilter(hk.Module):
+  """Filter that is only applied to a part of the state."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      filter_module: StepFilterModule,
+      mask: typing.Pytree,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
+    self.mask = mask
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    mask = type(u_next)(**self.mask)  # convert to same structure.
+    return jax.tree_util.tree_map(
+        lambda x, y, b: self.filter_fn(x, y) if b else y,
+        u, u_next, mask)
+
+
+@gin.register
+class FilterFromTransform(hk.Module):
+  """Filter module that wraps a transform module."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      transform_module: TransformModule,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.transform_fn = transform_module(
+        coords, dt, physics_specs, aux_features)
+
+  def __call__(
+      self,
+      u: typing.PyTreeState,
+      u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    return self.transform_fn(u_next)
+
+
+@gin.register
+class FixGlobalMeanFilter(hk.Module):
+  """Filter that removes the change in the global mean of certain keys."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      keys: tuple[str, ...] = ('log_surface_pressure',),
+      name: Optional[str] = None,
+  ):
+    del aux_features  # unused.
+    super().__init__(name=name)
+    self.keys = keys
+
+  def __call__(
+      self, u: typing.PyTreeState, u_next: typing.PyTreeState
+  ) -> typing.PyTreeState:
+    u_dict, _ = pytree_utils.as_dict(u)
+    u_dict, _ = pytree_utils.flatten_dict(u_dict)
+    u_next_dict, from_dict_fn = pytree_utils.as_dict(u_next)
+    u_next_dict, _ = pytree_utils.flatten_dict(u_next_dict)
+    for key in self.keys:
+      global_mean = u_dict[key][..., 0]
+      u_next_dict[key] = u_next_dict[key].at[..., 0].set(global_mean)
+    u_next_dict = pytree_utils.unflatten_dict(u_next_dict)
+    return from_dict_fn(u_next_dict)
+
+
+#  =============================================================================
+#  Filters that act on modal variables without time-step context.
+#  =============================================================================
+
+
+@gin.register
+class DataNoFilter(hk.Module):
+  """Filter module that performs no filtering."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return inputs
+
+
+@gin.register
+class DataExponentialFilter(hk.Module):
+  """Filter that removes high frequency components from a modal data."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      attenuation: float = 16,
+      order: int = 18,
+      cutoff: float = 0,
+      name: Optional[str] = None,
+  ):
+    """See `filtering.exponential_filter` for details."""
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.filter_fn = filtering.exponential_filter(
+        coords.horizontal, attenuation, order, cutoff)
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return self.filter_fn(inputs)
+
+
+@gin.register
+class PerVariableDataFilter(hk.Module):
+  """Filter module that applies different filters for each variable."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      per_variable_filters: Dict[str, Any],
+      name: Optional[str] = None,
+  ):
+    """See `filtering.exponential_filter` for details."""
+    super().__init__(name=name)
+    self.filter_fns = jax.tree_util.tree_map(
+        lambda m: m(coords, dt, physics_specs, aux_features),
+        per_variable_filters)
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    inputs_dict, from_dict_fn = pytree_utils.as_dict(inputs)
+    return from_dict_fn(jax.tree_util.tree_map(
+        lambda x, fn: fn(x), inputs_dict, self.filter_fns))

model/legacy/forcings.py

@@ -0,0 +1,292 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines `forcing` modules that produce time-dependent focing values."""
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import Any, Optional, Union
+
+from dinosaur import coordinate_systems
+from dinosaur import scales
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import transforms
+import numpy as np
+
+tree_leaves = jax.tree_util.tree_leaves
+tree_map = jax.tree_util.tree_map
+units = scales.units
+
+Pytree = typing.Pytree
+ForcingData = typing.ForcingData
+ForcingFn = typing.ForcingFn
+Forcing = typing.Forcing
+TransformModule = typing.TransformModule
+Quantity = units.Quantity
+QuantityOrStr = Union[str, scales.Quantity]
+
+
+# _FORCING_ERRORS global will store errors obtained during a io_callback.
+# The user can periodically call _check_errors to see if errors have accumulated
+# TODO(langmore) Use a more universal mechanism (not just in forcings.py) to
+# handle errors, if we like this, then make public.
+_FORCING_ERRORS = []
+
+# pylint: disable=logging-fstring-interpolation
+
+
+class ForcingDataError(Exception):
+  """To raise when an error is encountered with forcing data."""
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class NoForcing(hk.Module):
+  """Module that returns an empty Forcing object."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      time_axis: int = 0,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del coords, dt, physics_specs, aux_features, time_axis
+
+  def __call__(
+      self,
+      forcing_data: ForcingData,
+      sim_time: float,
+  ) -> Forcing:
+    """Returns forcings at the specified sim_time."""
+    del forcing_data, sim_time
+    return {}
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class DynamicDataForcing(hk.Module):
+  """Modules that returns forcing values by querying time-varying data.
+
+  Input to __call__ `sim_time` must match a value in forcing_data['sim_time']
+  within dt_tolerance, or else it returns nan for all pytree values.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      inputs_to_units_mapping: dict[str, str],
+      forcing_transform: TransformModule = transforms.IdentityTransform,
+      time_axis: int = 0,
+      data_time_step: float | QuantityOrStr | None = None,
+      dt_tolerance: Union[float, QuantityOrStr] = '1 hour',
+      # TODO(langmore) Remove checking once bug arising from http://cl/624039690
+      # is fixed.
+      check_sim_time_errors: bool = False,
+      name: Optional[str] = None,
+  ):
+    logging.info(f'[NGCM] Initializing DynamicDataForcing with {dt_tolerance=}')
+    # TODO(shoyer): remove data_time_step entirely, once we're sure that no
+    # saved checkpoints that we care about will break.
+    del data_time_step  # no longer used
+    super().__init__(name=name)
+    self.time_axis = time_axis
+    self.nondim_transform_fn = transforms.NondimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords=None,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+    self.forcing_transform_fn = forcing_transform(
+        coords, dt, physics_specs, aux_features
+    )
+    if isinstance(dt_tolerance, (str, scales.Quantity)):  # pyrefly: ignore[invalid-argument]
+      dt_tolerance = physics_specs.nondimensionalize(
+          scales.Quantity(dt_tolerance)
+      )
+    self.dt_tolerance = dt_tolerance
+    self._check_sim_time_errors = check_sim_time_errors
+
+  def __call__(
+      self,
+      forcing_data: ForcingData,
+      sim_time: float,
+  ) -> Forcing:
+    """Returns forcings at the specified sim_time."""
+    forcing_data = self.nondim_transform_fn(forcing_data)
+
+    times = forcing_data['sim_time']
+    approx_index = jnp.interp(sim_time, times, jnp.arange(times.size))
+    index = jnp.round(approx_index).astype('int32')
+
+    # Slice leaf values by index
+    field_index_fn = functools.partial(
+        jax.lax.dynamic_index_in_dim,
+        index=index,
+        axis=self.time_axis,
+        keepdims=False,
+    )
+    _assert_no_scalars(forcing_data)
+    forcing = tree_map(field_index_fn, forcing_data)
+
+    # Replace leaf values with nan if forcing['sim_time'] does not match
+    # the requested sim_time value within dt_tolerance.
+    abs_error = jnp.abs(forcing['sim_time'] - sim_time)
+    is_valid = abs_error < self.dt_tolerance
+    forcing = jax.tree_util.tree_map(
+        lambda x: jnp.where(is_valid, x, jnp.nan), forcing
+    )
+
+    # Also add errors (if any) to _FORCING_ERRORS so _check_errors can be called
+    # to raise.
+    if self._check_sim_time_errors:
+      jax.experimental.io_callback(
+          _check_sim_time_close_to_forcing_sim_time,
+          None,  # Returns None
+          sim_time=sim_time,
+          forcing_sim_time=forcing['sim_time'],
+          tolerance=self.dt_tolerance,
+      )
+    return self.forcing_transform_fn(forcing)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class PersistenceDataForcing(hk.Module):
+  """Modules that returns forcing using first time index of forcing_data."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      inputs_to_units_mapping: dict[str, str],
+      forcing_transform: TransformModule = transforms.IdentityTransform,
+      time_axis: int = 0,
+      name: Optional[str] = None,
+  ):
+    logging.info('[NGCM] Initializing PersistenceDataForcing')
+    super().__init__(name=name)
+    self.time_axis = time_axis
+    self.nondim_transform_fn = transforms.NondimensionalizeTransform(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords=None,
+        inputs_to_units_mapping=inputs_to_units_mapping,
+    )
+    self.forcing_transform_fn = forcing_transform(
+        coords, dt, physics_specs, aux_features
+    )
+
+  def __call__(
+      self,
+      forcing_data: ForcingData,
+      sim_time: float,
+  ) -> Forcing:
+    """Returns forcings from the first time index of sim_time."""
+    del sim_time  # unused
+    forcing_data = self.nondim_transform_fn(forcing_data)
+    idx = 0
+
+    # Slice leaf values by index
+    field_index_fn = functools.partial(
+        jax.lax.dynamic_index_in_dim,
+        index=idx,
+        axis=self.time_axis,
+        keepdims=False,
+    )
+    _assert_no_scalars(forcing_data)
+    forcing = tree_map(field_index_fn, forcing_data)
+    return self.forcing_transform_fn(forcing)
+
+
+@gin.register(denylist=['coords', 'dt', 'physics_specs', 'aux_features'])
+class IncrementSSTForcingTransform(hk.Module):
+  """Transform Forcing by uniformly incrementing sea surface temperature."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      temperature_change: Quantity,
+      key: str = 'sea_surface_temperature',
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del coords, dt, aux_features  # unused
+    self.temperature_change = physics_specs.nondimensionalize(
+        units.Quantity(temperature_change)  # pyrefly: ignore[not-callable]
+    )
+    self.key = key
+
+  def __call__(self, forcing: Forcing) -> Forcing:
+    assert isinstance(forcing, dict)
+    forcing = forcing.copy()
+    forcing[self.key] = forcing[self.key] + self.temperature_change
+    return forcing
+
+
+def _assert_no_scalars(tree: Pytree):
+  dims = tree_map(lambda x: len(jnp.shape(x)), tree)
+  if not all(d > 0 for d in tree_leaves(dims)):
+    raise ValueError(f'Scalar shapes encountered: {dims=}')
+
+
+# TODO(langmore) Use a more universal mechanism (not just in forcings.py) to
+# handle errors, if we like this, then make public.
+def _check_sim_time_close_to_forcing_sim_time(
+    sim_time: np.ndarray,
+    forcing_sim_time: np.ndarray,
+    tolerance: float,
+) -> None:
+  """Checks |sim_time - forcing_sim_time| < tolerance add to _FORCING_ERRORS."""
+  abs_error = np.abs(forcing_sim_time - sim_time)
+  if abs_error < tolerance:
+    return
+  err_msg = (
+      f'{sim_time=} differed from {forcing_sim_time=} by {abs_error=} which is '
+      f'> {tolerance=}'
+  )
+  _FORCING_ERRORS.append(err_msg)
+
+
+# TODO(langmore) Use a more universal mechanism (not just in forcings.py) to
+# handle errors, if we like this, then make public.
+def _check_errors(  # pylint: disable=dangerous-default-value
+    max_to_print: int = 4,
+    err_list: list[str] = _FORCING_ERRORS,
+) -> None:
+  """Check err_list and raise ForcingDataError if nonempty."""
+  n_err = len(err_list)
+  if n_err:
+    raise ForcingDataError(
+        f'ForcingDataError found: {n_err} exceptions: '
+        f'The first {min(n_err, max_to_print)} are: '
+        f'{", ".join(err_list[:max_to_print])}'
+    )

model/legacy/gin_utils.py

@@ -0,0 +1,81 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helper functions for processing and parsing gin configurations."""
+
+import contextlib
+import logging
+import threading
+import gin
+
+
+_GIN_LOCK = threading.RLock()
+
+
+def _remove_unknown_reference(gin_config_str: str) -> str:
+  """Removes unknown references form `gin_config_str`."""
+  # this happens when we have gin MACROS reference not imported objects.
+  return '\n'.join([
+      line for line in gin_config_str.splitlines()
+      if 'gin.config._UnknownConfigurable' not in line
+  ])
+
+
+def parse_gin_config(
+    physics_config_str: str,
+    model_config_str: str,
+    override_physics_configs_from_data: bool,
+    gin_bindings: list[str],
+):
+  """Parses physics_config_str, model_config_str and gin_bindings in order.
+
+  We use skip unknown parameters in model_config_str to avoid errors associated
+  with irrelevant training parameters that refer to configurables only imported
+  for training.
+
+  Args:
+    physics_config_str: gin configuration string for physics_specifications
+      object that stores relevant physics constants.
+    model_config_str: gin configuration string of the model.
+    override_physics_configs_from_data: whether to reparse `physics_config_str`
+      after processing `model_config_str`.
+    gin_bindings: additional gin configuration strings that will be parsed last.
+  """
+  gin.parse_config(physics_config_str)
+  gin.parse_config(model_config_str, skip_unknown=True)
+  if override_physics_configs_from_data:
+    gin.parse_config(physics_config_str)
+  gin.parse_config(gin_bindings)
+  logging.info('Evaluating model with the following config:\n %s',
+               gin.config_str())
+
+
+@contextlib.contextmanager
+def specific_config(
+    gin_config: str,
+    clear_current: bool = True,
+    skip_unknown: bool = True,
+):
+  """Context manager for evaluation of functions with `gin_config`."""
+  with _GIN_LOCK:
+    # avoid splitting long lines into multiples that may contain unknown refs.
+    current_config = gin.config_str(max_line_length=len(gin.config_str()))
+    current_config = _remove_unknown_reference(current_config)
+    if clear_current:
+      gin.clear_config()
+    try:
+      gin.parse_config(gin_config, skip_unknown=skip_unknown)
+      yield
+    finally:
+      gin.clear_config()
+      gin.parse_config(current_config)

model/legacy/initializers.py

@@ -0,0 +1,124 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of custom initializers for NN parameters."""
+
+from typing import Any, Optional, Sequence
+
+import gin
+import haiku as hk
+import jax
+import numpy as np
+
+
+# Registering default initializers.
+Constant = gin.external_configurable(hk.initializers.Constant)
+VarianceScaling = gin.external_configurable(hk.initializers.VarianceScaling)
+Orthogonal = gin.external_configurable(hk.initializers.Orthogonal)
+
+
+def _compute_fans(
+    shape: Sequence[int],
+    fan_in_axes: Optional[Sequence[int]] = None,
+) -> tuple[int, int]:
+  """Computes the number of input and output units for a weight shape."""
+  # adapted from dm-haiku/_src/initializers.py
+  if len(shape) < 1:
+    fan_in = fan_out = 1
+  elif len(shape) == 1:
+    fan_in = fan_out = shape[0]
+  elif len(shape) == 2:
+    fan_in, fan_out = shape
+  else:
+    if fan_in_axes is not None:
+      # Compute fan-in using user-specified fan-in axes.
+      fan_in = np.prod([shape[i] for i in fan_in_axes])
+      fan_out = np.prod([s for i, s in enumerate(shape)
+                         if i not in fan_in_axes])
+    else:
+      # If no axes specified, assume convolution kernels (2D, 3D, or more.)
+      # kernel_shape: (..., input_depth, depth)
+      receptive_field_size = np.prod(shape[:-2])
+      fan_in = shape[-2] * receptive_field_size
+      fan_out = shape[-1] * receptive_field_size
+  return fan_in, fan_out  # pyrefly: ignore[bad-return]
+
+
+@gin.register
+class ReducingVarianceScaling(hk.initializers.Initializer):  # pyrefly: ignore[invalid-inheritance]
+  """Initializer that result in variance that reduces as width increases.
+
+  Initializes weights that result in features with expected variance of
+  `scale / n`, where `n` corresponds to the width of the layer. This initializer
+  can be used in the output layer to achieve µ parameterization [1].
+
+  References:
+    [1]: https://arxiv.org/abs/2203.03466
+  """
+
+  def __init__(
+      self,
+      scale=1.0,
+      mode='fan_in',
+      distribution='truncated_normal',
+      fan_in_axes=None,
+  ):
+    """Constructs `ReducingVarianceScaling` initializer.
+
+    Args:
+      scale: Variance scale for a width == 1 initialization.
+      mode: One of ``fan_in``, ``fan_out``, ``fan_avg``
+      distribution: Random distribution to use. One of ``truncated_normal``,
+        ``normal`` or ``uniform``.
+      fan_in_axes: Optional sequence of int specifying which axes of the shape
+        are part of the fan-in. If none provided, then the weight is assumed
+        to be like a convolution kernel, where all leading dimensions are part
+        of the fan-in, and only the trailing dimension is part of the fan-out.
+        Useful if instantiating multi-headed attention weights.
+    """
+    if scale < 0.0:
+      raise ValueError('`scale` must be a positive float.')
+    if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
+      raise ValueError('Invalid `mode` argument:', mode)
+    distribution = distribution.lower()
+    if distribution not in {'normal', 'truncated_normal', 'uniform'}:
+      raise ValueError('Invalid `distribution` argument:', distribution)
+    self.scale = scale
+    self.mode = mode
+    self.distribution = distribution
+    self.fan_in_axes = fan_in_axes
+
+  def __call__(self, shape: Sequence[int], dtype: Any) -> jax.Array:
+    scale = self.scale
+    fan_in, fan_out = _compute_fans(shape, self.fan_in_axes)
+    if self.mode == 'fan_in':
+      scale /= max(1.0, fan_in) ** 2
+    elif self.mode == 'fan_out':
+      scale /= max(1.0, fan_out) ** 2
+    else:
+      scale /= max(1.0, (fan_in + fan_out) / 2.0) ** 2
+
+    if self.distribution == 'truncated_normal':
+      stddev = np.sqrt(scale)
+      # Adjust stddev for truncation.
+      # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+      distribution_stddev = np.asarray(.87962566103423978, dtype=dtype)
+      stddev = stddev / distribution_stddev
+      return hk.initializers.TruncatedNormal(stddev=stddev)(shape, dtype)
+    elif self.distribution == 'normal':
+      stddev = np.sqrt(scale)
+      return hk.initializers.RandomNormal(stddev=stddev)(shape, dtype)
+    else:
+      limit = np.sqrt(3.0 * scale)
+      uniform_init = hk.initializers.RandomUniform(minval=-limit, maxval=limit)
+      return uniform_init(shape, dtype)

model/legacy/integrators.py

@@ -0,0 +1,36 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines configurable time integrators to be used by models."""
+
+from typing import Callable
+from dinosaur import time_integration
+from dinosaur import typing
+import gin
+
+
+TimeIntegrator = Callable[
+    [time_integration.ImplicitExplicitODE, typing.Numeric], typing.TimeStepFn]
+
+
+backward_forward_euler = gin.external_configurable(
+    time_integration.backward_forward_euler)
+crank_nicolson_rk2 = gin.external_configurable(
+    time_integration.crank_nicolson_rk2)
+crank_nicolson_rk3 = gin.external_configurable(
+    time_integration.crank_nicolson_rk3)
+crank_nicolson_rk4 = gin.external_configurable(
+    time_integration.crank_nicolson_rk4)
+imex_rk_sil3 = gin.external_configurable(time_integration.imex_rk_sil3)
+semi_implicit_leapfrog = gin.external_configurable(
+    time_integration.semi_implicit_leapfrog)

model/legacy/layers.py

@@ -0,0 +1,443 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Basic neural network layers for whirl/gcm codebase."""
+
+from typing import Callable, Optional, Sequence, Tuple
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+from model.legacy import initializers  # pylint: disable=unused-import
+
+Array = typing.Array
+GatingFactory = typing.GatingFactory
+TowerFactory = typing.TowerFactory
+MLP = gin.external_configurable(hk.nets.MLP)
+
+# nonlinearities
+relu = gin.external_configurable(jax.nn.relu)
+gelu = gin.external_configurable(jax.nn.gelu)
+silu = gin.external_configurable(jax.nn.silu)
+
+
+@gin.register(denylist=['output_size'])
+class MlpUniform(hk.nets.MLP):
+  """MLP network with same output size in each hidden layer."""
+
+  def __init__(
+      self,
+      output_size: int,
+      num_hidden_units: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      num_hidden_layers: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      w_init: Optional[hk.initializers.Initializer] = None,
+      b_init: Optional[hk.initializers.Initializer] = None,
+      with_bias: bool = True,
+      activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
+      activate_final: bool = False,
+      w_init_final: Optional[hk.initializers.Initializer] = None,
+      b_init_final: Optional[hk.initializers.Initializer] = None,
+      name: Optional[str] = None,
+  ):
+    hidden_output_sizes = [num_hidden_units] * num_hidden_layers
+    super().__init__(
+        hidden_output_sizes,
+        w_init=w_init,
+        b_init=b_init,
+        with_bias=with_bias,
+        activation=activation,
+        activate_final=True,  # last layer added explicitly.
+        name=name,
+    )
+    self.linear_final = hk.Linear(
+        output_size=output_size,
+        w_init=w_init_final,
+        b_init=b_init_final,
+        with_bias=with_bias,
+        name='linear_%d' % num_hidden_layers,
+    )
+    self.activate_linear_final = activate_final
+
+  def __call__(
+      self,
+      inputs: jax.Array,
+      dropout_rate: Optional[float] = None,
+      rng: Optional[jax.Array] = None,
+  ) -> jax.Array:
+    out = super().__call__(inputs, dropout_rate=dropout_rate, rng=rng)
+    out = self.linear_final(out)
+    if self.activate_linear_final:
+      out = self.activation(out)
+    return out
+
+
+@gin.register(denylist=['output_size'])
+class ConvLonLat(hk.Module):
+  """Two dimensional convolutional neural network."""
+
+  def __init__(
+      self,
+      output_size: int,
+      kernel_shape: Tuple[int, int] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      with_bias: bool = True,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self._padding = []
+    for kernel_size in kernel_shape:
+      pad_left = kernel_size // 2
+      self._padding.append((pad_left, kernel_size - pad_left - 1))
+    # Use padding='VALID': since padding is done in call, haiku trims
+    self._conv_module = hk.Conv2D(
+        output_channels=output_size,
+        kernel_shape=kernel_shape,
+        with_bias=with_bias,
+        padding='VALID',
+        data_format='NCHW',
+    )
+    # NCHW = batch (ignored), channels (sigma), height (lon), width (lat)
+
+  def __call__(self, inputs: Array) -> Array:
+    """Applies convolution to inputs."""
+    # Padding order is z, x, y
+    # Periodic padding in longitude (x)
+    # Zero padding in latitude (y)
+    inputs = jnp.pad(inputs, [(0, 0), self._padding[0], (0, 0)], mode='wrap')
+    # TODO(pnorgaard): consider rotated mirror padding to simulate wrapping
+    # around the N/S poles.
+    inputs = jnp.pad(
+        inputs, [(0, 0), (0, 0), self._padding[1]], mode='constant'
+    )
+    return self._conv_module(inputs)
+
+
+@gin.register
+class ConvLevel(hk.Conv1D):
+  """1D convolution in the vertical (convolution on atmospheric columns)."""
+
+  def __init__(
+      self,
+      output_channels: int,
+      kernel_shape: int,
+      dilation_rate: int = 1,
+      padding: str = 'SAME',
+      with_bias: bool = True,
+      w_init: Optional[hk.initializers.Initializer] = None,
+      b_init: Optional[hk.initializers.Initializer] = None,
+      data_format: str = 'NCW',
+      name: Optional[str] = None,
+  ):
+    super().__init__(
+        output_channels=output_channels,
+        kernel_shape=kernel_shape,
+        rate=dilation_rate,
+        padding=padding,
+        with_bias=with_bias,
+        w_init=w_init,
+        b_init=b_init,
+        data_format=data_format,
+        name=name,
+    )
+
+
+@gin.register
+class VerticalConvNet(hk.Module):
+  """1D CNN in the vertical (convolution on atmospheric columns)."""
+
+  def __init__(
+      self,
+      output_size: int,
+      channels: Sequence[int],
+      kernel_shapes: int | Sequence[int],
+      dilation_rates: int | Sequence[int],
+      padding: str = 'SAME',
+      with_bias: bool = True,
+      w_init: Optional[hk.initializers.Initializer] = None,
+      b_init: Optional[hk.initializers.Initializer] = None,
+      data_format: str = 'NCW',
+      activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
+      activate_final: bool = False,
+      w_init_final: Optional[hk.initializers.Initializer] = None,
+      b_init_final: Optional[hk.initializers.Initializer] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    n_hidden = len(channels)
+    if isinstance(kernel_shapes, int):
+      kernel_shapes = [kernel_shapes] * (n_hidden + 1)  # +1 for output layer.
+    if isinstance(dilation_rates, int):
+      dilation_rates = [dilation_rates] * (n_hidden + 1)  # +1 for output layer.
+    channels = list(channels) + [output_size]
+    if len(set([len(channels), len(kernel_shapes), len(dilation_rates)])) != 1:
+      raise ValueError(
+          f'Missing kernel|dilation specs for {n_hidden + 1} '
+          f'layers, got {kernel_shapes=}, {dilation_rates=}.'
+      )
+    w_inits = [w_init] * n_hidden + [w_init_final]
+    b_inits = [b_init] * n_hidden + [b_init_final]
+    params = zip(channels, kernel_shapes, dilation_rates, w_inits, b_inits)
+    self.layers = []
+    for c, kernel, dilation, w_init_i, b_init_i in params:
+      self.layers.append(
+          ConvLevel(
+              output_channels=c,
+              kernel_shape=kernel,
+              dilation_rate=dilation,
+              padding=padding,
+              with_bias=with_bias,
+              w_init=w_init_i,
+              b_init=b_init_i,
+              data_format=data_format,
+          )
+      )
+    self.activation = activation
+    self.activate_final = activate_final
+
+  def __call__(self, inputs: Array) -> Array:
+    out = inputs
+    num_layers = len(self.layers)
+    for i, layer in enumerate(self.layers):
+      out = layer(out)
+      if i < (num_layers - 1) or self.activate_final:
+        out = self.activation(out)
+    return out
+
+
+@gin.register
+class LevelTransformer(hk.Module):
+  """Network that uses attention mechanism across vertical levels.
+
+  This network is a simple variation of a transformer architecture. It is
+  configurable to represent either the encoder and decoder blocks. Contrary to
+  other layers, this module accepts additional optional arguments: `latents` and
+  `positional_encoding` that enable it to represent computations with more
+  complex dependency structure. By default these arguments have value `None`, in
+  which case the network uses `inputs` and performs self-attention calculation
+  throughout. If `latents` are provided, then they are used for key and value
+  calculations for all attention blocks. If `positional_encoding` is provided,
+  then it is used to produce the first set of queries in an attention block.
+  Additionally this module supports extension with gating mechanism, generally
+  resembling GTrXL transformer from https://arxiv.org/pdf/1910.06764.pdf.
+
+  Attributes:
+    output_size: desired number of channels in the output of the module.
+    latent_size: latent representation size. Must be divisible by `num_heads`.
+    n_layers: number of transformer blocks in the network.
+    num_heads: number of attention heads in each attention layer.
+    key_size: size of key/query vectors to use for computing attention scores.
+    widening_factor: widening factor in dense layer at the end of each block.
+    activation: activation function to apply between linear transforms.
+    input_projection_net: network or layer to be used to project inputs into
+      initial latent representation. If set to `None`, then input projection is
+      skipped entirely (only possible if input size == latent_size).
+    skip_final_projection: whether to skip final projection layer. If set to
+      `True`, then requested `output_size` must be equal to `latent_size`.
+    gating_module: gating mechanism to use to combine residual connection and
+      dense updates. Defaults to residual connections.
+    name: optional name for the module.
+  """
+
+  def __init__(
+      self,
+      output_size: int,
+      latent_size: int,
+      n_layers: int,
+      num_heads: int,
+      key_size: int,
+      widening_factor: int = 2,
+      activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.gelu,
+      input_projection_net: TowerFactory = hk.Linear,
+      skip_final_projection: bool = False,
+      gating_module: GatingFactory = lambda: lambda x, y: x + y,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    value_size, reminder = divmod(latent_size, num_heads)
+    if reminder != 0:
+      raise ValueError(f'{latent_size=} is not divisible by {num_heads=}.')
+
+    self.output_size = output_size
+    self.latent_size = latent_size
+    self.n_layers = n_layers
+    self.num_heads = num_heads
+    self.key_size = key_size
+    self.value_size = value_size
+    self.wide_latent_size = widening_factor * latent_size
+    self.activation = activation
+    self.w_init = hk.initializers.VarianceScaling(2 / self.n_layers)
+    self.gating_fn = gating_module()
+
+    if input_projection_net is not None:
+      self.project_input_fn = input_projection_net(latent_size)
+    else:
+
+      def skip_with_check_fn(inputs):
+        _, d = inputs.shape
+        if d != latent_size:
+          raise ValueError(
+              f'{inputs.shape=} not compatible with {latent_size=}'
+              ' Specify projection module in the transformer.'
+          )
+        return inputs
+
+      self.project_input_fn = skip_with_check_fn
+    if skip_final_projection:
+      if output_size != self.latent_size:
+        raise ValueError(
+            f'Unable to skip projection for {output_size=}, '
+            f'{self.latent_size=}.'
+        )
+      self.final_projection = lambda x: x
+    else:
+      self.final_projection = hk.Linear(output_size)
+
+  @hk.transparent
+  def layer_norm(self, x: jnp.ndarray) -> jnp.ndarray:
+    """Applies a unique LayerNorm to x with default settings."""
+    ln = hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)
+    return ln(x)
+
+  def __call__(
+      self,
+      inputs: Array,
+      latents: Optional[Array] = None,
+      positional_encoding: Optional[Array] = None,
+  ) -> Array:
+    """Applies transformer layer to inputs. See class docstring for details."""
+    inputs = jnp.transpose(inputs)  # transpose to [levels, channels].
+    h = self.project_input_fn(inputs)
+    if latents is not None:
+      latents = jnp.transpose(latents)
+    if positional_encoding is not None:
+      init_query_input = jnp.transpose(positional_encoding)
+      special_query_stage = 0  # uses `positional_encoding` for first query.
+    else:
+      special_query_stage = -1  # ensures we pass `h_norm` to query in h_attn.
+    h_dense = None  # not used in the first layer.
+    last_layer_id = self.n_layers - 1
+    for layer_id in range(self.n_layers - 1):
+      # connects residual updates from the previous layer; skipped first time.
+      h = self.gating_fn(h, h_dense) if h_dense is not None else h
+      # apply layer norm before the attention block, as in GTrXL.
+      h_norm = self.layer_norm(h)  # pyrefly: ignore[bad-argument-type]
+      attn_block = hk.MultiHeadAttention(
+          num_heads=self.num_heads,
+          key_size=self.key_size,
+          value_size=self.value_size,
+          model_size=self.latent_size,
+          w_init=self.w_init,
+      )
+      # attend to `latents` if in decoding stage, otherwise use self-attention.
+      h_attn = attn_block(
+          query=init_query_input if layer_id == special_query_stage else h_norm,  # pyrefly: ignore[unbound-name]
+          key=latents if latents is not None else h_norm,
+          value=latents if latents is not None else h_norm,
+      )
+      # connects residual updates from attention layer.
+      h = self.gating_fn(h, h_attn)
+      if layer_id != last_layer_id:
+        dense_block = hk.Sequential([
+            hk.Linear(self.wide_latent_size, w_init=self.w_init),
+            self.activation,
+            hk.Linear(self.latent_size, w_init=self.w_init),
+        ])
+        h_dense = dense_block(self.layer_norm(h))  # pyrefly: ignore[bad-argument-type]
+
+    h_dense = self.final_projection(h)
+    h_dense = jnp.transpose(h_dense)  # transpose back to [channels, levels].
+    return h_dense
+
+
+@gin.register(denylist=['output_size'])
+class LevelBiLSTM(hk.Module):
+  """Applies a bidirectional LSTM to inputs.
+
+    This network is a bi-directional LSTM. This module accepts additional
+    optional argument, window_size which determines the number of positional
+    features the LSTM will use at each step. By default this argument have
+    value `1`, in which case the network uses features from a single level at
+    each step.
+
+  Attributes:
+    output_size: desired number of channels in the output of the module.
+    hidden_size: size of the hidden state in the LSTM.
+    n_layers: number of bi-directional LSTM layers in the network.
+    final_activation: optional activation to be applied to the output.
+    window_size: number of (local) features the LSTM will use at each step.
+    name: optional name for the module.
+  """
+  def __init__(
+      self,
+      output_size: int,
+      hidden_size: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      n_layers: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      final_activation: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None,
+      window_size: int = 1,
+      name='lstm'):
+    super().__init__(name=name)
+    self.hidden_size = hidden_size
+    self.n_layers = n_layers
+    self.final_projection = hk.Linear(output_size)
+    self.final_activation = final_activation
+    self.window_size = window_size
+
+    self.fw_lstms = []
+    self.bw_lstms = []
+    for i in range(n_layers):
+      self.fw_lstms.append(hk.LSTM(hidden_size, name=f"{name}_fw_{i}"))
+      self.bw_lstms.append(hk.LSTM(hidden_size, name=f"{name}_bw_{i}"))
+
+  def sliding_window_reshape(self, data):
+    """Reshapes data to include local vertical features."""
+    levels_num = data.shape[0]
+    pad_start = (self.window_size - 1) // 2
+    pad_end = self.window_size - 1 - pad_start
+    padded_data = jnp.pad(data, [(pad_start, pad_end)] + [(0, 0)])
+    feature_indices = (
+        jnp.arange(self.window_size)[jnp.newaxis, :]
+        + jnp.arange(levels_num)[:, jnp.newaxis]
+    )
+    windowed_data = padded_data[feature_indices, ...]
+    windowed_data = jnp.reshape(
+        windowed_data,
+        [
+            windowed_data.shape[0],
+            windowed_data.shape[2] * windowed_data.shape[1],
+        ],
+    )
+    return windowed_data
+
+  def __call__(self, inputs):
+    inputs = jnp.transpose(inputs)  # transpose to [levels, channels].
+    if self.window_size > 1:
+      inputs = self.sliding_window_reshape(inputs)
+    for i in range(self.n_layers):
+      #TODO(janniyuval): initializing from previous hidden state?
+      fw_initial_state = self.fw_lstms[i].initial_state(None)
+      bw_initial_state = self.bw_lstms[i].initial_state(None)
+
+      fw_outputs, _ = hk.dynamic_unroll(
+          self.fw_lstms[i], inputs, fw_initial_state
+      )
+      bw_outputs, _ = hk.dynamic_unroll(
+          self.bw_lstms[i], inputs, bw_initial_state, reverse=True
+      )
+      outputs = jnp.concatenate([fw_outputs, bw_outputs], axis=-1)
+      inputs = outputs
+    h_dense = self.final_projection(outputs)  # pyrefly: ignore[unbound-name]
+    if self.final_activation is not None:
+      h_dense = self.final_activation(h_dense)
+    h_dense = jnp.transpose(h_dense)  # transpose back to [channels, levels].
+    return h_dense

model/legacy/mappings.py

@@ -0,0 +1,206 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules that transform data between pytrees."""
+
+from typing import Callable, Optional, Sequence
+from dinosaur import pytree_utils
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+from model.legacy import transforms
+
+
+Array = typing.Array
+Tower = Callable[[int], Callable[..., Array]]
+MappingModule = Callable[[typing.Pytree], typing.Pytree]
+
+
+@gin.register(denylist=['output_shapes'])
+class NodalMapping(hk.Module):
+  """Maps the pytree of nodal features to a pytree of specified structure.
+
+  This module packs the pytree into a single array of shape (n, lon, lat),
+  passes it to a NN tower, and unpacks the result into a pytree with the
+  structure of output_shapes, typically (m, lon, lat).
+  """
+
+  def __init__(
+      self,
+      output_shapes: typing.Pytree,
+      tower_factory: Tower = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    feature_axis = -3  # default column axis.
+    output_size = sum([x[feature_axis]
+                       for x in jax.tree_util.tree_leaves(output_shapes)])
+    # tower preserves the last two spatial dimensions.
+    self.tower = tower_factory(output_size)
+    self.output_shapes = output_shapes
+    self.feature_axis = feature_axis
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    array = pytree_utils.pack_pytree(inputs, self.feature_axis)
+    if array.ndim != 3:
+      raise ValueError(f'Expected input array with ndim=3, got {array.shape=}')
+    outputs = self.tower(array)
+    if outputs.ndim != 3:
+      raise ValueError(f'Expected outputs with ndim=3, got {outputs.shape=}')
+    return pytree_utils.unpack_to_pytree(
+        outputs, self.output_shapes, self.feature_axis)
+
+
+@gin.register(denylist=['output_shapes'])
+class NodalVolumeMapping(hk.Module):
+  """Maps the pytree of nodal volume features to a pytree of given structure.
+
+  This module stacks the input pytree into an array of shape
+  (channel, level, lon, lat), passes it to a NN tower.  The output from the NN
+  is expected to have shape (n, level, lon, lat), and gets unpacked to a pytree
+  with the structure of output_shapes, e.g.
+    output_shapes = {
+        'var_1': jnp.asarray((level, lon, lat)),
+        'var_2': jnp.asarray((level, lon, lat)),
+        ...,
+        'var_n': jnp.asarray((level, lon, lat)),
+    }
+
+  The leaves of the input pytree must have the same shape, e.g. (1, lon, lat) or
+  (level, lon, lat). To mix shapes, broadcast before passing to the mapping.
+  """
+
+  def __init__(
+      self,
+      output_shapes: typing.Pytree,
+      tower_factory: Tower = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      name: Optional[str] = None
+  ):
+    super().__init__(name=name)
+    feature_axis = 0
+    output_size = len(jax.tree_util.tree_leaves(output_shapes))
+
+    # tower preserves the last two spatial dimensions.
+    self.tower = tower_factory(output_size)
+    self.output_shapes = output_shapes
+    self.feature_axis = feature_axis
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    array = pytree_utils.stack_pytree(inputs, axis=self.feature_axis)
+    if array.ndim != 4:
+      raise ValueError(f'Expected input array with ndim=4, got {array.shape=}')
+    outputs = self.tower(array)
+    if outputs.ndim != 4:
+      raise ValueError(f'Expected outputs with ndim=4, got {outputs.shape=}')
+    return pytree_utils.unstack_to_pytree(
+        outputs, self.output_shapes, axis=self.feature_axis
+    )
+
+
+@gin.register
+class NodalVolumeTransformerMapping(hk.Module):
+  """Maps the pytree of nodal volume features to a pytree of given structure.
+
+  Similar to NodalVolumeMapping, but uses certain features as positional
+  encoding arguments to the underlying transformer networks. Inputs are
+  expected to be of shape (channel, level, lon, lat), which are split into
+  encoder inputs, decoder inputs and positional encodings, which are then passed
+  to transformer towers. The output of the NN is expected to have shape
+  (n, level*, lon, lat), and gets unpacked to a pytree with the structure of
+  output_shapes, e.g.
+    output_shapes = {
+        'out_1': jnp.asarray((level*, lon, lat)),
+        'out_2': jnp.asarray((level*, lon, lat)),
+        ...,
+        'out_n': jnp.asarray((level*, lon, lat)),
+    }
+  Note: the output number of levels `level*` is equal to those defined by the
+  `decoder_inputs_selection_module`. In case it is empty, level* == level.
+
+  The leaves of the encoder/decoder pytrees must have the same shape, e.g.
+  (1, lon, lat) or (level, lon, lat) or (level*, lon, lat). To mix shapes,
+  broadcast before passing to the mapping.
+  """
+
+  def __init__(
+      self,
+      output_shapes: typing.Pytree,
+      encoder_transformer_tower_factory: Tower = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      decoder_transformer_tower_factory: Tower = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      latent_size: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      encoder_inputs_selection_module=gin.REQUIRED,
+      decoder_inputs_selection_module=transforms.EmptyTransform,
+      encoder_pos_encoding_module=transforms.EmptyTransform,
+      decoder_pos_encoding_module=transforms.EmptyTransform,
+      name: Optional[str] = None
+  ):
+    super().__init__(name=name)
+    feature_axis = 0
+    output_size = len(jax.tree_util.tree_leaves(output_shapes))
+    self.encoder_tower = encoder_transformer_tower_factory(latent_size)
+    self.decoder_tower = decoder_transformer_tower_factory(output_size)
+    self.output_shapes = output_shapes
+    self.feature_axis = feature_axis
+    self.get_encoder_inputs_fn = encoder_inputs_selection_module()  # pyrefly: ignore[not-callable]
+    self.get_decode_inputs_fn = decoder_inputs_selection_module()
+    self.encoder_positional_encodings_fn = encoder_pos_encoding_module()
+    self.decoder_positional_encodings_fn = decoder_pos_encoding_module()
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    enc_inputs = self.get_encoder_inputs_fn(inputs)
+    dec_inputs = self.get_decode_inputs_fn(inputs)
+    enc_array = pytree_utils.stack_pytree(enc_inputs, axis=self.feature_axis)
+    dec_array = pytree_utils.stack_pytree(dec_inputs, axis=self.feature_axis)
+    enc_pos_encoding = pytree_utils.stack_pytree(
+        self.encoder_positional_encodings_fn(inputs), axis=self.feature_axis)
+    dec_pos_encoding = pytree_utils.stack_pytree(
+        self.decoder_positional_encodings_fn(inputs), axis=self.feature_axis)
+    input_ndims = set(
+        x.ndim
+        for x in [enc_array, dec_array, enc_pos_encoding, dec_pos_encoding]
+        if x is not None)
+    if input_ndims != {4}:
+      raise ValueError(f'Expected all inputs have ndim=4, got {input_ndims=}')
+    latents = self.encoder_tower(enc_array, None, enc_pos_encoding)
+    # if dec_array is None, use latents as `inputs` and provide no `latents`.
+    decoder_latents = None if dec_array is None else latents
+    # if dec_array is None, use `latents`, otherwise use dec_array as `inputs`.
+    dec_array = dec_array if dec_array is not None else latents
+    outputs = self.decoder_tower(dec_array, decoder_latents, dec_pos_encoding)
+    if outputs.ndim != 4:
+      raise ValueError(f'Expected outputs with ndim=4, got {outputs.shape=}')
+    return pytree_utils.unstack_to_pytree(
+        outputs, self.output_shapes, axis=self.feature_axis
+    )
+
+
+@gin.register(denylist=['output_shapes'])
+class ParallelMapping(hk.Module):
+  """Maps a pytree to a pytree by additively compbining multiple mappings.
+
+  Outputs of `mappings` must be compatible with each other.
+  """
+
+  def __init__(
+      self,
+      output_shapes: typing.Pytree,
+      mappings: Sequence[MappingModule] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.mapping_fns = [m(output_shapes) for m in mappings]
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    results = [mapping_fn(inputs) for mapping_fn in self.mapping_fns]
+    return jax.tree_util.tree_map(lambda *args: sum(args), *results)

model/legacy/model_builder.py

@@ -0,0 +1,744 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines AbstractModel API, standard implementations and helper functions."""
+from __future__ import annotations
+
+import collections
+import dataclasses
+import functools
+from typing import Any, Callable, Dict, Optional, Sequence, Union
+from dinosaur import coordinate_systems
+from dinosaur import layer_coordinates
+from dinosaur import scales
+from dinosaur import sigma_coordinates
+from dinosaur import spherical_harmonic
+from dinosaur import time_integration
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax.numpy as jnp
+
+from model.legacy import correctors  # pylint: disable=unused-import
+from model.legacy import decoders  # pylint: disable=unused-import
+from model.legacy import embeddings  # pylint: disable=unused-import
+from model.legacy import encoders  # pylint: disable=unused-import
+from model.legacy import equations  # pylint: disable=unused-import
+from model.legacy import features  # pylint: disable=unused-import
+from model.legacy import filters  # pylint: disable=unused-import
+from model.legacy import forcings  # pylint: disable=unused-import
+from model.legacy import gin_utils
+from model.legacy import layers  # pylint: disable=unused-import
+from model.legacy import mappings  # pylint: disable=unused-import
+from model.legacy import model_utils
+from model.legacy import physics_specifications
+from model.legacy import steps  # pylint: disable=unused-import
+from model.legacy import stochastic  # pylint: disable=unused-import
+from model.legacy import towers  # pylint: disable=unused-import
+from model.legacy import transforms  # pylint: disable=unused-import
+import numpy as np
+import xarray
+# Note: many unused imports are needed to load configurable components;
+
+
+DEFAULT_REFERENCE_TEMPERATURE = 288
+DEFAULT_REFERENCE_DATETIME_STR = '1979-01-01T00'
+
+Array = typing.Array
+AuxFeatures = typing.AuxFeatures
+DataState = typing.DataState
+PyTreeState = typing.PyTreeState
+ModelState = typing.ModelState
+ForcingData = typing.ForcingData
+Forcing = typing.Forcing
+Numeric = typing.Numeric
+QuantityOrStr = Union[str, scales.Quantity]
+# Specifying the full signatures of Callable would get somewhat onerous
+# pylint: disable=g-bare-generic
+
+# Overzealous linter is getting confused by ABC typing.
+# pylint: disable=function-missing-types
+# pylint: disable=missing-arg-types
+
+# Register data to xarray conversion methods.
+data_to_xarray = gin.external_configurable(
+    xarray_utils.data_to_xarray, 'data_to_xarray'
+)
+# TODO(dkochkov) Remove this legacy name when no best checkpoints rely on it.
+primitive_eq_to_xarray = gin.external_configurable(
+    xarray_utils.data_to_xarray, 'primitive_eq_to_xarray'
+)
+data_to_xarray_with_renaming = gin.external_configurable(
+    xarray_utils.data_to_xarray_with_renaming, 'data_to_xarray_with_renaming'
+)
+dynamic_covariate_data_to_xarray = gin.external_configurable(
+    xarray_utils.dynamic_covariate_data_to_xarray,
+    'dynamic_covariate_data_to_xarray',
+)
+
+# Register xarray to data conversion methods.
+xarray_to_shallow_water = gin.external_configurable(
+    xarray_utils.xarray_to_shallow_water_eq_data, 'xarray_to_shallow_water'
+)
+xarray_to_primitive_eq = gin.external_configurable(
+    xarray_utils.xarray_to_primitive_eq_data, 'xarray_to_primitive_eq'
+)
+xarray_to_primitive_eq_with_time = gin.external_configurable(
+    xarray_utils.xarray_to_primitive_equations_with_time_data,
+    'xarray_to_primitive_eq_with_time',
+)
+xarray_to_weatherbench_data = gin.external_configurable(
+    xarray_utils.xarray_to_weatherbench_data, 'xarray_to_weatherbench_data'
+)
+xarray_to_data_with_renaming = gin.external_configurable(
+    xarray_utils.xarray_to_data_with_renaming, 'xarray_to_data_with_renaming'
+)
+xarray_to_dynamic_covariate_data = gin.external_configurable(
+    xarray_utils.xarray_to_dynamic_covariate_data,
+    'xarray_to_dynamic_covariate_data',
+)
+xarray_to_state_and_dynamic_covariate_data = gin.external_configurable(
+    xarray_utils.xarray_to_state_and_dynamic_covariate_data,
+    'xarray_to_state_and_dynamic_covariate_data',
+)
+coordinate_system_from_dataset = gin.external_configurable(
+    xarray_utils.coordinate_system_from_dataset,
+    'coordinate_system_from_dataset',
+    allowlist=['truncation', 'spherical_harmonics_impl'],
+)
+
+# Register grids and coordinates for instantiation of coordinate systems.
+Grid = gin.external_configurable(
+    spherical_harmonic.Grid, denylist=['spmd_mesh']
+)
+GridWithWavenumbers = gin.external_configurable(
+    spherical_harmonic.Grid.with_wavenumbers, 'GridWithWavenumbers'
+)
+GridT21 = gin.external_configurable(spherical_harmonic.Grid.T21, 'GridT21')
+GridT31 = gin.external_configurable(spherical_harmonic.Grid.T31, 'GridT31')
+GridT42 = gin.external_configurable(spherical_harmonic.Grid.T42, 'GridT42')
+GridT85 = gin.external_configurable(spherical_harmonic.Grid.T85, 'GridT85')
+GridT106 = gin.external_configurable(spherical_harmonic.Grid.T106, 'GridT106')
+GridT119 = gin.external_configurable(spherical_harmonic.Grid.T119, 'GridT119')
+GridT170 = gin.external_configurable(spherical_harmonic.Grid.T170, 'GridT170')
+GridT213 = gin.external_configurable(spherical_harmonic.Grid.T213, 'GridT213')
+GridTL31 = gin.external_configurable(spherical_harmonic.Grid.TL31, 'GridTL31')
+GridTL63 = gin.external_configurable(spherical_harmonic.Grid.TL63, 'GridTL63')
+GridTL95 = gin.external_configurable(spherical_harmonic.Grid.TL95, 'GridTL95')
+GridTL127 = gin.external_configurable(
+    spherical_harmonic.Grid.TL127, 'GridTL127'
+)
+GridTL159 = gin.external_configurable(
+    spherical_harmonic.Grid.TL159, 'GridTL159'
+)
+GridTL179 = gin.external_configurable(
+    spherical_harmonic.Grid.TL179, 'GridTL179'
+)
+GridTL255 = gin.external_configurable(
+    spherical_harmonic.Grid.TL255, 'GridTL255'
+)
+RealSphericalHarmonics = gin.external_configurable(
+    spherical_harmonic.RealSphericalHarmonics,
+)
+RealSphericalHarmonicsWithZeroImag = gin.external_configurable(
+    spherical_harmonic.RealSphericalHarmonicsWithZeroImag,
+    denylist=['spmd_mesh'],
+)
+LayerCoordinates = gin.external_configurable(layer_coordinates.LayerCoordinates)
+SigmaCoordinates = gin.external_configurable(sigma_coordinates.SigmaCoordinates)
+SigmaCoordinatesEquidistant = gin.external_configurable(
+    sigma_coordinates.SigmaCoordinates.equidistant,
+    'SigmaCoordinatesEquidistant',
+)
+CoordinateSystem = gin.external_configurable(
+    coordinate_systems.CoordinateSystem, denylist=['spmd_mesh']
+)
+
+# Register vertical interpolation methods
+centered_vertical_advection = gin.external_configurable(
+    sigma_coordinates.centered_vertical_advection
+)
+upwind_vertical_advection = gin.external_configurable(
+    sigma_coordinates.upwind_vertical_advection
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class ModelSpecs(collections.abc.Mapping):
+  """Specification of model configuration.
+
+  Attributes:
+    coords: horizontal and vertical grid data.
+    dt: nondimensionalized model time step.
+    physics_specs: physical constants and definition of custom units.
+    aux_features: additional static data.
+  """
+
+  coords: coordinate_systems.CoordinateSystem
+  dt: float
+  physics_specs: Any
+  aux_features: typing.AuxFeatures
+
+  def __len__(self):
+    return len(dataclasses.fields(self))
+
+  def __iter__(self):
+    return iter(f.name for f in dataclasses.fields(self))
+
+  def __getitem__(self, key):
+    return getattr(self, key)
+
+
+@gin.configurable(
+    allowlist=(
+        'model_time_step',
+        'custom_coords',
+        'reference_temperature',
+        'reference_datetime_str',
+    )
+)
+def get_model_specs(
+    data_coords: coordinate_systems.CoordinateSystem,
+    physics_specs: Any,
+    aux_features: typing.AuxFeatures,
+    model_time_step: Optional[Union[float, QuantityOrStr]] = None,
+    custom_coords: Optional[coordinate_systems.CoordinateSystem] = None,
+    reference_temperature: Optional[float | Sequence] = None,
+    reference_datetime_str: Optional[str] = None,
+) -> ModelSpecs:
+  """Returns specifications for a WhirlModel configuration.
+
+  Provides gin hooks, and in some cases defaults, for model specification
+  formerly encoded in aux_features.
+
+  Args:
+    data_coords: coordinate system in which states are represented in the data.
+    physics_specs: physical constants and definition of custom units.
+    aux_features: auxiliary features that come with the dataset.
+    model_time_step: duration of the outer time-step in our model, i.e., the
+      time by which the state is advanced in a single model.advance call.
+    custom_coords: optional coordinate system to be used by the model instead of
+      data_coords.
+    reference_temperature: reference temperature to use for sigma coordinates.
+      Must be None if already defined in aux_features.  Default value of 288
+      used if None and also not in aux_features.
+    reference_datetime_str: reference datetime for which nondimensionalized time
+      is set to 0. Must be None if already defined in aux_features. Default
+      value of '1979-01-01T00' used if None and also not in aux_features.
+
+  Returns:
+    Configured specification of coordinate system, time-step, physical constants
+    and units, and aux_features and for our hybrid ML/physics model.
+  """
+  if model_time_step is None:
+    raise ValueError('must provide model_time_step or outer_time_step')
+
+  if custom_coords is None:
+    coords = data_coords
+  else:
+    coords = dataclasses.replace(custom_coords, spmd_mesh=data_coords.spmd_mesh)
+
+  if aux_features.get(xarray_utils.REF_TEMP_KEY) is None:
+    if reference_temperature is None:
+      ones = np.ones(coords.vertical.layers, np.float32)
+      ref_temps = DEFAULT_REFERENCE_TEMPERATURE * ones
+      aux_features[xarray_utils.REF_TEMP_KEY] = ref_temps
+    else:
+      ones = np.ones(coords.vertical.layers, np.float32)
+      ref_temps = np.asarray(reference_temperature)
+      if ref_temps.ndim == 1 and ref_temps.shape[0] != coords.vertical.layers:
+        raise ValueError(
+            '`ref_temps` must be a scalar or a sequence with '
+            f'{coords.vertical.layers=} elements, got {ref_temps.shape=}'
+        )
+      ref_temps = ref_temps * ones
+      aux_features[xarray_utils.REF_TEMP_KEY] = ref_temps
+  else:  # cannot set ref temp if already specified in aux_data
+    if reference_temperature is not None:
+      raise ValueError(
+          'reference temperature already specified in aux_features'
+      )
+
+  if aux_features.get(xarray_utils.REFERENCE_DATETIME_KEY) is None:
+    if reference_datetime_str is None:
+      reference_datetime = np.datetime64(DEFAULT_REFERENCE_DATETIME_STR)
+      aux_features[xarray_utils.REFERENCE_DATETIME_KEY] = reference_datetime
+    else:
+      reference_datetime = np.datetime64(reference_datetime_str)
+      aux_features[xarray_utils.REFERENCE_DATETIME_KEY] = reference_datetime
+  else:  # cannot set ref datetime if already specified in aux_data
+    if reference_datetime_str is not None:
+      raise ValueError('reference datetime already specified in aux_data')
+
+  if isinstance(model_time_step, (str, scales.Quantity)):  # pyrefly: ignore[invalid-argument]
+    dt = physics_specs.nondimensionalize(scales.Quantity(model_time_step))
+  else:
+    dt = model_time_step
+
+  return ModelSpecs(
+      coords=coords,
+      dt=dt,
+      physics_specs=physics_specs,
+      aux_features=aux_features,
+  )
+
+
+def _identity(x):
+  return x
+
+
+class DynamicalSystem(hk.Module):
+  """Abstract class for modeling dynamical systems."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      output_coords: coordinate_systems.CoordinateSystem,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.dt = dt
+    self.physics_specs = physics_specs
+    self.aux_features = aux_features
+    self.input_coords = input_coords
+    self.output_coords = output_coords
+
+  def encode(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    """Encodes input trajectory `x` with `forcing` to the model state."""
+    raise NotImplementedError('Model subclass did not define encode')
+
+  def decode(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    """Decodes a model state `x` with `forcing` to a data representation."""
+    raise NotImplementedError('Model subclass did not define decode')
+
+  def advance(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    """Returns a model state `x` with `forcing` advanced by `self.dt`."""
+    raise NotImplementedError('Model subclass did not define advance')
+
+  def forcing_fn(self, forcing_data: ForcingData, sim_time: Numeric) -> Forcing:
+    """Returns forcing at sim_time, possibly using `forcing_data`."""
+    raise NotImplementedError('Model subclass did not define forcing_fn')
+
+  def trajectory(
+      self,
+      x,
+      outer_steps: int,
+      inner_steps: int = 1,
+      *,
+      forcing_data: ForcingData,
+      start_with_input: bool = False,
+      post_process_fn: Callable = _identity,
+  ):
+    """Returns a final model state and trajectory."""
+
+    def step_fn(x: PyTreeState) -> PyTreeState:
+      # if x does not have `sim_time`, expect forcing_fn to handle sim_time=None
+      if isinstance(x, typing.ModelState):
+        sim_time = getattr(x.state, 'sim_time', None)
+      else:
+        sim_time = getattr(x, 'sim_time', None)
+      forcing = self.forcing_fn(forcing_data, sim_time)  # pyrefly: ignore[bad-argument-type]
+      x, forcing = self.coords.with_dycore_sharding((x, forcing))
+      y = self.advance(x, forcing)
+      y = self.coords.with_dycore_sharding(y)
+      return y
+
+    return trajectory_from_step(
+        step_fn,
+        outer_steps,
+        inner_steps,
+        start_with_input=start_with_input,
+        post_process_fn=post_process_fn,
+    )(x)
+
+
+@gin.configurable
+class ModularStepModel(DynamicalSystem):
+  """Dynamical model based on independent encoder/decoder/step components."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      output_coords: coordinate_systems.CoordinateSystem,
+      advance_module=gin.REQUIRED,
+      encoder_module=gin.REQUIRED,
+      decoder_module=gin.REQUIRED,
+      forcing_module=forcings.NoForcing,
+      name: Optional[str] = None,
+  ):
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords,
+        output_coords,
+        name=name,
+    )
+    self.advance_fn = advance_module(coords, dt, physics_specs, aux_features)  # pyrefly: ignore[not-callable]
+    self.encoder_fn = encoder_module(  # pyrefly: ignore[not-callable]
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+    self.decoder_fn = decoder_module(  # pyrefly: ignore[not-callable]
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+    self.forcing_fn = forcing_module(coords, dt, physics_specs, aux_features)
+
+  def encode(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    return self.encoder_fn(x, forcing)
+
+  def decode(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    return self.decoder_fn(x, forcing)
+
+  def advance(self, x: PyTreeState, forcing: Forcing) -> PyTreeState:
+    return self.advance_fn(x, forcing)
+
+
+@gin.configurable
+class StochasticModularStepModel(DynamicalSystem):
+  """Dynamical model with modular components and stochasticity.
+
+  This instance of DynamicalSystem works with ModelState
+  representation of the model state. The `advance_module` initializes a
+  RandomnessModule. This must be compatible with ModelState.
+  Since randomness initialization might depend on the timestep at which it is
+  evolved, RandomnessModule module is initialized with `num_substeps`.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Dict[str, Any],
+      input_coords: coordinate_systems.CoordinateSystem,
+      output_coords: coordinate_systems.CoordinateSystem,
+      advance_module=gin.REQUIRED,
+      encoder_module=gin.REQUIRED,
+      decoder_module=gin.REQUIRED,
+      forcing_module=forcings.NoForcing,
+      name: Optional[str] = None,
+  ):
+    super().__init__(
+        coords,
+        dt,
+        physics_specs,
+        aux_features,
+        input_coords,
+        output_coords,
+        name=name,
+    )
+    self.advance_fn = advance_module(coords, dt, physics_specs, aux_features)  # pyrefly: ignore[not-callable]
+    self.encoder_fn = encoder_module(  # pyrefly: ignore[not-callable]
+        coords, dt, physics_specs, aux_features, input_coords
+    )
+    self.decoder_fn = decoder_module(  # pyrefly: ignore[not-callable]
+        coords, dt, physics_specs, aux_features, output_coords
+    )
+    self.forcing_fn = forcing_module(coords, dt, physics_specs, aux_features)
+
+  def encode(
+      self,
+      x: DataState,
+      forcing: Forcing,
+  ) -> ModelState:
+    """Encodes model state and creates a new perturbation."""
+    model_state = self.encoder_fn(x, forcing=forcing)
+    # encoder_fn returns `ModelState` that contains prognostic state
+    # and initial values for memory, diagnostics and randomness.
+    return self.advance_fn.finalize_state(model_state, forcing)
+
+  def decode(self, x: ModelState, forcing: Forcing) -> typing.Pytree:
+    """Returns model state with perturbation component removed."""
+    # TODO(langmore) Consider propagating decoding fields so decoder noise at
+    # different lead times is correlated.
+    return self.decoder_fn(x, forcing=forcing)
+
+  def advance(
+      self,
+      x: ModelState,
+      forcing: Forcing,
+  ) -> ModelState:
+    """Advances model state."""
+    return self.advance_fn(x, forcing)
+
+
+@gin.configurable(
+    allowlist=(
+        'checkpoint_step',
+        'checkpoint_multistep',
+        'checkpoint_post_process',
+    )
+)
+def trajectory_from_step(
+    step_fn: Callable,
+    outer_steps: int,
+    inner_steps: int,
+    *,
+    start_with_input: bool,
+    post_process_fn: Callable,
+    checkpoint_step: bool = True,
+    checkpoint_multistep: bool = False,
+    checkpoint_post_process: bool = True,
+) -> Callable:
+  """Returns a function that accumulates repeated applications of `step_fn`.
+
+  Compute a trajectory by repeatedly calling `step_fn()`
+  `outer_steps * inner_steps` times.
+
+  Args:
+    step_fn: function that takes a state and returns state after one time step.
+    outer_steps: number of steps to save in the generated trajectory.
+    inner_steps: number of repeated calls to step_fn() between saved steps.
+    start_with_input: if True, output the trajectory at steps [0, ..., steps-1]
+      instead of steps [1, ..., steps].
+    post_process_fn: function to apply to trajectory outputs.
+    checkpoint_step: whether to use `jax.checkpoint` on `step_fn`.
+    checkpoint_multistep: weather to use `jax.checkpoint` on `step_fn` repeated
+      steps between outputting observations used in the loss. Multi-step
+      checkpointing is off by default; turn it on to trade off ~25% increased
+      computed for ~25% less memory usage.
+    checkpoint_post_process: whether to use `jax.checkpoint` on
+      `post_process_fn`. `checkpoint_post_process` is a no-op if multi-step
+      checkpointing is enabled.
+
+  Returns:
+    A function that takes an initial state and returns a tuple consisting of:
+      (1) the final frame of the trajectory.
+      (2) trajectory of length `outer_steps` representing time evolution.
+  """
+  if checkpoint_step:
+    step_fn = hk.remat(step_fn)
+
+  if checkpoint_post_process:
+    post_process_fn = hk.remat(post_process_fn)
+
+  if checkpoint_multistep:
+
+    def outer_scan_fn(f, init, xs, length=None):
+      return hk.scan(hk.remat(f), init, xs, length=length)
+
+  else:
+    outer_scan_fn = hk.scan
+
+  return time_integration.trajectory_from_step(
+      step_fn,
+      outer_steps,
+      inner_steps,
+      start_with_input=start_with_input,
+      post_process_fn=post_process_fn,
+      inner_scan_fn=hk.scan,
+      outer_scan_fn=outer_scan_fn,
+  )
+
+
+@gin.configurable(allowlist=('model_cls', 'to_xarray_fn', 'from_xarray_fn'))
+class WhirlModel:
+  """Class that holds a Haiku model class and xarray conversion methods."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Optional[AuxFeatures] = None,
+      input_coords: Optional[coordinate_systems.CoordinateSystem] = None,
+      output_coords: Optional[coordinate_systems.CoordinateSystem] = None,
+      model_cls: Callable[[], DynamicalSystem] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      to_xarray_fn: Optional[Callable[..., xarray.Dataset]] = None,
+      from_xarray_fn: Optional[Callable[..., DataState]] = None,
+  ):
+    """Constructs pre-defined model functions and holds conversion functions.
+
+    Args:
+      coords: horizontal and vertical descritization.
+      dt: time step of the model.
+      physics_specs: object describing the scales and physical constants.
+      aux_features: dictionary holding static features that the model may use.
+      input_coords: horizontal and vertical descritization of the input data. if
+        `None`, uses `coords`. Default `None.
+      output_coords: horizontal and vertical descritization for the output data.
+        if `None`, uses `coords`. Default `None.
+      model_cls: model Haiku class that implements encode/advance/decode fns.
+      to_xarray_fn: function that converts decoded data slices to xarray.
+      from_xarray_fn: function that extracts data slices from xarray.
+    """
+    if aux_features is None:
+      aux_features = {}
+    if input_coords is None:
+      input_coords = coords
+    if output_coords is None:
+      output_coords = coords
+    self._coords = coords
+    self._data_coords = input_coords  # by data coords we refer to model inputs.
+    specs = ModelSpecs(coords, dt, physics_specs, aux_features)
+    model_cls = functools.partial(
+        model_cls,
+        coords=coords,
+        dt=dt,
+        physics_specs=physics_specs,
+        aux_features=aux_features,
+        input_coords=input_coords,
+        output_coords=output_coords,
+    )
+
+    def forcing_fwd(forcing_data, sim_time):
+      return model_cls().forcing_fn(forcing_data, sim_time)  # pytype: disable=wrong-keyword-args
+
+    forcing_fn = hk.transform(forcing_fwd).apply
+    encode_fwd = lambda x, forcing: model_cls().encode(x, forcing)  # pytype: disable=wrong-keyword-args
+    encode_fn = hk.transform(encode_fwd).apply
+    decode_fwd = lambda x, forcing: model_cls().decode(x, forcing)  # pytype: disable=wrong-keyword-args
+    decode_fn = hk.transform(decode_fwd).apply
+    advance_fwd = lambda x, forcing: model_cls().advance(x, forcing)  # pytype: disable=wrong-keyword-args
+    advance_fn = hk.transform(advance_fwd).apply
+    if to_xarray_fn is not None:
+      to_xarray_fn = functools.partial(to_xarray_fn, coords=output_coords)
+    self.forcing_fn = forcing_fn
+    self.encode_fn = encode_fn
+    self.decode_fn = decode_fn
+    self.advance_fn = advance_fn
+    self.specs = specs
+    self.model_cls = model_cls
+    self.to_xarray_fn = to_xarray_fn
+    self.from_xarray_fn = from_xarray_fn
+
+  @property
+  def coords(self) -> coordinate_systems.CoordinateSystem:
+    return self._coords
+
+  @property
+  def data_coords(self) -> coordinate_systems.CoordinateSystem:
+    return self._data_coords
+
+  def init_params(
+      self,
+      rng: Array,
+      input_trajectory: typing.DataState,
+      forcing_data: ForcingData,
+  ) -> typing.Params:
+    """Returns model parameters by initializing encode/advance/decode fn."""
+
+    def fwd(x):
+      model = self.model_cls()  # pytype: disable=wrong-keyword-args
+      decode = model_utils.with_forcing(
+          model.decode, model.forcing_fn, forcing_data
+      )
+      advance = model_utils.with_forcing(
+          model.advance, model.forcing_fn, forcing_data
+      )
+      encode = model_utils.with_forcing(
+          model.encode, model.forcing_fn, forcing_data
+      )
+      return decode(advance(encode(x)))
+
+    hk_model = hk.transform(fwd)
+    return hk_model.init(rng, input_trajectory)
+
+
+def get_whirl_model(
+    data_ds: xarray.Dataset,
+    model_config_str: str,
+    additional_gin_bindings: Optional[list[str]] = None,
+) -> WhirlModel:
+  """Returns a configured WhirlModel."""
+  if additional_gin_bindings is None:
+    additional_gin_bindings = []
+
+  try:
+    data_aux_features = xarray_utils.aux_features_from_xarray(data_ds)
+  except KeyError:
+    data_aux_features = {}
+
+  if 'physics_config_str' in data_ds.attrs:
+    physics_config_str = data_ds.attrs['physics_config_str']
+  else:
+    physics_config_str = ''  # empty string is equivalent to skipping.
+
+  gin.enter_interactive_mode()
+  gin.clear_config()
+  gin_utils.parse_gin_config(
+      physics_config_str,
+      model_config_str,
+      override_physics_configs_from_data=True,
+      gin_bindings=additional_gin_bindings,
+  )
+
+  data_coords = coordinate_system_from_dataset(data_ds)
+  physics_specs = physics_specifications.get_physics_specs()
+  model_specs = get_model_specs(data_coords, physics_specs, data_aux_features)
+  return WhirlModel(
+      coords=model_specs.coords,
+      dt=model_specs.dt,
+      physics_specs=model_specs.physics_specs,
+      aux_features=model_specs.aux_features,
+      input_coords=data_coords,
+      output_coords=data_coords,
+  )
+
+
+_ECMWF_CUTOFFS = {
+    # On Palmer 2009 (http://shortn/_56HCcQwmSS) page 4, the cutoffs for
+    # perturbations are given. Here we translate them to sigma levels.
+    # low_cutoffs: (100hPa, 50hPa)
+    'low_cutoffs': (0.05, 0.1),  # Will not be accurate over topography.
+    # high_cutoffs: (1300m, 300m)
+    'high_cutoffs': (0.86, 0.965),
+}
+
+
+def _piecewise_squasher(
+    sigma: Array,
+    low_cutoffs: Sequence[float],
+    high_cutoffs: Sequence[float],
+) -> Array:
+  """Piecewise linear values used to "squash" values by sigma level.
+
+  See function χ definition at: http://screen/5V3jzU7ZFA4vVJP
+
+  Args:
+    sigma: 1-D array of values for sigma levels. Should be in [0, 1].
+    low_cutoffs: σ=low_cutoffs[0] is when χ starts linearly increasing from 0.
+      σ=low_cutoffs[1] is when χ levels out at 1
+    high_cutoffs: σ=high_cutoffs[0] is when χ starts linearly decreasing from 1.
+      σ=high_cutoffs[1] is when χ reaches 0.
+
+  Returns:
+    Values χ of shape `sigma.shape + (1, 1)` that should be multiplied by
+      arrays of shape (n_levels, K, L) to "squash" high/low σ values.
+  """
+  if sigma.ndim != 1:
+    raise ValueError(f'{sigma.shape=} but should have been a 1-D array')
+  if len(low_cutoffs) != 2:
+    raise ValueError(f'{len(low_cutoffs)=} but should have been 2.')
+  if len(high_cutoffs) != 2:
+    raise ValueError(f'{len(high_cutoffs)=} but should have been 2.')
+
+  low_func = (sigma - low_cutoffs[0]) / (low_cutoffs[1] - low_cutoffs[0])
+  high_func = (high_cutoffs[1] - sigma) / (high_cutoffs[1] - high_cutoffs[0])
+
+  # lower_bound is a function equal to the squasher between
+  #   low_cutoffs[0] and high_cutoffs[1].
+  # It becomes negative outside that range.
+  lower_bound = jnp.minimum(1.0, jnp.minimum(low_func, high_func))
+  return jnp.maximum(0.0, lower_bound)[:, jnp.newaxis, jnp.newaxis]

model/legacy/model_utils.py

@@ -0,0 +1,527 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helper methods for constructing augmented trajectory functions."""
+
+import dataclasses
+import functools
+from typing import Any, Callable, Sequence, Tuple
+from dinosaur import coordinate_systems
+from dinosaur import pytree_utils
+from dinosaur import typing
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+Array = typing.Array
+DynamicalSystem = Any  # to prevent circular dependency on model_builder
+Pytree = typing.Pytree
+
+tree_map = jax.tree_util.tree_map
+tree_leaves = jax.tree_util.tree_leaves
+
+# Linter confused by wrapped functions
+# pylint: disable=g-bare-generic
+
+
+def with_preprocessing(
+    fn: Callable[..., Pytree],
+    preprocess_fn: Callable,
+) -> Callable[..., Pytree]:
+  """Generates a function that computes `fn` on `preprocess_fn(x)`."""
+  @functools.wraps(fn)
+  def apply_fn(x, *args, **kwargs):
+    return fn(preprocess_fn(x), *args, **kwargs)
+
+  return apply_fn
+
+
+def with_post_processing(
+    fn: Callable[..., Pytree],
+    post_process_fn: Callable,
+) -> Callable[..., Pytree]:
+  """Generates a function that applies `post_process_fn` to outputs of `fn`."""
+  @functools.wraps(fn)
+  def apply_fn(*args, **kwargs):
+    return post_process_fn(*fn(*args, **kwargs))
+
+  return apply_fn
+
+
+def with_forcing(
+    fn: Callable[..., Pytree],
+    forcing_fn: typing.ForcingFn,
+    forcing_data: typing.ForcingData,
+) -> Callable[..., Pytree]:
+  """Converts fn(x, forcing) to fn(x)."""
+  # evaluates forcing=forcing_fn(forcing_data, x.sim_time)
+  # when x does not have sim_time, forcing_fn will get sim_time=None
+  @functools.wraps(fn)
+  def wrapped(x, forcing_fn=forcing_fn):
+    # handle dataclass or dict for state data
+    if dataclasses.is_dataclass(x):
+      if isinstance(x, typing.ModelState):
+        sim_time = dataclasses.asdict(x.state).get('sim_time', None)
+      else:
+        sim_time = dataclasses.asdict(x).get('sim_time', None)
+    else:
+      sim_time = x.get('sim_time', None)
+    # handle sim_time of ndim 0 or 1
+    if sim_time is not None:
+      sim_time = jax.numpy.asarray(sim_time)
+      if sim_time.ndim:
+        forcing_fn = jax.vmap(forcing_fn, in_axes=(None, 0))
+    forcing = forcing_fn(forcing_data, sim_time)  # pyrefly: ignore[bad-argument-type]
+    return fn(x, forcing=forcing)
+  return wrapped
+
+
+def with_split_input(
+    fn: Callable[..., Pytree],
+    split_index: int,
+    time_axis: int = 0,
+) -> Callable[..., Pytree]:
+  """Decorates `fn` to be evaluated on first `split_index` time slices.
+
+  The returned function is a generalization to pytrees of the function:
+  `fn(x[:split_index], *args, **kwargs)`
+
+  Args:
+    fn: function to be transformed.
+    split_index: number of input elements along the time axis to use.
+    time_axis: axis corresponding to time dimension in `x` to decorated `fn`.
+
+  Returns:
+    decorated `fn` that is evaluated on only `split_index` first time slices of
+    provided inputs.
+  """
+  @functools.wraps(fn)
+  def apply_fn(x, *args, **kwargs):
+    init, _ = pytree_utils.split_along_axis(x, split_index, axis=time_axis)
+    return fn(init, *args, **kwargs)
+
+  return apply_fn
+
+
+def with_input_included(
+    trajectory_fn: typing.TrajectoryFn,
+    time_axis: int = 0,
+    num_last_input_frames_to_trim: int = 0,
+) -> typing.TrajectoryFn:
+  """Returns a `trajectory_fn` that concatenates inputs `x` to trajectory."""
+  if num_last_input_frames_to_trim > 0:
+    num_last_input_frames_to_trim = -num_last_input_frames_to_trim
+  else:
+    num_last_input_frames_to_trim = None  # pyrefly: ignore[bad-assignment]
+  inputs_time_slice = slice(None, num_last_input_frames_to_trim)
+  @functools.wraps(trajectory_fn)
+  def _trajectory(x, *args, **kwargs):
+    final, unroll = trajectory_fn(x, *args, **kwargs)
+    x_concat = pytree_utils.slice_along_axis(x, time_axis, inputs_time_slice)
+    return final, pytree_utils.concat_along_axis([x_concat, unroll], time_axis)
+
+  return _trajectory
+
+
+def trajectory_with_inputs_and_forcing(
+    model: DynamicalSystem,
+    num_init_frames: int,
+    start_with_input: bool = False,
+) -> typing.TrajectoryFn:
+  """Returns trajectory_fn that comuptes model trajectory from target data.
+
+  Wraps the default model.trajectory_fn to operate on data representation. It
+  corresponds to slicing `num_init_frames` from the inputs, encoding and
+  unrolling the trajectory.
+
+  Args:
+    model: model of a dynamical system used to obtain the trajectory.
+    num_init_frames: number of time frames used from the physics trajectory to
+      initialize the model state.
+    start_with_input: whether the firest decoded step in the output trajectory
+      should correspond to last input time or first future output.
+
+  Returns:
+    Trajectory function that operates on target data trajectory by encoding
+    the `initial_frames` inputs and unrolls trajectory in a model space.
+  """
+  def _trajectory_fn(x, forcing_data, outer_steps, inner_steps=1):
+
+    # configure the model.trajectory function with a decoder on the output.
+    trajectory_fn = functools.partial(
+        model.trajectory,
+        outer_steps=outer_steps,
+        inner_steps=inner_steps,
+        forcing_data=forcing_data,
+        start_with_input=start_with_input)
+    # add preprocessing to encode input to model state.
+    encode_fn = with_forcing(model.encode, model.forcing_fn, forcing_data)
+    trajectory_fn = with_preprocessing(trajectory_fn, encode_fn)
+    trajectory_fn = with_split_input(trajectory_fn, num_init_frames)
+    return trajectory_fn(x)
+
+  return _trajectory_fn
+
+
+def trajectory_with_inputs_and_forcing_and_stop_gradients(
+    model: DynamicalSystem,
+    num_init_frames: int,
+    start_with_input: bool = False,
+    stop_gradient_outer_steps: Sequence[int] = (),
+) -> typing.TrajectoryFn:
+  """Returns trajectory_fn that comuptes model trajectory from target data.
+
+  This extension of `trajectory_with_inputs_and_forcing` allows adding stop
+  gradients to the trajectory at designated steps. For example, if
+  `stop_gradient_outer_steps = [2]`, then gradients along the trajectory stop
+  at t=2. This does not mean that gradients with respect to X[2] will be zero.
+  It simply means that, for t > 2, gradients of X[t] with respect to X[2] will
+  be zero.
+
+  Wraps the default model.trajectory_fn to operate on data representation. It
+  corresponds to slicing `num_init_frames` from the inputs, encoding and
+  unrolling the trajectory.
+
+  Args:
+    model: model of a dynamical system used to obtain the trajectory.
+    num_init_frames: number of time frames used from the physics trajectory to
+      initialize the model state.
+    start_with_input: whether the firest decoded step in the output trajectory
+      should correspond to last input time or first future output.
+    stop_gradient_outer_steps: Tuple (possibly empty) indicating outer steps at
+      which to place stop gradients.
+
+  Returns:
+    Trajectory function that operates on target data trajectory by encoding
+    the `initial_frames` inputs and unrolls trajectory in a model space.
+    Decoding is not done by this function.
+  """
+  stop_gradient_outer_steps = list(sorted(stop_gradient_outer_steps))
+  if num_init_frames != 1:
+    raise ValueError(f'{num_init_frames=} is not supported yet.')
+
+  if stop_gradient_outer_steps and min(stop_gradient_outer_steps) <= 0:
+    raise ValueError(
+        f'{stop_gradient_outer_steps=} contained non-positive values'
+    )
+
+  expand_dim0 = lambda tree: tree_map(lambda x_i: x_i[jnp.newaxis], tree)
+  concat_dim0 = lambda trees: pytree_utils.concat_along_axis(trees, axis=0)
+  slice_dim0 = lambda tree, idx: pytree_utils.slice_along_axis(
+      tree, axis=0, idx=idx
+  )
+
+  def concat_trajectories_with_stop_grads(
+      x, forcing_data, outer_steps, inner_steps=1
+  ):
+    if (
+        stop_gradient_outer_steps
+        and max(stop_gradient_outer_steps) > outer_steps
+    ):
+      raise ValueError(
+          f'{stop_gradient_outer_steps=} contained values > {outer_steps=}'
+      )
+    outer_steps_seq = list(stop_gradient_outer_steps)
+    if not outer_steps_seq or outer_steps_seq[-1] != outer_steps:
+      outer_steps_seq.append(outer_steps)
+
+    # The first leg needs to encode the input. So use
+    # trajectory_with_inputs_and_forcing, which does the encoding.
+    final_state, first_leg = trajectory_with_inputs_and_forcing(
+        model,
+        num_init_frames=num_init_frames,
+        start_with_input=start_with_input,
+    )(
+        x,
+        forcing_data=forcing_data,
+        outer_steps=outer_steps_seq[0],
+        inner_steps=inner_steps,
+    )
+
+    # At this point, sections contains times [0, ..., outer_steps_seq[0]]
+    sections = [
+        first_leg,
+    ]
+
+    # Subsequent legs do not need encoding, so use model.trajectory directly.
+    trajectory_fn = functools.partial(
+        model.trajectory,
+        inner_steps=inner_steps,
+        forcing_data=forcing_data,
+        start_with_input=start_with_input,
+    )
+    for i in range(1, len(outer_steps_seq)):
+      # outer_steps_seq[-1] may or may not be in stop_gradient_outer_steps.
+      # The other steps will be by construction.
+      assert set(outer_steps_seq[:-1]).issubset(stop_gradient_outer_steps)
+      stop_grad_at_start = outer_steps_seq[i - 1] in stop_gradient_outer_steps
+
+      initial_state = final_state
+
+      # this_leg contains times [outer_steps_seq[0]+1, ..., outer_steps_seq[1]]
+      final_state, this_leg = trajectory_fn(
+          jax.lax.stop_gradient(initial_state)
+          if stop_grad_at_start
+          else initial_state,
+          outer_steps=outer_steps_seq[i] - outer_steps_seq[i - 1],
+      )
+
+      if stop_grad_at_start and start_with_input:
+        # Replace the initial point that had a stop gradient on it.
+        this_leg = concat_dim0([
+            expand_dim0(initial_state),
+            slice_dim0(this_leg, idx=slice(1, None)),
+        ])
+      sections.append(this_leg)
+
+    return final_state, concat_dim0(sections)
+
+  return concat_trajectories_with_stop_grads
+
+
+def decoded_trajectory_with_forcing(
+    model: DynamicalSystem,
+    start_with_input: bool = False,
+) -> typing.TrajectoryFn:
+  """Returns trajectory_fn that comuptes decoded trajectory values.
+
+  Args:
+    model: model of a dynamical system used to obtain the trajectory.
+    start_with_input: whether the firest decoded step in the output trajectory
+      should correspond to last input time or first future output.
+
+  Returns:
+    Trajectory function that additionally decodes trajectory values.
+  """
+  def _trajectory_fn(x, forcing_data, outer_steps, inner_steps=1):
+
+    # configure the model.trajectory function with a decoder on the output.
+    trajectory_fn = functools.partial(
+        model.trajectory,
+        forcing_data=forcing_data,
+        post_process_fn=with_forcing(model.decode,
+                                     model.forcing_fn, forcing_data),
+        start_with_input=start_with_input)
+    return trajectory_fn(x, outer_steps, inner_steps)
+
+  return _trajectory_fn
+
+
+def decoded_trajectory_with_inputs_and_forcing(
+    model: DynamicalSystem,
+    num_init_frames: int,
+    start_with_input: bool = False,
+) -> typing.TrajectoryFn:
+  """Returns trajectory_fn operating on decoded input and forcing data.
+
+  The returned function uses `num_init_frames` of the physics space trajectory
+  provided as an input to model.encode_fn to initialize the model state, then
+  unrolls the trajectory of specified length that is decoded to the physics
+  space using `model.decode_fn`.
+
+  Args:
+    model: model of a dynamical system used to obtain the trajectory.
+    num_init_frames: number of time frames used from the physics trajectory to
+      initialize the model state.
+    start_with_input: whether the firest decoded step in the output trajectory
+      should correspond to last input time or first future output.
+
+  Returns:
+    Trajectory function that operates on physics space trajectories
+    and returns unrolls in physics space.
+  """
+  def _trajectory_fn(x, forcing_data, outer_steps, inner_steps=1):
+
+    # configure the model.trajectory function with a decoder on the output.
+    trajectory_fn = decoded_trajectory_with_forcing(model, start_with_input)
+    trajectory_fn = functools.partial(
+        trajectory_fn,
+        forcing_data=forcing_data,
+        outer_steps=outer_steps,
+        inner_steps=inner_steps)
+    # add preprocessing to encode input to model state.
+    trajectory_fn = with_preprocessing(
+        trajectory_fn, with_forcing(model.encode,
+                                    model.forcing_fn, forcing_data))
+    # concatenate input trajectory to output trajectory for easier comparison.
+    trajectory_fn = with_input_included(
+        trajectory_fn, num_last_input_frames_to_trim=int(start_with_input))
+    # make trajectories operate on full examples by splitting the init.
+    trajectory_fn = with_split_input(trajectory_fn, num_init_frames)
+    return trajectory_fn(x)
+
+  return _trajectory_fn
+
+
+def process_trajectory(
+    input_trajectory: Pytree,
+    process_fn: Callable[[Pytree], Pytree],
+) -> Pytree:
+  """Processes trajectory by applying `process_fn` along time axis."""
+  step_fn = lambda c, x: tuple([None, hk.remat(process_fn)(x)])
+  _, out = hk.scan(step_fn, None, xs=input_trajectory)
+  return out
+
+
+def _maybe_to_nodal_with_physics_sharding(x, /, coords):
+  x = coordinate_systems.maybe_to_nodal(x, coords)
+  x = coords.with_physics_sharding(x)
+  return x
+
+
+def _maybe_to_modal_with_physics_sharding(x, /, coords):
+  x = coordinate_systems.maybe_to_modal(x, coords)
+  x = coords.with_physics_sharding(x)
+  return x
+
+
+def compute_prediction_representations(
+    predicted_trajectory: typing.Pytree,
+    forcing_data: typing.ForcingData,
+    model: DynamicalSystem,
+) -> typing.TrajectoryRepresentations:
+  """Computes TrajectoryRepresentations for predicted trajectory.
+
+  Args:
+    predicted_trajectory: predictions on `model.coords` coordinates.
+    forcing_data: forcing data to be used for encode/decode transformations.
+    model: model used for conversion between representations.
+
+  Returns:
+    `TrajectoryRepresentations` for predictions.
+  """
+  decode_fn = with_forcing(model.decode, model.forcing_fn, forcing_data)
+  data_to_nodal = functools.partial(
+      _maybe_to_nodal_with_physics_sharding, coords=model.output_coords)
+  data_to_modal = functools.partial(
+      _maybe_to_modal_with_physics_sharding, coords=model.output_coords)
+  model_to_nodal = functools.partial(
+      _maybe_to_nodal_with_physics_sharding, coords=model.coords)
+  model_to_modal = functools.partial(
+      _maybe_to_modal_with_physics_sharding, coords=model.coords)
+  predicted_data_trajectory = process_trajectory(
+      predicted_trajectory, decode_fn)
+  # Note: we pass original prediction to the decoder, but use dict for outputs.
+  if isinstance(predicted_trajectory, typing.ModelState):
+    predicted_trajectory = predicted_trajectory.state
+  if dataclasses.is_dataclass(predicted_trajectory):
+    # Losses operate on dicts: convert struct to dict if needed.
+    predicted_trajectory = predicted_trajectory.asdict()
+  return typing.TrajectoryRepresentations(
+      data_nodal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          predicted_data_trajectory, data_to_nodal),
+      data_modal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          predicted_data_trajectory, data_to_modal),
+      model_nodal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          predicted_trajectory, model_to_nodal),
+      model_modal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          predicted_trajectory, model_to_modal),
+  )
+
+
+def compute_target_representations(
+    target_trajectory: typing.Pytree,
+    forcing_data: typing.ForcingData,
+    model: DynamicalSystem,
+) -> typing.TrajectoryRepresentations:
+  """Computes TrajectoryRepresentations for target trajectory.
+
+  Note: currently this method only supports models that use a single time slice
+  for initialization.
+
+  Args:
+    target_trajectory: target trajectory on `model.output_coords` coordinates.
+    forcing_data: forcing data to be used for encode/decode transformations.
+    model: model used for conversion between representations.
+
+  Returns:
+    `TrajectoryRepresentations` for predictions.
+  """
+  encode_slice_fn = with_forcing(model.encode, model.forcing_fn, forcing_data)
+  encode_fn = lambda tree: encode_slice_fn(  # pylint: disable=g-long-lambda.
+      jax.tree_util.tree_map(lambda x: jnp.expand_dims(x, 0), tree))
+  data_to_nodal = functools.partial(
+      _maybe_to_nodal_with_physics_sharding, coords=model.output_coords)
+  data_to_modal = functools.partial(
+      _maybe_to_modal_with_physics_sharding, coords=model.output_coords)
+  model_to_nodal = functools.partial(
+      _maybe_to_nodal_with_physics_sharding, coords=model.coords)
+  model_to_modal = functools.partial(
+      _maybe_to_modal_with_physics_sharding, coords=model.coords)
+  target_model_trajectory = process_trajectory(
+      target_trajectory, encode_fn)
+  if isinstance(target_model_trajectory, typing.ModelState):
+    target_model_trajectory = target_model_trajectory.state
+  if dataclasses.is_dataclass(target_model_trajectory):
+    # Losses operate on dicts: convert struct to dict if needed.
+    target_model_trajectory = target_model_trajectory.asdict()
+  return typing.TrajectoryRepresentations(
+      data_nodal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          target_trajectory, data_to_nodal),
+      data_modal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          target_trajectory, data_to_modal),
+      model_nodal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          target_model_trajectory, model_to_nodal),
+      model_modal_trajectory=process_trajectory(  # pyrefly: ignore[unexpected-keyword]
+          target_model_trajectory, model_to_modal),
+  )
+
+
+def compute_prediction_and_target_representations(
+    predicted_model_trajectory: typing.Pytree,
+    target_data_trajectory: typing.Pytree,
+    forcing_data: typing.ForcingData,
+    model: DynamicalSystem,
+) -> Tuple[typing.TrajectoryRepresentations, typing.TrajectoryRepresentations]:
+  """Computes TrajectoryRepresentations for predicted and target trajectories.
+
+  Note: currently this method only supports models that use a single time slice
+  for initialization. While computing all terms seems wasteful, once jit-ed
+  all unused computations are optimized away. It is also tempting to compute
+  all representations at once, but as of 2023-02-28 compiler doesn't manage to
+  remove unused computation from a single primitive.
+
+  Args:
+    predicted_model_trajectory: predictions on `model.coords` coordinates.
+    target_data_trajectory: target data on `model.output_coords` coordinates.
+    forcing_data: forcing data to be used for encode/decode transformations.
+    model: model used for conversion between representations.
+
+  Returns:
+    Tuple of `TrajectoryRepresentations` for predictions and targets.
+  """
+  prediction_representations = compute_prediction_representations(
+      predicted_model_trajectory, forcing_data, model)
+  target_representations = compute_target_representations(
+      target_data_trajectory, forcing_data, model)
+  return prediction_representations, target_representations
+
+
+@jax.custom_jvp
+def safe_sqrt(x: Array) -> jax.Array:
+  """Sqrt(x) with gradient = 0 for x near 0."""
+  return jnp.sqrt(x)
+
+
+@safe_sqrt.defjvp
+def safe_sqrt_jvp(
+    primals: Array,
+    tangents: Array,
+) -> tuple[jax.Array, jax.Array]:
+  (x,) = primals
+  (x_dot,) = tangents
+  primal_out = safe_sqrt(x)
+  eps = jnp.finfo(x.dtype).eps
+  safe_x = jnp.where(x > eps, x, 1.0)
+  tangent_out = jnp.where(x > eps, x_dot / (2 * safe_sqrt(safe_x)), 0)
+  return primal_out, tangent_out

model/legacy/optimization.py

@@ -0,0 +1,179 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Configurable optimizers from JAX."""
+import collections
+import re
+from typing import Sequence
+
+import gin
+import optax
+
+
+gin.external_configurable(optax.adabelief, module='optax')
+gin.external_configurable(optax.adam, module='optax')
+gin.external_configurable(optax.adamw, module='optax')
+
+gin.external_configurable(optax.constant_schedule, module='optax')
+gin.external_configurable(optax.join_schedules, module='optax')
+gin.external_configurable(optax.piecewise_constant_schedule, module='optax')
+gin.external_configurable(optax.exponential_decay, module='optax')
+gin.external_configurable(
+    optax.warmup_exponential_decay_schedule, module='optax'
+)
+
+
+class OptimizerError(Exception):
+  """Raised if a custom Whirl optimizer encounters an error."""
+
+
+@gin.configurable
+def optimizer(value):
+  return value
+
+
+OptState = collections.namedtuple('OptState', ['state', 'params'])
+
+
+@gin.register
+def piecewise_constant_schedule_specified_by_rates(
+    rates: Sequence[float],
+    boundaries: Sequence[int],
+) -> optax.Schedule:
+  """Schedule that is piecewise constant and specified by rates (not scales).
+
+  This is similar to optax.piecewise_constant_schedule, which requires users
+  to specify "scales" (ratio of old LR to new LR).
+
+  Args:
+    rates: Length K sequence of learning rates. `rates[i]` is used for steps
+      `0 <= step < boundaries[1]`, for i=0, and
+      `boundaries[i-1] <= step < boundaries[i]`, for 0 < i < len(boundaries)
+      `boundaries[i-1] <= step < ∞`, for i = len(boundaries)
+    boundaries: Length K-1 sequence of boundaries.
+
+  Returns:
+    Schedule to pass to optax optimizers.
+  """
+  return optax.join_schedules(
+      schedules=[optax.constant_schedule(r) for r in rates],
+      boundaries=boundaries,
+  )
+
+
+@gin.register
+def delayed_constant_schedule(
+    turn_on_step: int,
+    rate: float,
+) -> optax.Schedule:
+  """Schedule that is zero until `turn_on_step` then `rate` thereafter."""
+  return piecewise_constant_schedule_specified_by_rates(
+      rates=[0., rate],
+      boundaries=[turn_on_step],
+  )
+
+
+@gin.register
+def top_level_multi_adam(
+    top_level_keys: Sequence[str] = (),
+    learning_rates: Sequence[optax.ScalarOrSchedule] = (),
+    default_learning_rate: optax.ScalarOrSchedule = 1e-4,
+    b1: float = 0.9,
+    b2: float = 0.95,
+    eps: float = 1e-6,
+    raise_if_keys_not_found: bool = True,
+) -> optax.GradientTransformation:
+  """Uses an Adam optimizer with different learning rates for different params.
+
+  Args:
+    top_level_keys: Keys to use non-default learning rates for. A key starting
+      with 'REGEX_', such as 'REGEX_cats' will use re.search to find keys, e.g.
+      re.search('cats', key).
+    learning_rates: Learning rates to use leafs under the `top_level_keys`.
+    default_learning_rate: Learning rate to use for keys not in `learning_rates`
+    b1: Exponential decay to track the first moment of past gradients.
+    b2: Exponential decay to track the second moment of past gradients.
+    eps:  A small constant applied to denominator outside of the square root to
+      avoid dividing by zero when rescaling.
+    raise_if_keys_not_found: Whether to raise if some `top_level_keys` are not
+      found in params.
+
+  Returns:
+    optax optimizer with learning rate based on top level key in params dict.
+  """
+  if len(top_level_keys) != len(learning_rates):
+    raise ValueError(
+        f'{top_level_keys=} had different length than {learning_rates=}'
+    )
+  if '' in top_level_keys:
+    raise ValueError('An empty string "" was found in `top_level_keys`.')
+
+  default_label = 'DEFAULT_LABEL'
+  if default_label in top_level_keys:
+    raise ValueError(f'{default_label=} should not be in `top_level_keys`')
+
+  def find_matching_top_level_key(param_name: str) -> str:
+    """Searches for param_name in top_level_keys, returns the matching key."""
+    prefix = 'REGEX_'
+    matches = []
+    for k in top_level_keys:
+      if k.startswith(prefix) and re.search(k.lstrip(prefix), param_name):
+        matches.append(k)
+      elif k == param_name:
+        matches.append(k)
+    if not matches:
+      return default_label
+    elif len(matches) == 1:
+      return matches[0]
+    else:
+      raise ValueError(
+          f'{param_name=} had more than 1 ({len(matches)}) match '
+          f'({matches}). Only one `top_level_keys` should match, or else we '
+          'cannot choose a unique learning rate for these parameters.'
+      )
+
+  def get_prefix_labels(params):
+    """Makes prefix labels to help optax match params with learning rates."""
+    # E.g. if top_level_keys = ['module_A', 'REGEX_special'],
+    # and params.keys() = ['module_A', 'special_A', 'special_B', 'module_C'],
+    # labels = {
+    #   'module_A': 'module_A',
+    #   'special_A': 'REGEX_special', 'special_B': 'REGEX_special',
+    #   'module_C': 'DEFAULT_LABEL', 'module_D': 'DEFAULT_LABEL',...
+    # }
+    # E.g. labels tells optax to use the learning rate 'REGEX_special' for
+    # parameters under the prefix 'module_C'.
+    labels = {
+        param_name: find_matching_top_level_key(param_name)
+        for param_name in params
+    }
+    top_level_keys_that_matched = [
+        k for k in labels.values() if k != default_label
+    ]
+    missing_keys = set(top_level_keys).difference(top_level_keys_that_matched)
+    if raise_if_keys_not_found and missing_keys:
+      raise OptimizerError(
+          f'{missing_keys=} not found in params: {sorted(params)}'
+      )
+    return labels
+
+  def make_adam(lr):
+    return optax.adam(lr, b1=b1, b2=b2, eps=eps)
+
+  return optax.multi_transform(
+      transforms={  # pyrefly: ignore[bad-argument-type]
+          k: make_adam(lr) for k, lr in zip(top_level_keys, learning_rates)
+      }
+      | {default_label: make_adam(default_learning_rate)},
+      param_labels=get_prefix_labels,
+  )

model/legacy/orographies.py

@@ -0,0 +1,131 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules responsible for orography processing and initialization."""
+
+from typing import Any, Callable, Mapping, Optional, Sequence
+from dinosaur import coordinate_systems
+from dinosaur import primitive_equations
+from dinosaur import scales
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+import haiku as hk
+import jax.numpy as jnp
+import numpy as np
+
+
+units = scales.units
+OrographyModule = Callable[..., typing.Array]
+FilterModule = Callable[..., typing.PyTreeFilterFn]
+
+
+@gin.register
+class ClippedOrography(hk.Module):
+  """Module that initializes orography by converting to modal and clipping."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      wavenumbers_to_clip: int = 1,
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs
+    super().__init__(name=name)
+    self.coords = coords
+    self.wavenumbers_to_clip = wavenumbers_to_clip
+    self.nodal_orography = aux_features.get(
+        xarray_utils.OROGRAPHY, np.zeros(coords.horizontal.nodal_shape))
+
+  def __call__(self) -> typing.Array:
+    """Returns orography converted to modal representation with clipping."""
+    return primitive_equations.truncated_modal_orography(
+        self.nodal_orography, self.coords, self.wavenumbers_to_clip)
+
+
+@gin.register
+class FilteredCustomOrography(hk.Module):
+  """Module that initializes orography from external data."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      orography_data_path: str,
+      filter_modules: Sequence[FilterModule] = tuple(),
+      renaming_dict: Optional[Mapping[str, str]] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    ds = xarray_utils.ds_from_path_or_aux(orography_data_path, aux_features)
+    if renaming_dict is not None:
+      ds = ds.rename(renaming_dict)
+    nodal_orography = xarray_utils.nodal_orography_from_ds(ds)
+    # TODO(dkochkov) Insist on having units specified in variable attrs.
+    self.nodal_orography = physics_specs.nondimensionalize(
+        nodal_orography * units.meter)  # pyrefly: ignore[unsupported-operation]
+    self.coords = coords
+    # Note: here we explicitly use linear truncation to preserve full signal.
+    # Smoothing is then achieved by interpolation to self.coords and filtering.
+    self.input_coords = xarray_utils.coordinate_system_from_dataset(
+        ds, truncation=xarray_utils.LINEAR, spmd_mesh=coords.spmd_mesh,
+        spherical_harmonics_impl=self.coords.horizontal.spherical_harmonics_impl
+    )
+    self.filter_fns = [
+        module(coords, dt, physics_specs, aux_features)
+        for module in filter_modules]
+
+  def __call__(self) -> typing.Array:
+    """Returns orography converted to modal representation with filtering."""
+    return primitive_equations.filtered_modal_orography(
+        self.nodal_orography, self.coords, self.input_coords, self.filter_fns)
+
+
+@gin.register
+class LearnedOrography(hk.Module):
+  """Module that uses learned parameters to correct orography."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      base_orography_module: OrographyModule,
+      correction_scale: float,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.base_orography_fn = base_orography_module(
+        coords, dt, physics_specs, aux_features)
+    self.scale = correction_scale
+    # coords.horizontal.modal_shape can change based upon the required amount of
+    # padding for a particular implementation of spherical harmonics, but the
+    # mask should always have the same number of non-zero elements in the same
+    # order.
+    self.correction = hk.get_parameter(
+        'orography', (coords.horizontal.mask.sum(),), jnp.float32,
+        init=hk.initializers.Constant(0.0))
+
+  def __call__(self) -> typing.Array:
+    """Returns orography in modal representation."""
+    mask = self.coords.horizontal.mask
+    correction_2d = jnp.zeros(self.coords.horizontal.modal_shape)
+    correction_2d = correction_2d.at[mask].set(self.correction)
+    return self.base_orography_fn() + correction_2d * self.scale  # pytype: disable=not-callable  # jax-ndarray

model/legacy/parameterizations.py

@@ -0,0 +1,171 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Physics parameterization modules that compute non-dynamical tendencies."""
+
+from typing import Any, Callable, Optional
+from dinosaur import coordinate_systems
+from dinosaur import pytree_utils
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+from model.legacy import features
+from model.legacy import mappings
+from model.legacy import transforms
+
+FeaturesModule = features.FeaturesModule
+Forcing = typing.Forcing
+MappingModule = mappings.MappingModule
+StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
+TransformModule = typing.TransformModule
+
+
+@gin.register
+class DirectNeuralParameterization(hk.Module):
+  """Computes modal physics tendencies from the input state and forcing."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: mappings.MappingModule,
+      tendency_transform_module: TransformModule,
+      prediction_mask: Optional[typing.Pytree] = None,
+      filter_module: Optional[StepFilterModule] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.prediction_mask = prediction_mask
+    self.modal_to_nodal_features_fn = modal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features)
+    self.nodal_mapping_module = nodal_mapping_module
+    self.tendency_transform_fn = tendency_transform_module(
+        coords, dt, physics_specs, aux_features)
+    if filter_module is not None:
+      self.filter_fn = filter_module(
+          coords, dt, physics_specs, aux_features)
+    else:
+      self.filter_fn = lambda _, y: y  # no filtering.
+
+  def __call__(
+      self,
+      inputs: typing.PyTreeState,
+      memory: Optional[typing.Pytree] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.Pytree] = None,
+      forcing: Optional[Forcing] = None,
+  ) -> typing.PyTreeState:
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    if memory is not None:
+      memory, _ = pytree_utils.as_dict(memory)
+    prediction_mask = self.prediction_mask
+    if prediction_mask is None:
+      prediction_mask = pytree_utils.tree_map_over_nonscalars(
+          lambda _: True, inputs, scalar_fn=lambda _: False  # pyrefly: ignore[bad-argument-type]
+      )
+    prediction_shapes = jax.tree_util.tree_map(
+        lambda x, y: x if y else None,
+        coordinate_systems.get_nodal_shapes(inputs, self.coords),
+        prediction_mask,
+    )
+    net = self.nodal_mapping_module(prediction_shapes)
+    nodal_inputs = self.modal_to_nodal_features_fn(
+        inputs, memory=memory, diagnostics=diagnostics, randomness=randomness,
+        forcing=forcing,
+    )
+    nodal_tendencies = net(nodal_inputs)
+    nodal_tendencies = self.tendency_transform_fn(nodal_tendencies)
+    modal_tendencies = self.coords.horizontal.to_modal(nodal_tendencies)
+    modal_tendencies = self.filter_fn(inputs, modal_tendencies)
+    return from_dict_fn(modal_tendencies)
+
+
+@gin.register
+class DivCurlNeuralParameterization(hk.Module):
+  """Computes modal physics tendencies via `u, v` → `δ, ζ`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      modal_to_nodal_features_module: FeaturesModule,
+      nodal_mapping_module: mappings.MappingModule,
+      tendency_transform_module: TransformModule,
+      prediction_mask: Optional[typing.Pytree] = None,
+      filter_module: Optional[StepFilterModule] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.prediction_mask = prediction_mask
+    self.modal_to_nodal_features_fn = modal_to_nodal_features_module(
+        coords, dt, physics_specs, aux_features)
+    self.nodal_mapping_module = nodal_mapping_module
+    self.tendency_transform_fn = tendency_transform_module(
+        coords, dt, physics_specs, aux_features)
+    self.get_nodal_shape_fn = (
+        lambda x: coordinate_systems.get_nodal_shapes(x, coords))
+    self.to_div_curl_fn = transforms.ToModalWithDivCurlTransform(
+        coords, dt, physics_specs, aux_features)
+    if filter_module is not None:
+      self.filter_fn = filter_module(
+          coords, dt, physics_specs, aux_features)
+    else:
+      self.filter_fn = lambda _, y: y  # no filtering.
+
+  def __call__(
+      self,
+      inputs: typing.PyTreeState,
+      memory: Optional[typing.Pytree] = None,
+      diagnostics: Optional[typing.Pytree] = None,
+      randomness: Optional[typing.Pytree] = None,
+      forcing: Optional[Forcing] = None,
+  ) -> typing.PyTreeState:
+    inputs = self.coords.with_dycore_sharding(inputs)
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    if memory is not None:
+      memory = self.coords.with_dycore_sharding(memory)
+      memory, _ = pytree_utils.as_dict(memory)
+    prediction_mask = self.prediction_mask
+    if prediction_mask is None:
+      prediction_mask = pytree_utils.tree_map_over_nonscalars(
+          lambda _: True, inputs, scalar_fn=lambda _: False  # pyrefly: ignore[bad-argument-type]
+      )
+    prediction_shapes = jax.tree_util.tree_map(
+        lambda x, y: self.get_nodal_shape_fn(x) if y else None,
+        inputs,
+        prediction_mask,
+    )
+    prediction_shapes['u'] = prediction_shapes.pop('divergence')
+    prediction_shapes['v'] = prediction_shapes.pop('vorticity')
+    net = self.nodal_mapping_module(prediction_shapes)
+    nodal_inputs = self.modal_to_nodal_features_fn(
+        inputs, memory=memory, diagnostics=diagnostics, randomness=randomness,
+        forcing=forcing,
+    )
+    nodal_inputs = self.coords.dycore_to_physics_sharding(nodal_inputs)
+    nodal_tendencies = net(nodal_inputs)
+    nodal_tendencies = self.coords.physics_to_dycore_sharding(nodal_tendencies)
+    nodal_tendencies = self.tendency_transform_fn(nodal_tendencies)
+    modal_tendencies = self.to_div_curl_fn(nodal_tendencies)
+    modal_tendencies = self.filter_fn(inputs, modal_tendencies)
+    outputs = from_dict_fn(modal_tendencies)
+    outputs = self.coords.with_dycore_sharding(outputs)
+    return outputs

model/legacy/perturbations.py

@@ -0,0 +1,230 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of perturbation modules."""
+import abc
+from typing import Any, Callable
+from dinosaur import coordinate_systems
+from dinosaur import pytree_utils
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+import gin
+import jax
+import jax.numpy as jnp
+from model.legacy import transforms
+
+Pytree = typing.Pytree
+PerturbationFn = Callable[..., Pytree]
+PerturbationModule = Callable[..., PerturbationFn]
+
+
+_ALLOWED_PERTURBATION_BASIS = (
+    # Converts vorticity/divergence to u/v then perturbs.
+    'uv',
+
+    # Perturbs in whatever the state is in (typically vorticity/divergence).
+    'generic',
+)
+
+# We ♥ λ's
+# pylint: disable=g-long-lambda
+
+
+@gin.register
+class NoPerturbation:
+  """No-op perturbation that introduces no perturbation to `inputs`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+  ):
+    """Initializes a random field."""
+    del coords, dt, physics_specs, aux_features  # unused.
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      state: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Updates the state of a random field."""
+    del state, randomness  # unused.
+    return inputs
+
+
+class BasePerturbation(abc.ABC):
+  """Base class for perturbations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      randomness_transform_module: transforms.TransformModule = (
+          transforms.IdentityTransform
+      ),
+      return_modal: bool = True,
+      perturbation_basis: str = 'generic',
+  ):
+    """Initializes module to perturb random fields.
+
+    Args:
+      coords: Model coordinate system.
+      dt: Time step.
+      physics_specs:
+      aux_features:
+      randomness_transform_module: Module that transforms jax.Array of random
+        variables before converting to nodal.
+      return_modal: Whether results should be returned in modal space.
+      perturbation_basis: Whether to perturb wind in "uv" or "generic" basis.
+    """
+    self.coords = coords
+    self.randomness_transform_fn = randomness_transform_module(
+        coords, dt, physics_specs, aux_features
+    )
+    self.return_modal = return_modal
+    self.to_modal = coords.horizontal.to_modal
+    self.to_nodal = coords.horizontal.to_nodal
+    self.maybe_to_modal = lambda tr: coordinate_systems.maybe_to_modal(
+        tr, coords
+    )
+    self.maybe_to_nodal = lambda tr: coordinate_systems.maybe_to_nodal(
+        tr, coords
+    )
+    if perturbation_basis not in _ALLOWED_PERTURBATION_BASIS:
+      raise ValueError(
+          f'{perturbation_basis=} which was not in '
+          f'{_ALLOWED_PERTURBATION_BASIS=}'
+      )
+    self.perturbation_basis = perturbation_basis
+
+  def __call__(
+      self,
+      inputs: typing.Pytree,
+      state: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Updates the state of a random field."""
+    del state  # unused.
+    # TODO(dkochkov) allow pytree randomness in addition to broadcasting option.
+
+    if self.perturbation_basis == 'generic':
+      return self._perturb_in_generic_coordinates(inputs, randomness)
+    elif self.perturbation_basis == 'uv':
+      return self._perturb_in_uv_coordinates(inputs, randomness)
+
+  @abc.abstractmethod
+  def _perturb_core(
+      self,
+      inputs: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Perturbs inputs using randomness."""
+
+  def _perturb_in_generic_coordinates(
+      self,
+      inputs: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Perturb `inputs` in (vorticity, divergence) coordinate system."""
+    nodal_inputs = self.maybe_to_nodal(inputs)
+    nodal_randomness = self.maybe_to_nodal(randomness)
+
+    nodal_randomness = self.randomness_transform_fn(
+        pytree_utils.tree_map_over_nonscalars(
+            # Broadcast randomness so that self.randomness_transform_fn can use
+            # the shape of x to determine what to do.
+            lambda x: jnp.broadcast_to(nodal_randomness, x.shape),
+            nodal_inputs,
+            scalar_fn=jnp.zeros_like,
+        )
+    )
+
+    perturbed_nodal_inputs = self._perturb_core(nodal_inputs, nodal_randomness)
+    if self.return_modal:
+      return self.to_modal(perturbed_nodal_inputs)
+    else:
+      return perturbed_nodal_inputs
+
+  def _perturb_in_uv_coordinates(
+      self,
+      inputs: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Perturb `inputs` in (u, v) coordinate system."""
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+
+    # Remove vorticity/divergence from inputs and replace with u/v.
+    vordiv = self.maybe_to_modal({
+        'vorticity': inputs.pop('vorticity'),
+        'divergence': inputs.pop('divergence'),
+    })
+    u_nodal, v_nodal = spherical_harmonic.vor_div_to_uv_nodal(
+        grid=self.coords.horizontal,
+        vorticity=vordiv['vorticity'],
+        divergence=vordiv['divergence'],
+        clip=True,
+    )
+    nodal_inputs = self.maybe_to_nodal(inputs)  # Recall we popped vor/div.
+    nodal_inputs['u'] = u_nodal
+    nodal_inputs['v'] = v_nodal
+
+    # Perturb in u/v space
+    nodal_randomness = self.maybe_to_nodal(randomness)
+    nodal_randomness = self.randomness_transform_fn(
+        pytree_utils.tree_map_over_nonscalars(
+            # Broadcast randomness so that self.randomness_transform_fn can use
+            # the shape of x to determine what to do.
+            lambda x: jnp.broadcast_to(nodal_randomness, x.shape),
+            nodal_inputs,
+            scalar_fn=jnp.zeros_like,
+        )
+    )
+    perturbed_nodal_inputs = self._perturb_core(nodal_inputs, nodal_randomness)
+
+    # Transform perturbed u/v to vor/div (modal).
+    vorticity, divergence = spherical_harmonic.uv_nodal_to_vor_div_modal(
+        grid=self.coords.horizontal,
+        u_nodal=perturbed_nodal_inputs.pop('u'),
+        v_nodal=perturbed_nodal_inputs.pop('v'),
+        clip=True,
+    )
+
+    # Insert vorticity/divergence into perturbed_inputs in the right space.
+    if self.return_modal:
+      perturbed_inputs = self.to_modal(perturbed_nodal_inputs)
+      perturbed_inputs['vorticity'] = vorticity
+      perturbed_inputs['divergence'] = divergence
+    else:
+      perturbed_inputs = perturbed_nodal_inputs.copy()
+      perturbed_inputs['vorticity'] = self.to_nodal(vorticity)
+      perturbed_inputs['divergence'] = self.to_nodal(divergence)
+
+    return from_dict_fn(perturbed_inputs)
+
+
+@gin.register
+class ProportionalPerturbation(BasePerturbation):
+  """Perturbation that scales inputs by 1 + randomness."""
+
+  def _perturb_core(
+      self,
+      inputs: typing.Pytree,
+      randomness: typing.Pytree,
+  ) -> typing.Pytree:
+    """Multiplies inputs by (1 + randomness)."""
+    return jax.tree_util.tree_map(lambda x, y: x * (1 + y), inputs, randomness)

model/legacy/physics_specifications.py

@@ -0,0 +1,115 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PhysicsSpecifications contain physical parameters of dynamical systems.
+
+To ensure that all model components the expected PhysicsSpecs all modules
+(except specializing on a particular equation) must instantiate
+PhysicsSpecs objects using `get_physics_specs`, which should be configured
+appropriately via `gin`.
+"""
+
+from typing import Sequence, Union
+from dinosaur import primitive_equations
+from dinosaur import scales
+from dinosaur import shallow_water
+import gin
+import numpy as np
+
+
+# TODO(jamieas): consolidate with `PrimitiveEquationSpecs`. In particular,
+# decide whether 'specs' should have units or be nondimensionalized.
+QuantityOrStr = Union[str, scales.Quantity]
+
+
+GET_DEFAULT_SCALE = gin.external_configurable(
+    lambda: scales.DEFAULT_SCALE, name='GET_DEFAULT_SCALE')
+GET_ATMOSPHERIC_SCALE = gin.external_configurable(
+    lambda: scales.ATMOSPHERIC_SCALE, name='GET_ATMOSPHERIC_SCALE')
+
+
+@gin.configurable
+def get_physics_specs(construct_fn=gin.REQUIRED):
+  """Returns physical parameters object generated by `construct_fn`."""
+  return construct_fn()  # pyrefly: ignore[not-callable]
+
+
+@gin.register
+def shallow_water_specs_constructor(
+    density_vals: Union[Sequence[float], np.ndarray],
+    density_units: QuantityOrStr = scales.WATER_DENSITY,
+    radius_si: QuantityOrStr = scales.RADIUS,
+    angular_velocity_si: QuantityOrStr = scales.ANGULAR_VELOCITY,
+    gravity_acceleration_si: QuantityOrStr = scales.GRAVITY_ACCELERATION,
+    scale: scales.Scale = scales.DEFAULT_SCALE
+) -> shallow_water.ShallowWaterSpecs:
+  """Constructs `ShallowWaterSpecs` using gin-configurable parameters.
+
+  Args:
+    density_vals: density values for each layer of the shallow water system.
+    density_units: units in which `density_vals` are specified.
+    radius_si: radius of the domain specified with units attached.
+    angular_velocity_si: angular velocity of the domain with units attached.
+    gravity_acceleration_si: gravity on the surface with units attached.
+    scale: a scale object specifying the scales to use for nondimensionalizing.
+
+  Returns:
+    ShallowWaterSpecs object containing physical parameters of the system.
+  """
+  densities = np.asarray(density_vals) * scales.Quantity(density_units)
+  return shallow_water.ShallowWaterSpecs.from_si(
+      densities=densities,  # pyrefly: ignore[unexpected-keyword]
+      radius_si=scales.Quantity(radius_si),
+      angular_velocity_si=scales.Quantity(angular_velocity_si),
+      gravity_acceleration_si=scales.Quantity(gravity_acceleration_si),
+      scale=scale)
+
+
+@gin.register
+def primitive_eq_specs_constructor(
+    radius_si: QuantityOrStr = scales.RADIUS,
+    angular_velocity_si: QuantityOrStr = scales.ANGULAR_VELOCITY,
+    gravity_acceleration_si: QuantityOrStr = scales.GRAVITY_ACCELERATION,
+    ideal_gas_constant_si: QuantityOrStr = scales.IDEAL_GAS_CONSTANT,
+    water_vapor_gas_constant_si: QuantityOrStr = scales.IDEAL_GAS_CONSTANT_H20,
+    water_vapor_isobaric_heat_capacity_si: QuantityOrStr = (
+        scales.WATER_VAPOR_CP),
+    kappa_si: QuantityOrStr = scales.KAPPA,
+    scale: scales.Scale = scales.DEFAULT_SCALE,
+) -> primitive_equations.PrimitiveEquationsSpecs:
+  """Constructs `PrimitiveEquationsSpecs` using gin-configurable parameters.
+
+  Args:
+    radius_si: radius of the domain with units attached.
+    angular_velocity_si: angular velocity of the domain with units attached.
+    gravity_acceleration_si: gravity on the surface with units attached.
+    ideal_gas_constant_si: the gas constant with units attached.
+    water_vapor_gas_constant_si: the gas constant for vapor with units attached.
+    water_vapor_isobaric_heat_capacity_si: isobaric heat capacity of vapor with
+      units attached.
+    kappa_si: `ideal_gas_constant / Cp` where  Cp is the isobaric heat capacity.
+    scale: a scale object specifying the scales to use for nondimensionalizing.
+
+  Returns:
+    PrimitiveEquationsSpecs object containing physical parameters of the system.
+  """
+  return primitive_equations.PrimitiveEquationsSpecs.from_si(
+      radius_si=scales.Quantity(radius_si),
+      angular_velocity_si=scales.Quantity(angular_velocity_si),
+      gravity_acceleration_si=scales.Quantity(gravity_acceleration_si),
+      ideal_gas_constant_si=scales.Quantity(ideal_gas_constant_si),
+      water_vapor_gas_constant_si=scales.Quantity(water_vapor_gas_constant_si),
+      water_vapor_isobaric_heat_capacity_si=scales.Quantity(
+          water_vapor_isobaric_heat_capacity_si),
+      kappa_si=scales.Quantity(kappa_si),
+      scale=scale)

model/legacy/steps.py

@@ -0,0 +1,332 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modules that parameterize composed time-steppers."""
+
+import abc
+import functools
+from typing import Any, Callable, Optional, Sequence
+from dinosaur import coordinate_systems
+from dinosaur import primitive_equations
+from dinosaur import time_integration
+from dinosaur import typing
+import gin
+import haiku as hk
+from model.legacy import diagnostics
+from model.legacy import integrators
+from model.legacy import perturbations
+from model.legacy import stochastic
+
+DiagnosticModule = diagnostics.DiagnosticModule
+Forcing = typing.Forcing
+Pytree = typing.Pytree
+ModelState = typing.ModelState
+EquationModule = Callable[..., time_integration.ImplicitExplicitODE]
+CorrectorModule = typing.CorrectorModule
+PerturbationModule = perturbations.PerturbationModule
+RandomnessModule = stochastic.RandomnessModule
+PyTreeStepFilterModule = typing.PyTreeStepFilterModule
+TimeIntegrator = integrators.TimeIntegrator
+TransformModule = typing.TransformModule
+
+
+class BaseStep(abc.ABC):
+  """Base class for Step modules."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      diagnostics_module: DiagnosticModule = diagnostics.NoDiagnostics,
+      randomness_module: RandomnessModule = stochastic.NoRandomField,
+  ):
+    self.diagnostics_fn = diagnostics_module(
+        coords, dt, physics_specs, aux_features)
+    self.randomness_fn = randomness_module(
+        coords, dt, physics_specs, aux_features)
+
+  @abc.abstractmethod
+  def __call__(
+      self,
+      state: ModelState,
+      forcing: typing.Forcing,
+  ) -> ModelState:
+    """Computes the state of the system evolved in time by `self.dt`."""
+
+  def finalize_state(
+      self,
+      x: ModelState,
+      forcing: typing.Forcing,
+  ) -> ModelState:
+    """Finalizes initialization of a model state `x`, encoded from data.
+
+    This method ensures that state has all of the `ModelState` fields
+    initialized in a way compatible with this step function. This includes
+    populating initial `diagnostics`, `memory` and `randomness` fields.
+
+    Args:
+      x: Initial values for the model state typically provided by the encoder.
+        forcing: Data covariates from the same time slice as `x`.
+
+    Returns:
+      Initialized model state.
+    """
+    x.randomness = self.randomness_fn.unconditional_sample(
+        hk.maybe_next_rng_key()
+    )
+    x.diagnostics = self.diagnostics_fn(
+        x, physics_tendencies=None, forcing=forcing)
+    return x
+
+
+@gin.register
+class EquationStep(BaseStep, hk.Module):
+  """Step module that advances the state by integrating an equation in time."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      equation_module: EquationModule,
+      time_integrator: TimeIntegrator = integrators.imex_rk_sil3,
+      filter_modules: Sequence[PyTreeStepFilterModule] = tuple(),
+      checkpoint_explicit_terms: bool = True,
+      name: Optional[str] = None,
+  ):
+    hk.Module.__init__(self, name=name)
+    BaseStep.__init__(self, coords, dt, physics_specs, aux_features)
+    equation = equation_module(coords, dt, physics_specs, aux_features)
+    if checkpoint_explicit_terms:
+      equation = time_integration.ImplicitExplicitODE.from_functions(
+          hk.remat(equation.explicit_terms),
+          equation.implicit_terms,
+          equation.implicit_inverse)  # pyrefly: ignore[bad-argument-type]
+    step_fn = time_integrator(equation, dt)
+    filter_fns = [
+        module(coords, dt, physics_specs, aux_features)
+        for module in filter_modules]
+    self.dt = dt
+    self.step_fn = time_integration.step_with_filters(step_fn, filter_fns)
+
+  def __call__(
+      self,
+      x: ModelState,
+      forcing: Optional[typing.Forcing] = None,
+  ) -> ModelState:
+    """Computes the state of the system evolved in time by `dt`."""
+    del forcing
+    next_state = time_integration.maybe_fix_sim_time_roundoff(
+        self.step_fn(x.state), self.dt)
+    return ModelState(next_state)  # pyrefly: ignore[bad-argument-count]
+
+
+@gin.register
+class RepeatedStep(BaseStep, hk.Module):
+  """Step module that consists of repeated substeps of the same form."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      inner_step_module: typing.StepModule,
+      num_inner_steps: int = 1,
+      name: Optional[str] = None,
+  ):
+    hk.Module.__init__(self, name=name)
+    BaseStep.__init__(self, coords, dt, physics_specs, aux_features)
+    inner_dt = dt / num_inner_steps
+    self.step_fn = inner_step_module(
+        coords, inner_dt, physics_specs, aux_features)
+    self.num_inner_steps = num_inner_steps
+
+  def __call__(
+      self,
+      state: ModelState,
+      forcing: typing.Forcing,
+  ) -> ModelState:
+    """Computes the state of the system evolved in time by `dt`."""
+    step_fn = functools.partial(self.step_fn, forcing=forcing)
+    step_fn = time_integration.repeated(step_fn, self.num_inner_steps, hk.scan)
+    return step_fn(state)
+
+
+@gin.register
+class CustomCoordsStep(BaseStep, hk.Module):
+  """Step module that uses gin-configured coordinates instead of coords.
+
+  This class currently supports model states in spectral representation. It
+  could be easily extended to nodal-state models by converting to modal space
+  prior to spectral interpolation and back after the timestep if performed.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      step_module: typing.StepModule,
+      custom_coords: coordinate_systems.CoordinateSystem = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      name: Optional[str] = None,
+  ):
+    hk.Module.__init__(self, name=name)
+    BaseStep.__init__(self, coords, dt, physics_specs, aux_features)
+    self.step_fn = step_module(
+        custom_coords, dt, physics_specs, aux_features)
+    self.to_custom_coords_fn = coordinate_systems.get_spectral_interpolate_fn(
+        coords, custom_coords)
+    self.from_custom_coords_fn = coordinate_systems.get_spectral_interpolate_fn(
+        custom_coords, coords)
+
+  def __call__(
+      self,
+      x: typing.PyTreeState,
+      forcing: typing.Forcing,
+  ) -> typing.PyTreeState:
+    del forcing  # currently not supported.
+    x = self.to_custom_coords_fn(x)
+    custom_out = self.step_fn(x, None)
+    return self.from_custom_coords_fn(custom_out)
+
+
+@gin.register
+class StochasticPhysicsParameterizationStep(BaseStep, hk.Module):
+  """Step module that uses stochastic physics tendencies with dycore."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      corrector_module: CorrectorModule,
+      physics_parameterization_module: typing.ParameterizationModule,
+      num_substeps: int = 1,
+      diagnostics_module: DiagnosticModule = diagnostics.NoDiagnostics,
+      randomness_module: RandomnessModule = stochastic.ZerosRandomField,
+      perturbation_module: PerturbationModule = perturbations.NoPerturbation,
+      checkpoint_substep: bool = False,
+      name: Optional[str] = None,
+  ):
+    hk.Module.__init__(self, name=name)
+    BaseStep.__init__(
+        self, coords, dt, physics_specs, aux_features,
+        diagnostics_module=diagnostics_module,
+        randomness_module=randomness_module)
+    inner_dt = dt / num_substeps
+    self.num_substeps = num_substeps
+    self.corrector_fn = corrector_module(
+        coords, inner_dt, physics_specs, aux_features)
+    self.physics_parameterization_fn = physics_parameterization_module(
+        coords, inner_dt, physics_specs, aux_features)
+    self.perturbation_fn = perturbation_module(
+        coords, inner_dt, physics_specs, aux_features)
+    self.checkpoint_substep = checkpoint_substep
+    self.coords = coords
+
+  def finalize_state(
+      self,
+      x: ModelState,
+      forcing: typing.Forcing,
+  ) -> ModelState:
+    """Finalizes initialization of a model state `x`, encoded from data.
+
+    This method ensures that state has all of the `ModelState` fields
+    initialized in a way compatible with this step function. This includes
+    populating initial `diagnostics`, `memory` and `randomness` fields.
+
+    This is called by StochasticModularStepModel.encode, after encoding the data
+
+    Args:
+      x: Initial values for the model state typically provided by the encoder.
+      forcing: Data covariates from the same time slice as `x`.
+
+    Returns:
+      Initialized model state.
+    """
+    # TODO(dkochkov) Consider adding an option of not overriding randomness.
+    x.randomness = self.randomness_fn.unconditional_sample(
+        hk.maybe_next_rng_key()
+    )
+    pp_tendency = self.physics_parameterization_fn(
+        x.state, x.memory, x.diagnostics, x.randomness.nodal_value, forcing
+    )
+    x.diagnostics = self.diagnostics_fn(x, pp_tendency, forcing)
+    return x
+
+  def __call__(
+      self,
+      state: ModelState,
+      forcing: typing.Forcing,
+  ) -> ModelState:
+    """Computes the state of the system evolved in time by `dt`."""
+
+    def step_fn(x):
+      x = self.coords.with_dycore_sharding(x)
+      # TODO(dkochkov) Consider passing `x` to physics_parameterization.
+      pp_tendency = self.physics_parameterization_fn(
+          x.state, x.memory, x.diagnostics, x.randomness.nodal_value, forcing
+      )
+
+      pp_tendency = self.perturbation_fn(
+          pp_tendency,
+          state=x.state,
+          randomness=x.randomness.nodal_value,
+      )
+
+      next_state = self.corrector_fn(x.state, pp_tendency, forcing)
+      # TODO(dkochkov) update stochastic modules to take optional state.
+      next_randomness = self.randomness_fn.advance(x.randomness)
+      next_memory = x.state if x.memory is not None else None
+      next_diagnostics = self.diagnostics_fn(x, pp_tendency, forcing)
+      x_next = ModelState(
+          state=next_state, memory=next_memory, diagnostics=next_diagnostics,  # pyrefly: ignore[unexpected-keyword]
+          randomness=next_randomness)  # pyrefly: ignore[unexpected-keyword]
+      x_next = self.coords.with_dycore_sharding(x_next)
+      return x_next
+
+    if self.checkpoint_substep:
+      step_fn = hk.remat(step_fn)
+    step_fn = time_integration.repeated(step_fn, self.num_substeps, hk.scan)
+    return step_fn(state)
+
+
+# TODO(dkochkov) Move vertical advection step to transforms.py.
+
+
+@gin.register
+class SemiLagrangianVerticalAdvectionStep(hk.Module):
+  """Step module that applies vertical advection for the primitive equations."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.dt = dt
+
+  def __call__(self, state):
+    return primitive_equations.semi_lagrangian_vertical_advection_step(
+        state, self.coords, self.dt
+    )

model/legacy/stochastic.py

@@ -0,0 +1,1225 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of stochastic modules."""
+
+import abc
+import dataclasses
+import enum
+import logging
+from typing import Any, Callable, Optional, Sequence, TypeVar, Union
+import zlib
+
+from dinosaur import coordinate_systems
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+tree_map = jax.tree_util.tree_map
+tree_leaves = jax.tree_util.tree_leaves
+
+Numeric = typing.Numeric
+Quantity = typing.Quantity
+_SOFTPLUS_INVERSE_1 = 0.5413248546129181
+
+# CoreRandomState is advanced by a RandomField, and .to_*_values(core_state)
+# produces the final (usable) random Array.
+CoreRandomState = typing.Pytree
+RandomnessState = typing.RandomnessState
+
+
+def _validate_randomness_state(state: RandomnessState) -> None:
+  """Validates that `state.core` is not `None`, raises an error otherwise."""
+  if state.core is None:
+    raise ValueError(
+        f'Got {state.core=} when value is expected. '
+        'Check how incoming randomness is initialized.'
+    )
+
+
+def make_positive_scalar(raw_parameter: typing.Array) -> jax.Array:
+  """Positive [batch] scalar values, maps 0 --> 1 using a softplus(...)."""
+  raw_parameter = jnp.asarray(raw_parameter)
+  return jax.nn.softplus(raw_parameter + _SOFTPLUS_INVERSE_1)
+
+
+# pylint: disable=logging-fstring-interpolation
+
+
+################################################################################
+# Single random fields that stand on their own.
+################################################################################
+
+
+class PreferredRepresentation(enum.Enum):
+  """The preferred (for computational reasons) representation of a field."""
+
+  NODAL = 'NODAL'
+  MODAL = 'MODAL'
+
+
+class RandomField(abc.ABC):
+  """Base class for random fields."""
+
+  def __init__(self, coords):
+    self.coords = coords
+
+  @property
+  @abc.abstractmethod
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    """The PreferredRepresentation for this field, or None if no preference."""
+
+  @abc.abstractmethod
+  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
+    """Sample the random field unconditionally."""
+
+  @abc.abstractmethod
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the core state of a random field."""
+
+  @abc.abstractmethod
+  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the modal rep. of the random field specified by this class."""
+
+  @abc.abstractmethod
+  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the nodal rep. of the random field specified by this class."""
+
+
+RandomnessModule = Callable[..., RandomField]
+
+
+_ADVANCE_SALT = zlib.crc32(b'advance')  # arbitrary uint32 value
+
+
+T = TypeVar('T', typing.PRNGKeyArray, None)
+
+
+def _prng_key_for_current_advance_step(
+    randomness: typing.RandomnessState,
+) -> typing.PRNGKeyArray | None:
+  """Get a PRNG Key suitable for randomness in the current advance step."""
+  if randomness.prng_key is None:
+    return None
+  salt = jnp.uint32(_ADVANCE_SALT) + jnp.uint32(randomness.prng_step)
+  return jax.random.fold_in(randomness.prng_key, salt)
+
+
+@gin.register
+class NoRandomField(RandomField):
+  """Module that disables randomness in a given module returning `None`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      prefer_nodal: bool = True,
+  ):
+    """Constructs a ZerosRandomField.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      prefer_nodal: Whether this field should prefer a nodal representation.
+    """
+    super().__init__(coords)
+    logging.info('[NGCM] Initializing NoRandomField')
+    del dt, physics_specs, aux_features, prefer_nodal  # unused.
+
+  @property
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    return None
+
+  def unconditional_sample(
+      self, rng: typing.PRNGKeyArray | None
+  ) -> RandomnessState:
+    """Returns a zeros initialized state."""
+    return RandomnessState(prng_key=rng, prng_step=0)  # pyrefly: ignore[unexpected-keyword]
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the state of a random gaussian field."""
+    return RandomnessState(
+        prng_key=state.prng_key, prng_step=state.prng_step + 1  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+    )
+
+  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    del core_state  # unused.
+    return None
+
+  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    del core_state  # unused.
+    return None
+
+
+@gin.register
+class ZerosRandomField(RandomField):
+  """Implements a constant random field identically equal to zero."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      prefer_nodal: bool = True,
+  ):
+    """Constructs a ZerosRandomField.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      prefer_nodal: Whether this field should prefer a nodal representation.
+    """
+    super().__init__(coords)
+    logging.info('[NGCM] Initializing ZerosRandomField')
+    del dt  # unused
+    del physics_specs  # unused.
+    del aux_features  # unused.
+    self._prefer_nodal = prefer_nodal
+
+  @property
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    if self._prefer_nodal:
+      return PreferredRepresentation.NODAL
+    else:
+      return PreferredRepresentation.MODAL
+
+  def unconditional_sample(
+      self, rng: typing.PRNGKeyArray | None
+  ) -> RandomnessState:
+    """Returns a zeros initialized state."""
+    if self._prefer_nodal:
+      core = jnp.zeros(self.coords.horizontal.nodal_shape)
+    else:
+      core = jnp.zeros(self.coords.horizontal.modal_shape)
+    return RandomnessState(
+        core=core,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=rng,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
+    )
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the state of a random gaussian field."""
+    _validate_randomness_state(state)
+    return RandomnessState(
+        core=jnp.zeros_like(state.core),  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
+        nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+    )
+
+  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the ready-for-use Zeros random field."""
+    return jnp.zeros(self.coords.horizontal.nodal_shape)
+
+  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the ready-for-use Zeros random field."""
+    return jnp.zeros(self.coords.horizontal.modal_shape)
+
+
+@gin.register
+class GaussianRandomField(RandomField):
+  """Implements gaussian random field with spatial and temporal correlations.
+
+  This type of random fields is used in SPPT (stochastic physics
+  parameterization tendencies) schemes, where each tendency due to physics
+  parameterizations are multiplicatively perturbed by the value of such field.
+
+  For implementation details see Appendix 8 in http://shortn/_56HCcQwmSS.
+
+  With x ∈ EarthSurface, this field U is initialized at t=0 with
+    U(0, x) = Σₖ Ψₖ(x) (1 - φ²)^(-0.5) σₖ γₖ σₖ ηₖ₀,
+  where Ψₖ is the kth spherical harmonic basis function, φ² is the one timestep
+  correlation, σₖ > 0 is a scaling factor, and ηₖ₀ are iid 1D unit Gaussians.
+
+  With `variance` an init kwarg,
+    E[U(0, x)] ≡ 0,
+    1 / (4πR²) ∫ Var(U(0, x))dx = variance,
+  regardless of coords (and the radius).
+
+  Further states are generated with the recursion
+    U(t + δ) = ϕ U(t) + σₖ ηₖₜ
+  This ensures that U is stationary.
+
+  In general,
+    Cov(U(t, x), U(t + δ, y)) = ϕᵟ Σₖ Ψₖ(x) Ψₖ(y) (γₖ)².
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      correlation_time: Union[jax.Array, Quantity, str] = gin.REQUIRED,
+      correlation_length: Union[jax.Array, Quantity, str] = gin.REQUIRED,
+      variance: Optional[Union[jax.Array, Quantity, str]] = gin.REQUIRED,
+      clip: float = 6.0,
+  ):
+    """Constructs a GaussianRandomField.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      correlation_time: timescale with units over which autoregressive process
+        decorrelates. Typical values in NWP range from hours to days.
+      correlation_length: lengthscale with units over which random field is
+        correlated. Typical values in NWP range from 500-2500 km.
+      variance: The average (over EarthSurface) variance of the random field If
+        None, this GRF always returns a zeros field and no RNGS are drawn.
+      clip: number of standard deviations at which to clip randomness to ensure
+        numerical stability.
+    """
+    del aux_features  # unused.
+    super().__init__(coords)
+    logging.info(
+        '[NGCM] Initializing GaussianRandomField (possibly via'
+        f' CenteredLognormalRandomField) with {variance=}, {correlation_time=},'
+        f' {correlation_length=}'
+    )
+
+    tau = maybe_nondimensionalize(correlation_time, physics_specs)
+    correlation_length = maybe_nondimensionalize(
+        correlation_length, physics_specs
+    )
+
+    # In sampling, phi appears as 1 - phi**2 = 1 - exp(-2 dt / tau)
+    self.one_minus_phi2 = -jnp.expm1(-2 * dt / tau)  # pyrefly: ignore[unsupported-operation]
+
+    self.phi = jnp.exp(-dt / tau)  # pyrefly: ignore[unsupported-operation]
+
+    self._variance = maybe_nondimensionalize(variance, physics_specs)  # σ²
+
+    # [Palmer] states correlation_length = sqrt(2κT) / R, therefore
+    self.kt = (correlation_length / self.coords.horizontal.radius) ** 2 / 2
+    self.clip = clip
+
+  @property
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    return PreferredRepresentation.MODAL
+
+  @property
+  def _surf_area(self) -> jax.Array:
+    """Surface area of sphere of radius self.coords.horizontal.radius."""
+    return 4 * jnp.pi * self.coords.horizontal.radius**2  # pytype: disable=bad-return-type  # jnp-type
+
+  def _sigma_array(self) -> jax.Array:
+    """Array of σₙ from Appendix 8 in [Palmer] http://shortn/_56HCcQwmSS."""
+    # n = [0, 1, ..., N]
+    n = self.coords.horizontal.modal_axes[1]  # total wavenumbers.
+
+    # Number of longitudinal wavenumbers at each total wavenumber n.
+    #  L = 2n + 1, except for the last entry.
+    n_longitudian_wavenumbers = self.coords.horizontal.mask.sum(axis=0)
+
+    # sigmas_unnormed[n] is proportional to the standard deviation for each
+    # longitudinal wavenumbers at each total wavenumber n.
+    sigmas_unnormed = jnp.exp(-0.5 * self.kt * n * (n + 1))
+
+    # The sum of unnormalized variance for all longitudinal wavenumbers at each
+    # total wavenumber.
+    sum_unnormed_vars = jnp.sum(n_longitudian_wavenumbers * sigmas_unnormed**2)
+
+    # This is analogous to F₀ from [Palmer].
+    # (normalization * sigmas_unnormed)² would sum to 1. The leading factor
+    #   self._integrated_grf_variance * (1 - self.phi ** 2)
+    # ensures that the AR(1) process has variance self._integrated_grf_variance.
+    # We do not include the extra fator of 2 in the denominator. I do not know
+    # why [Palmer] has this factor.
+    normalization = jnp.sqrt(
+        self._integrated_grf_variance()  # pyrefly: ignore[unsupported-operation]
+        * self.one_minus_phi2
+        / sum_unnormed_vars
+    )
+
+    # The factor of coords.horizontal.radius appears because our basis vectors
+    # have L2 norm = radius.  See http://screen/9FYVXZ5cMHoGDZk
+    return normalization * sigmas_unnormed / self.coords.horizontal.radius
+
+  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
+    """Returns a randomly initialized state for the autoregressive process."""
+    modal_shape = self.coords.horizontal.modal_shape
+    rng, next_rng = jax.random.split(rng)
+    if self.variance is None:
+      return RandomnessState(
+          core=jnp.zeros(modal_shape),  # pyrefly: ignore[unexpected-keyword]
+          nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
+          modal_value=jnp.zeros(modal_shape),  # pyrefly: ignore[unexpected-keyword]
+          prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
+          prng_step=0,  # pyrefly: ignore[unexpected-keyword]
+      )
+    sigmas = self._sigma_array()
+    weights = jnp.where(
+        self.coords.horizontal.mask,
+        jax.random.truncated_normal(rng, -self.clip, self.clip, modal_shape),
+        jnp.zeros(modal_shape),
+    )
+    core = self.one_minus_phi2 ** (-0.5) * sigmas * weights
+    return RandomnessState(
+        core=core,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=self.to_nodal_values(core),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=self.to_modal_values(core),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
+    )
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the CoreRandomState of a random gaussian field."""
+    _validate_randomness_state(state)
+    if self.variance is None:
+      return RandomnessState(
+          core=jnp.zeros_like(state.core),  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
+          nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
+          modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
+          prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
+          prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+      )
+    modal_shape = self.coords.horizontal.modal_shape
+    rng = _prng_key_for_current_advance_step(state)
+    eta = jax.random.truncated_normal(rng, -self.clip, self.clip, modal_shape)  # pyrefly: ignore[bad-argument-type]
+    next_core = state.core * self.phi + self._sigma_array() * jnp.where(  # pyrefly: ignore[unsupported-operation]
+        self.coords.horizontal.mask, eta, jnp.zeros(modal_shape)
+    )
+    return RandomnessState(
+        core=next_core,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=self.to_nodal_values(next_core),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=self.to_modal_values(next_core),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+    )
+
+  @property
+  def variance(self) -> Numeric | None:
+    """An estimate of pointwise (in nodal space) variance of this random field.
+
+    This random field is defined in spectral space, and has no precise
+    pointwise variance quantity. However, it does have a precise integrated
+    variance, which is used to define the field.
+
+    If we assume the field is stationary (with higher spectral
+    precision it is near stationary), then the average of this quantity is a
+    good pointwise estimate. So define
+      σ² := (1 / (4πR²)) ∫ Var(U(0, x))dx
+          = (1 / (4πR²)) integrated_grf_variance
+
+    Therefore the init parameter `variance` can be used to define
+      `_integrated_grf_variance := variance * surf_area`
+    and then `_integrated_grf_variance` is used to define this field. The result
+    is a field with pointwise variance close to the init kwarg `variance`.
+
+    Returns:
+      Numeric estimate of pointwise variance.
+    """
+    return self._variance
+
+  def _integrated_grf_variance(self) -> Numeric | None:
+    """Integral of the GRF's variance over the earth's surface."""
+    if self.variance is None:
+      return self.variance
+    return self.variance * self._surf_area
+
+  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the ready-for-use Gaussian random field."""
+    return core_state
+
+  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Returns the ready-for-use Gaussian random field."""
+    return self.coords.horizontal.to_nodal(core_state)
+
+
+@gin.register
+class GaussianRandomFieldModule(GaussianRandomField, hk.Module):
+  """Module wrapper of GaussianRandomField with trainable parameters."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      initial_correlation_time: Union[Quantity, str] = gin.REQUIRED,
+      initial_correlation_length: Union[Quantity, str] = gin.REQUIRED,
+      initial_variance: Optional[Union[Quantity, str]] = gin.REQUIRED,
+      variance_bound: Optional[Union[Quantity, str]] = gin.REQUIRED,
+      tune_variance: bool = True,
+      clip: float = 6.0,
+      name: Optional[str] = None,
+  ):
+    """Constructs a GaussianRandomFieldModule.
+
+    Stochastic parameters are initialized at provided `initial_*` values.
+    This hk.Module can then be used to tune values.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      initial_correlation_time: timescale with units over which autoregressive
+        process decorrelates. Typical values in NWP range from hours to days.
+      initial_correlation_length: lengthscale with units over which random field
+        is correlated. Typical values in NWP range from 500-2500 km.
+      initial_variance: The average (over EarthSurface) variance of the random
+        field. If None, this GRF always returns a zeros field and no RNGS will
+        be drawn
+      variance_bound: If provided, an upper bound on tuned variance values.
+      tune_variance: Whether variance should be a tunable hk.parameter, or fixed
+      clip: number of standard deviations at which to clip randomness to ensure
+        numerical stability.
+      name: Something no one cares about and we just use None.
+    """
+    # You must call hk.Module.__init__ before initializing this class.
+    hk.Module.__init__(self, name=name)
+
+    correlation_time_raw = hk.get_parameter(
+        'correlation_time_raw', shape=(), init=hk.initializers.Constant(0.0)
+    )
+    correlation_length_raw = hk.get_parameter(
+        'correlation_length_raw', shape=(), init=hk.initializers.Constant(0.0)
+    )
+
+    if tune_variance:
+      variance_raw = hk.get_parameter(
+          'variance_raw', shape=(), init=hk.initializers.Constant(0.0)
+      )
+    else:
+      variance_raw = 0.0
+
+    initial_variance = maybe_nondimensionalize(initial_variance, physics_specs)
+    _assert_positive_or_none(initial_variance, 'initial_variance')
+
+    if initial_variance is None:
+      variance = None
+    elif variance_bound in {None, 'None'}:  # Allow strings for gin.
+      variance = convert_hk_param_to_positive_scalar(
+          variance_raw, initial_variance  # pyrefly: ignore[bad-argument-type]
+      )
+    else:
+      variance_bound = maybe_nondimensionalize(variance_bound, physics_specs)
+      _assert_positive_or_none(variance_bound, 'variance_bound')
+      _assert_positive_or_none(
+          variance_bound - initial_variance, 'variance_bound - initial_variance'  # pyrefly: ignore[unsupported-operation]
+      )
+      variance = convert_hk_param_to_bounded_scalar(
+          variance_raw,  # pyrefly: ignore[bad-argument-type]
+          initial_variance,
+          low=0.0,
+          high=variance_bound,  # pyrefly: ignore[bad-argument-type]
+      )
+
+    # We call GaussianRandomFieldModule.__init__ rather than super().__init__
+    # since we don't want to call hk.Module.__init__ twice... although doing
+    # that didn't hurt anything.
+    GaussianRandomField.__init__(
+        self,
+        coords=coords,
+        dt=dt,
+        physics_specs=physics_specs,
+        aux_features=aux_features,
+        correlation_time=convert_hk_param_to_positive_scalar(
+            correlation_time_raw,
+            maybe_nondimensionalize(initial_correlation_time, physics_specs),  # pyrefly: ignore[bad-argument-type]
+        ),
+        correlation_length=convert_hk_param_to_positive_scalar(
+            correlation_length_raw,
+            maybe_nondimensionalize(initial_correlation_length, physics_specs),  # pyrefly: ignore[bad-argument-type]
+        ),
+        variance=variance,
+        clip=clip,
+    )
+
+
+################################################################################
+# Single random fields that are derived from "stand on their own" fields.
+################################################################################
+
+
+@gin.register
+class CenteredLognormalRandomField(GaussianRandomField):
+  """A lognormal random field shifted to have mean zero."""
+
+  @property
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    return PreferredRepresentation.NODAL
+
+  def _integrated_grf_variance(self) -> jax.Array | None:
+    """Integrated variance of the associated GRF (not this Lognormal field)."""
+    if self.variance is None:
+      return None
+    # If Z ~ Normal(μ, σ²), then X ~ exp(Z) has
+    #  variance = (exp(σ²) - 1) exp(2μ + σ²).
+    # We have centered this field, which involved setting μ = -σ² / 2.
+    # => variance = exp(σ²) - 1,
+    # and thus
+    #  σ² = log(1 + variance)
+    return jnp.log1p(self.variance) * self._surf_area
+
+  def to_nodal_values(self, core_state: CoreRandomState) -> jax.Array:
+    """Returns the ready-for-use Lognormal random field."""
+    if self.variance is None:
+      grf_variance = 0.0
+    else:
+      grf_variance = self._integrated_grf_variance() / self._surf_area  # pyrefly: ignore[unsupported-operation]
+    # If Z ~ Normal(μ, σ²), then X ~ exp(Z) has mean exp(μ + σ²/2).
+    # To ensure E[X] = 1, we must set μ = -σ²/2.
+    x = self.coords.horizontal.to_nodal(core_state)  # ~ Normal(0, σ²)
+    return jnp.expm1(x - grf_variance / 2)  # ~ Exp(Normal(-σ²/2, σ²)) - 1
+
+  def to_modal_values(self, core_state: CoreRandomState) -> jax.Array:
+    """Returns the ready-for-use Lognormal random field."""
+    return self.coords.horizontal.to_modal(self.to_nodal_values(core_state))
+
+
+@gin.register
+class CenteredLognormalRandomFieldModule(
+    CenteredLognormalRandomField, GaussianRandomFieldModule
+):
+  """A lognormal random hk.Module field shifted to have mean zero."""
+
+
+################################################################################
+# Fields made from many different fields.
+################################################################################
+
+
+@gin.register
+class BatchGaussianRandomFieldModule(hk.Module):
+  """Batch of independent GaussianRandomFieldModules.
+
+  These GRFs are meant to be fed into a neural network as generic "signals".
+
+  The state arrays have leading batch dim indexing independent GRFs.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      variances: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      field_subset: Optional[Sequence[int]] = None,
+      n_fixed_fields: Optional[int] = None,
+      clip: float = 6.0,
+      name: Optional[str] = None,
+  ):
+    """Constructs a BatchGaussianRandomFieldModule.
+
+    Correlation scales are initialized to `initial_*` args and will be tuned
+    by Haiku optimizers. Variance will be fixed.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      initial_correlation_times: timescales with units over which autoregressive
+        process decorrelates. Typical values in NWP range from hours to days.
+      initial_correlation_lengths: lengthscale with units over which random
+        field is correlated. Typical values in NWP range from 500-2500 km.
+      variances: The average (over EarthSurface) variance of the random field.
+        These are fixed arrays (not tunable hk.parameters).
+      field_subset: Optional nonempty subset of indices into initial parameters.
+        Specifies which fields to construct. If None, use all fields. E.g.,
+        field_subset=[0, 5] means form 3 GRFs from the 0th and 5th parameter
+        values.
+      n_fixed_fields: Number of fields that use fixed parameters. These will
+        be fixed at the trailing `n_fixed_fields` initial correlations. The
+        total number of fields is unchanged, since these fixed fields replace
+        learnable fields.
+      clip: number of standard deviations at which to clip randomness to ensure
+        numerical stability.
+      name: Name to show in xprof.
+    """
+    ## You must call hk.Module.__init__ before initializing this class.
+    hk.Module.__init__(self, name=name)
+
+    lengths = [
+        len(initial_correlation_times),
+        len(initial_correlation_lengths),
+        len(variances),
+    ]
+    if len(set(lengths)) != 1:
+      raise ValueError(f'Argument lengths differed: {lengths=}')
+    n_fixed_fields = n_fixed_fields or 0
+
+    # Get subset of args using `field_subset`
+    if field_subset is not None:
+      if not field_subset:
+        raise ValueError(
+            '`field_subset` must be `None` or non-empty sequence. Found'
+            f' {field_subset=}'
+        )
+      get_subset = lambda seq: [seq[i] for i in field_subset]
+      initial_correlation_lengths = get_subset(initial_correlation_lengths)
+      initial_correlation_times = get_subset(initial_correlation_times)
+      variances = get_subset(variances)
+
+    logging.info(
+        '[NGCM] Initializing BatchGaussianRandomFieldModule with'
+        f' {initial_correlation_times=}, and {initial_correlation_lengths=},'
+        f' and {variances=}'
+    )
+
+    # Get Haiku parameters.
+    self._n_fields = len(variances)
+    self._variances = jnp.array(
+        [nondimensionalize(v, physics_specs) for v in variances]
+    )
+
+    initial_correlation_lengths = jnp.array([  # pyrefly: ignore[bad-assignment]
+        nondimensionalize(l, physics_specs) for l in initial_correlation_lengths
+    ])
+    correlation_lengths_raw = hk.get_parameter(
+        'correlation_lengths_raw',
+        shape=(self.n_fields - n_fixed_fields,),
+        init=hk.initializers.Constant(0.0),
+    )
+    if n_fixed_fields:
+      correlation_lengths_raw = jnp.concatenate([
+          correlation_lengths_raw, jnp.zeros([n_fixed_fields])])
+    self._correlation_lengths = convert_hk_param_to_positive_scalar(
+        correlation_lengths_raw, initial_correlation_lengths  # pyrefly: ignore[bad-argument-type]
+    )
+
+    initial_correlation_times = jnp.array(  # pyrefly: ignore[bad-assignment]
+        [nondimensionalize(t, physics_specs) for t in initial_correlation_times]
+    )
+    correlation_times_raw = hk.get_parameter(
+        'correlation_times_raw',
+        shape=(self.n_fields - n_fixed_fields,),
+        init=hk.initializers.Constant(0.0),
+    )
+    if n_fixed_fields:
+      correlation_times_raw = jnp.concatenate([
+          correlation_times_raw, jnp.zeros([n_fixed_fields])])
+    self._correlation_times = convert_hk_param_to_positive_scalar(
+        correlation_times_raw, initial_correlation_times  # pyrefly: ignore[bad-argument-type]
+    )
+
+    def make_rf(correlation_time, correlation_length, variance):
+      return GaussianRandomField(
+          coords=coords,
+          dt=dt,
+          physics_specs=physics_specs,
+          aux_features=aux_features,
+          correlation_time=correlation_time,
+          correlation_length=correlation_length,
+          variance=variance,
+          clip=clip,
+      )
+
+    self._make_rf = make_rf
+
+  @property
+  def n_fields(self) -> int:
+    return self._n_fields
+
+  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
+    """Sample the batch GRFs unconditionally."""
+    logging.info(
+        '[NGCM] Calling BatchGaussianRandomFieldModule.unconditional_sample'
+    )
+
+    def _unconditional_sample_one_rf(
+        key, correlation_time, correlation_length, variance
+    ):
+      rf = self._make_rf(correlation_time, correlation_length, variance)
+      return rf.unconditional_sample(key)
+
+    rngs = jax.random.split(rng, self.n_fields + 1)
+    rngs, next_rng = rngs[:-1], rngs[-1]
+    sample = jax.vmap(_unconditional_sample_one_rf)(
+        rngs,
+        self._correlation_times,
+        self._correlation_lengths,
+        self._variances,
+    )
+    # We have RNG keys and steps associated with each field from vmap, but
+    # RandomnessState should only have a single (scalar) RNG key/step.
+    return dataclasses.replace(sample, prng_key=next_rng, prng_step=0)
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the state of the batch of GRFs."""
+    logging.info('[NGCM] Calling BatchGaussianRandomFieldModule.advance')
+
+    def _advance_one_rf(state, correlation_time, correlation_length, variance):
+      rf = self._make_rf(correlation_time, correlation_length, variance)
+      return rf.advance(state)
+
+    rng = _prng_key_for_current_advance_step(state)
+    rngs = jax.random.split(rng, self.n_fields)  # pyrefly: ignore[bad-argument-type]
+    steps = jnp.ones(self.n_fields, int) * state.prng_step  # pyrefly: ignore[unsupported-operation]
+    advanced = jax.vmap(_advance_one_rf)(
+        dataclasses.replace(state, prng_key=rngs, prng_step=steps),  # pyrefly: ignore[bad-specialization]
+        self._correlation_times,
+        self._correlation_lengths,
+        self._variances,
+    )
+    return dataclasses.replace(
+        advanced, prng_key=state.prng_key, prng_step=state.prng_step + 1  # pyrefly: ignore[unsupported-operation]
+    )
+
+
+@gin.register
+class DictOfGaussianRandomFieldModules(hk.Module):
+  """Dictionary of independent GaussianRandomFieldModules.
+
+  These GRFs are meant to be fed into a neural network as generic "signals".
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      variances: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      field_names: Optional[Sequence[str]] = None,
+      field_subset: Optional[Sequence[int]] = None,
+      clip: float = 6.0,
+      name: Optional[str] = None,
+  ):
+    """Constructs a DictOfGaussianRandomFieldModules.
+
+    Correlation scales are initialized to `initial_*` args and will be tuned
+    by Haiku optimizers. Variance will be fixed.
+
+    Args:
+      coords: horizontal and vertical grid data.
+      dt: nondimensionalized model time step.
+      physics_specs: physical constants and definition of custom units.
+      aux_features: additional static data.
+      initial_correlation_times: timescales with units over which autoregressive
+        process decorrelates. Typical values in NWP range from hours to days.
+      initial_correlation_lengths: lengthscale with units over which random
+        field is correlated. Typical values in NWP range from 500-2500 km.
+      variances: The average (over EarthSurface) variance of the random field.
+        These are fixed arrays (not tunable hk.parameters).
+      field_names: Optional names to give the fields. If None, the fields are
+        named like "GRF0", "GRF1",...
+      field_subset: Optional nonempty subset of indices into initial parameters.
+        Specifies which fields to construct. If None, use all fields. E.g.,
+        field_subset=[0, 5] means form 3 GRFs from the 0th and 5th parameter
+        values.
+      clip: number of standard deviations at which to clip randomness to ensure
+        numerical stability.
+      name: Name to show in xprof.
+    """
+    ## You must call hk.Module.__init__ before initializing this class.
+    hk.Module.__init__(self, name=name)
+    logging.info(
+        '[NGCM] Initializing DictOfGaussianRandomFieldModules with'
+        f' {initial_correlation_times=}, and {initial_correlation_lengths=},'
+        f' and {variances=}'
+    )
+
+    field_names = field_names or [
+        f'GRF{i}' for i in range(len(initial_correlation_times))
+    ]
+
+    lengths = [
+        len(initial_correlation_times),
+        len(initial_correlation_lengths),
+        len(variances),
+        len(field_names),
+    ]
+    if len(set(lengths)) != 1:
+      raise ValueError(f'Argument lengths differed: {lengths=}')
+
+    if field_subset is not None:
+      if not field_subset:
+        raise ValueError(
+            '`field_subset` must be `None` or non-empty sequence. Found'
+            f' {field_subset=}'
+        )
+      subset = lambda seq: [seq[i] for i in field_subset]
+      field_names = subset(field_names)
+      initial_correlation_lengths = subset(initial_correlation_lengths)
+      initial_correlation_times = subset(initial_correlation_times)
+      variances = subset(variances)
+
+    self._field_names = tuple(field_names)
+
+    self._random_fields = {}
+    for tau, lam, var, field_name in zip(
+        initial_correlation_times,
+        initial_correlation_lengths,
+        variances,
+        self.field_names,
+        strict=True,
+    ):
+      self._random_fields[field_name] = GaussianRandomFieldModule(
+          coords,
+          dt,
+          physics_specs,
+          aux_features,
+          initial_correlation_time=tau,
+          initial_correlation_length=lam,
+          initial_variance=var,
+          tune_variance=False,
+          variance_bound=None,
+          clip=clip,
+          name=field_name,
+      )
+
+  @property
+  def n_fields(self) -> int:
+    return len(self._random_fields)
+
+  @property
+  def field_names(self) -> tuple[str, ...]:
+    return self._field_names
+
+  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
+    """Sample the random field unconditionally."""
+    core = {}
+    nodal_values = {}
+    modal_values = {}
+    *rngs, next_rng = jax.random.split(rng, self.n_fields + 1)
+    for (name, rf), sample_key in zip(self._random_fields.items(), rngs):
+      rvs = rf.unconditional_sample(sample_key)
+      core[name] = rvs.core
+      nodal_values[name] = rvs.nodal_value
+      modal_values[name] = rvs.modal_value
+    return RandomnessState(
+        core=core,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=nodal_values,  # pyrefly: ignore[unexpected-keyword]
+        modal_value=modal_values,  # pyrefly: ignore[unexpected-keyword]
+        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
+    )
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the core state of a random field."""
+    core = {}
+    nodal_values = {}
+    modal_values = {}
+    rng = _prng_key_for_current_advance_step(state)
+    rngs = jax.random.split(rng, self.n_fields)  # pyrefly: ignore[bad-argument-type]
+    for (name, rf), sample_key in zip(self._random_fields.items(), rngs):
+      # rvs is a RandomnessState.
+      rvs = rf.advance(
+          RandomnessState(state.core[name], prng_key=sample_key, prng_step=0)  # pyrefly: ignore[bad-argument-count, unexpected-keyword, unsupported-operation]
+      )
+      core[name] = rvs.core
+      nodal_values[name] = rvs.nodal_value
+      modal_values[name] = rvs.modal_value
+    return RandomnessState(
+        core=core,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=nodal_values,  # pyrefly: ignore[unexpected-keyword]
+        modal_value=modal_values,  # pyrefly: ignore[unexpected-keyword]
+        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+    )
+
+
+class SumOfRandomFields(RandomField):
+  """RandomField that is the sum of multiple fields."""
+
+  def __init__(self, random_fields: Sequence[RandomField]):
+    self._random_fields = list(random_fields)  # Shallow copy
+    coords = self._random_fields[0].coords
+    if any(rf.coords != coords for rf in self._random_fields):
+      raise ValueError(f'All fields must have the same coords. Found {coords=}')
+    super().__init__(coords)
+
+  @property
+  def preferred_representation(self) -> PreferredRepresentation | None:
+    n_nodal = sum(
+        rf.preferred_representation == PreferredRepresentation.NODAL
+        for rf in self._random_fields
+    )
+    n_modal = sum(
+        rf.preferred_representation == PreferredRepresentation.MODAL
+        for rf in self._random_fields
+    )
+    if n_nodal > n_modal:
+      return PreferredRepresentation.NODAL
+    elif n_nodal < n_modal:
+      return PreferredRepresentation.MODAL
+    return None
+
+  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
+    """Sample the random field unconditionally."""
+    rvs = []
+    *rngs, next_rng = jax.random.split(rng, len(self._random_fields) + 1)
+    for rf, sample_key in zip(self._random_fields, rngs, strict=True):
+      rvs.append(rf.unconditional_sample(sample_key).core)
+    return RandomnessState(
+        core=rvs,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=self.to_nodal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=self.to_modal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
+    )
+
+  def advance(self, state: RandomnessState) -> RandomnessState:
+    """Updates the core state of a random field."""
+    rvs = []
+    rng = _prng_key_for_current_advance_step(state)
+    rngs = jax.random.split(rng, len(self._random_fields))  # pyrefly: ignore[bad-argument-type]
+    for rf, s, k in zip(
+        self._random_fields, state.core, rngs, strict=True  # pyrefly: ignore[bad-argument-type]
+    ):
+      rs = RandomnessState(s, prng_key=k, prng_step=state.prng_step)  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
+      rvs.append(rf.advance(rs).core)
+    return RandomnessState(
+        core=rvs,  # pyrefly: ignore[unexpected-keyword]
+        nodal_value=self.to_nodal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
+        modal_value=self.to_modal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
+        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
+        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
+    )
+
+  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Finishes `core_state` by summing components."""
+    modal_sum = 0.0
+    nodal_sum = 0.0
+    for rf, s in zip(self._random_fields, core_state, strict=True):
+      if rf.preferred_representation == PreferredRepresentation.NODAL:
+        nodal_sum += rf.to_nodal_values(s)  # pyrefly: ignore[unsupported-operation]
+      elif rf.preferred_representation in [PreferredRepresentation.MODAL, None]:
+        modal_sum += rf.to_modal_values(s)  # pyrefly: ignore[unsupported-operation]
+    return modal_sum + self.coords.horizontal.to_modal(nodal_sum)
+
+  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
+    """Finishes `core_state` by summing components."""
+    modal_sum = 0.0
+    nodal_sum = 0.0
+    for rf, s in zip(self._random_fields, core_state, strict=True):
+      if rf.preferred_representation == PreferredRepresentation.MODAL:
+        modal_sum += rf.to_modal_values(s)  # pyrefly: ignore[unsupported-operation]
+      elif rf.preferred_representation in [PreferredRepresentation.NODAL, None]:
+        nodal_sum += rf.to_nodal_values(s)  # pyrefly: ignore[unsupported-operation]
+    return nodal_sum + self.coords.horizontal.to_nodal(modal_sum)
+
+
+class SumOfGaussianLikeRandomFields(SumOfRandomFields, abc.ABC):
+  """Base class for sum of independent Gaussian-like random fields."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      correlation_times: Sequence[
+          Union[jax.Array, Quantity, str]
+      ] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      correlation_lengths: Sequence[
+          Union[jax.Array, Quantity, str]
+      ] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      variances: Sequence[Union[jax.Array, Quantity, str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      clip: float = 6.0,
+  ):
+    """Constructs a SumOfGaussianLikeRandomFields."""
+    n_fields = len(correlation_times)
+    variances = variances or [None] * n_fields
+    random_fields = []
+    logging.info(
+        '[NGCM] Initializing SumOfGaussianLikeRandomFields with '
+        f'{variances=}, {correlation_times=}, {correlation_lengths=}'
+    )
+    for tau, lam, var in zip(
+        correlation_times, correlation_lengths, variances, strict=True
+    ):
+      random_fields.append(
+          self.get_cls_constructor()(
+              coords,
+              dt,
+              physics_specs,
+              aux_features,
+              correlation_time=tau,
+              correlation_length=lam,
+              variance=var,
+              clip=clip,
+          )
+      )
+
+    super().__init__(random_fields)
+
+  @abc.abstractmethod
+  def get_cls_constructor(self) -> type[GaussianRandomField]:
+    """Gets class constructor that is initialized with Gaussian-like kwargs."""
+
+
+class SumOfGaussianLikeRandomFieldsModule(
+    SumOfRandomFields, hk.Module, abc.ABC
+):
+  """Base class for sums of independent Gaussian-like RandomFieldModules."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: Any,
+      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      initial_variances: Optional[Sequence[Quantity | str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      variance_bounds: Optional[Sequence[Quantity | str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      clip: float = 6.0,
+      name: Optional[str] = None,
+  ):
+    """Constructs a SumOfGaussianLikeRandomFieldsModule."""
+    # You must call hk.Module.__init__ before initializing this class.
+    hk.Module.__init__(self, name=name)
+
+    n_fields = len(initial_correlation_times)
+    initial_variances = initial_variances or [None] * n_fields
+    variance_bounds = variance_bounds or [None] * n_fields
+    random_fields = []
+    for tau, lam, var, bound in zip(
+        initial_correlation_times,
+        initial_correlation_lengths,
+        initial_variances,
+        variance_bounds,
+        strict=True,
+    ):
+      random_fields.append(
+          self.get_cls_constructor()(
+              coords,
+              dt,
+              physics_specs,
+              aux_features,
+              initial_correlation_time=tau,
+              initial_correlation_length=lam,
+              initial_variance=var,
+              variance_bound=bound,
+              clip=clip,
+              name=name,
+          )
+      )
+    # We call SumOfRandomFields.__init__ rather than super().__init__
+    # since we don't want to call hk.Module.__init__ twice... although doing
+    # that didn't hurt anything.
+    SumOfRandomFields.__init__(self, random_fields)
+
+  @abc.abstractmethod
+  def get_cls_constructor(self) -> type[GaussianRandomFieldModule]:
+    """Gets class constructor that is initialized with Gaussian-like kwargs."""
+
+
+@gin.register
+class SumOfGaussianRandomFields(SumOfGaussianLikeRandomFields):
+
+  def get_cls_constructor(self) -> type[GaussianRandomField]:
+    return GaussianRandomField
+
+
+@gin.register
+class SumOfGaussianRandomFieldsModule(SumOfGaussianLikeRandomFieldsModule):
+  """A sum of independent GaussianRandomFieldModules."""
+
+  def get_cls_constructor(self) -> type[GaussianRandomFieldModule]:
+    return GaussianRandomFieldModule
+
+
+@gin.register
+class SumOfCenteredLognormalRandomFields(SumOfGaussianLikeRandomFields):
+
+  def get_cls_constructor(self) -> type[CenteredLognormalRandomField]:
+    return CenteredLognormalRandomField
+
+
+@gin.register
+class SumOfCenteredLognormalRandomFieldsModule(
+    SumOfGaussianLikeRandomFieldsModule
+):
+  """A sum of independent CenteredLognormalRandomFieldModules."""
+
+  def get_cls_constructor(self) -> type[CenteredLognormalRandomFieldModule]:
+    return CenteredLognormalRandomFieldModule
+
+
+################################################################################
+# Helper functions for creating fields.
+################################################################################
+
+
+def convert_hk_param_to_positive_scalar(
+    param: jax.Array,
+    initial_value: Numeric,
+) -> jax.Array:
+  """Converts [batch] scalar parameter to scalar value using Softplus."""
+  return initial_value * make_positive_scalar(param)  # pyrefly: ignore[bad-return]
+
+
+def _sigmoid(low: Numeric, high: Numeric, x: jax.Array) -> jax.Array:
+  """Numerically stable sigmoid, adapted from tfp.bijectors.Sigmoid."""
+  diff = high - low
+  left = low + diff * jax.nn.sigmoid(x)
+  right = high - diff * jax.nn.sigmoid(-x)
+  return jnp.where(x < 0, left, right)
+
+
+def _inv_sigmoid(low: Numeric, high: Numeric, x: jax.Array) -> jax.Array:
+  """Inverse sigmoid, adapted from tfp.bijectors.Sigmoid."""
+  return jnp.log(x - low) - jnp.log(high - x)
+
+
+def convert_hk_param_to_bounded_scalar(
+    param: jax.Array,
+    initial_value: Numeric,
+    low: Numeric,
+    high: Numeric,
+) -> jax.Array:
+  """Converts a [batch] scalar parameter to scalar value using Sigmoid."""
+  offset = _inv_sigmoid(low, high, initial_value)  # pyrefly: ignore[bad-argument-type]
+  return _sigmoid(low, high, offset + param)
+
+
+def nondimensionalize(
+    x: Union[typing.Numeric, Quantity, str],
+    physics_specs: Any,
+) -> typing.Numeric:
+  if isinstance(x, (Quantity, str)):  # pyrefly: ignore[invalid-argument]
+    return physics_specs.nondimensionalize(Quantity(x))
+  else:
+    return x
+
+
+def maybe_nondimensionalize(
+    x: Optional[Union[typing.Numeric, Quantity, str]],
+    physics_specs: Any,
+) -> None | typing.Numeric:
+  """Calls nondimensionalize on Quantity or str, otherwise passthrough."""
+  if x == 'None':  # Allow strings for gin
+    return None
+  return nondimensionalize(x, physics_specs)
+
+
+def _assert_positive_or_none(x: typing.Numeric | None, name: str) -> None:
+  if x is None:
+    return
+  if x <= 0:
+    raise ValueError(f'{name}={x} but should have been positive or None')

model/legacy/towers.py

@@ -0,0 +1,206 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Basic neural network towers for whirl/gcm codebase.
+
+A tower is a neural network that operates identically over the last two
+dimensions, i.e. (longitude, latitude).
+"""
+from collections import abc
+from typing import Callable, Optional, Tuple
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import layers
+
+Array = typing.Array
+TowerFactory = typing.TowerFactory
+LayerFactory = typing.LayerFactory
+MLP = gin.external_configurable(hk.nets.MLP)
+
+
+@gin.register(denylist=['output_size'])
+class ColumnTower(hk.Module):
+  """Column tower module parameterized by column_net_factory."""
+
+  def __init__(
+      self,
+      output_size: int,
+      column_net_factory: LayerFactory = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      checkpoint_tower: bool = False,
+      name: Optional[str] = None,
+  ):
+    """Tower that maps a column_net over two spatial dimensions."""
+    super().__init__(name=name)
+    column_net = column_net_factory(output_size)
+    vmap_last = lambda fn: hk.vmap(fn, in_axes=-1, out_axes=-1, split_rng=False)
+    column_tower = vmap_last(vmap_last(column_net))
+    if checkpoint_tower:
+      column_tower = hk.remat(column_tower)
+    self.column_tower = column_tower
+
+  def __call__(self, inputs: Array) -> Array:
+    """Applies Column tower to inputs."""
+    return self.column_tower(inputs)
+
+
+@gin.register(denylist=['output_size'])
+class ColumnTransformerTower(ColumnTower):
+  """Same as ColumnTower, but passes additional transformer inputs."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+
+  def __call__(
+      self,
+      inputs: Array,
+      latents: Optional[Array] = None,
+      positional_encoding: Optional[Array] = None,
+  ) -> Array:
+    """Applies Column tower to inputs."""
+    return self.column_tower(inputs, latents, positional_encoding)
+
+
+@gin.register(denylist=['output_size'])
+class VerticalConvTower(hk.Module):
+  """Tower that stacks up layers of Conv1D.
+
+  input shape: [in_channel, level, lon, lat],
+  output shape: [output_size, level, lon, lat].
+  """
+
+  def __init__(
+      self,
+      output_size: int,  # The number of channels in the last layer
+      channels: abc.Sequence[int] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      kernel_shape: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      with_bias: bool = True,
+      activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
+      activate_final: bool = False,
+      checkpoint_tower: bool = False,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.activation = activation
+    self.output_size = output_size
+    self.activate_final = activate_final
+    self.checkpoint_tower = checkpoint_tower
+
+    self.layers = []
+    channels = list(channels) + [self.output_size]
+    for channels_i in channels:
+      self.layers.append(layers.ConvLevel(
+          output_channels=channels_i,
+          kernel_shape=kernel_shape,
+          with_bias=with_bias))
+
+  def net(self, inputs: Array) -> Array:
+    out = inputs
+    num_layers = len(self.layers)
+    for i, layer in enumerate(self.layers):
+      out = layer(out)
+      if i < (num_layers - 1) or self.activate_final:
+        out = self.activation(out)
+    return out
+
+  def __call__(self, inputs: Array) -> Array:
+    vmap_last = lambda fn: hk.vmap(fn, in_axes=-1, out_axes=-1, split_rng=False)
+    tower_fn = vmap_last(vmap_last(self.net))
+    if self.checkpoint_tower:
+      tower_fn = hk.remat(tower_fn)
+    return tower_fn(inputs)
+
+
+@gin.register(denylist=['output_size'])
+class Conv2DTower(hk.Module):
+  """Two dimensional ConvNet tower module."""
+
+  def __init__(
+      self,
+      output_size: int,
+      num_hidden_units: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      num_hidden_layers: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      kernel_shape: Tuple[int, int] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      with_bias: bool = True,
+      activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
+      activate_final: bool = False,
+      name: Optional[str] = None,
+  ):
+    """Tower that stacks up layers of ConvLonLat."""
+    super().__init__(name=name)
+    self.activation = activation
+    self.activate_final = activate_final
+
+    output_sizes = [num_hidden_units] * num_hidden_layers + [output_size]
+    self.layers = []
+    for output_size in output_sizes:
+      self.layers.append(layers.ConvLonLat(
+          output_size=output_size,
+          kernel_shape=kernel_shape,
+          with_bias=with_bias))
+
+  def __call__(self, inputs: Array) -> Array:
+    """Applies ConvNet tower to inputs."""
+    num_layers = len(self.layers)
+    out = inputs
+    for i, layer in enumerate(self.layers):
+      out = layer(out)
+      if i < (num_layers - 1) or self.activate_final:
+        out = self.activation(out)
+    return out
+
+
+@gin.register(denylist=['output_size'])
+class EpdTower(hk.Module):
+  """EPD tower module parameterized by encode/process/decode factories."""
+
+  def __init__(
+      self,
+      output_size: int,
+      latent_size: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      num_process_blocks: int = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      encode_tower_factory: TowerFactory = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      process_tower_factory: TowerFactory = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      decode_tower_factory: TowerFactory = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
+      post_encode_activation: Optional[Callable[[Array], Array]] = None,
+      pre_decode_activation: Optional[Callable[[Array], Array]] = None,
+      final_activation: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.output_size = output_size
+    self.latent_size = latent_size
+    self.num_process_blocks = num_process_blocks
+    self.encode_tower_factory = encode_tower_factory
+    self.process_tower_factory = process_tower_factory
+    self.decode_tower_factory = decode_tower_factory
+    self.post_encode_activation = post_encode_activation
+    self.pre_decode_activation = pre_decode_activation
+    self.final_activation = final_activation
+
+  def __call__(self, inputs: Array) -> Array:
+    """Applies EPD tower to inputs."""
+    encoded = self.encode_tower_factory(self.latent_size)(inputs)
+    if self.post_encode_activation is not None:
+      encoded = self.post_encode_activation(encoded)
+    current = encoded
+    for _ in range(self.num_process_blocks):
+      current = current + self.process_tower_factory(self.latent_size)(current)
+    if self.pre_decode_activation is not None:
+      current = self.pre_decode_activation(current)
+    out = self.decode_tower_factory(self.output_size)(current)
+    if self.final_activation is not None:
+      return self.final_activation(out)
+    return out

model/legacy/transforms.py

@@ -0,0 +1,738 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformation modules that convert or pre/post process data structures."""
+
+import dataclasses
+import functools
+import re
+from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple
+from dinosaur import coordinate_systems
+from dinosaur import pytree_utils
+from dinosaur import sigma_coordinates
+from dinosaur import typing
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from model.legacy import filters
+import numpy as np
+
+
+KeyWithCosLatFactor = typing.KeyWithCosLatFactor
+TransformModule = typing.TransformModule
+
+
+@gin.register
+class EmptyTransform(hk.Module):
+  """Transform returns an empty dict."""
+
+  def __init__(self, *args, name: Optional[str] = None):
+    del args  # unused.
+    super().__init__(name=name)
+
+  def __call__(self, inputs) -> typing.Pytree:
+    return {}
+
+
+@gin.register
+class IdentityTransform(hk.Module):
+  """Transform does not modify inputs."""
+
+  def __init__(self, *args, name: Optional[str] = None, **kwargs):
+    del args, kwargs  # unused.
+    super().__init__(name=name)
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return inputs
+
+
+@gin.register
+class ShiftAndNormalize(hk.Module):
+  """Transforms inputs by shifting and normalizing values by `shifts/scales`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      shifts: typing.Pytree,
+      scales: typing.Pytree,
+      features_to_exclude: Sequence[str] = tuple(),
+      global_scale: Optional[float] = None,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.shifts = shifts
+    if global_scale is not None:
+      scales = jax.tree_util.tree_map(lambda x: x * global_scale, scales)
+    self.scales = scales
+
+  def __call__(self, inputs):
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    shifts = pytree_utils.replace_with_matching_or_default(
+        inputs, self.shifts, default=None, check_used_all_replace_keys=False)
+    scales = pytree_utils.replace_with_matching_or_default(
+        inputs, self.scales, default=None, check_used_all_replace_keys=False)
+    # if shifts/scales have missing values present in `inputs`, we insert `None`
+    # for the default. If corresponding `inputs` is not `None`, this will raise
+    # an error, as expected. This works because tree_map skips `None` values in
+    # the first argument, as long as all dictionary keys match.
+    result = jax.tree_util.tree_map(
+        lambda x, y, z: (x - y) / z, inputs, shifts, scales)
+    return from_dict_fn(result)
+
+
+@gin.register
+class InverseShiftAndNormalize(hk.Module):
+  """Inverse of the `ShiftAndNormalize` for the same `shifts/scales`."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      shifts: typing.Pytree,
+      scales: typing.Pytree,
+      global_scale: Optional[float] = None,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.shifts = shifts
+    if global_scale is not None:
+      scales = jax.tree_util.tree_map(lambda x: x * global_scale, scales)
+    self.scales = scales
+
+  def __call__(self, inputs):
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    shifts = pytree_utils.replace_with_matching_or_default(
+        inputs, self.shifts, default=None, check_used_all_replace_keys=False)
+    scales = pytree_utils.replace_with_matching_or_default(
+        inputs, self.scales, default=None, check_used_all_replace_keys=False)
+    # if shifts/scales have missing values present in `inputs`, we insert `None`
+    # for the default. If corresponding `inputs` is not `None`, this will raise
+    # an error, as expected. This works because tree_map skips `None` values in
+    # the first argument, as long as all dictionary keys match.
+    result = jax.tree_util.tree_map(
+        lambda x, y, z: (None if x is None else x * z + y),
+        inputs,
+        shifts,
+        scales,
+        is_leaf=lambda x: x is None,
+    )
+    return from_dict_fn(result)
+
+
+@gin.register
+class ToModalWithDivCurlTransform(hk.Module):
+  """Module that converts inputs to modal replacing velocity with div/curl."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    if 'u' not in inputs or 'v' not in inputs:
+      raise ValueError('Inputs to ToModalWithDivCurlTransform must include `u, '
+                       f'v`, got keys: {inputs.keys()}')
+    sec_lat = 1 / self.coords.horizontal.cos_lat
+    u, v = inputs.pop('u'), inputs.pop('v')
+    # here u,v stand for velocity / cos(lat), but the cos(lat) is cancelled in
+    # divergence and curl operators below.
+    inputs['u'] = u * sec_lat
+    inputs['v'] = v * sec_lat
+    to_modal_fn = lambda x: (self.coords.horizontal.to_modal(x)  # pylint: disable=g-long-lambda
+                             if x is not None else None)
+    modal_outputs = jax.tree_util.tree_map(to_modal_fn, inputs)
+    u, v = modal_outputs.pop('u'), modal_outputs.pop('v')
+    modal_outputs['divergence'] = self.coords.horizontal.div_cos_lat((u, v))
+    modal_outputs['vorticity'] = self.coords.horizontal.curl_cos_lat((u, v))
+    return modal_outputs
+
+
+@gin.register
+class ToModalDiffOperators(hk.Module):
+  """Module that returns grad and laplacian features of inputs fields.
+
+  To avoid accidental accumulation of the cos(lat) factors, features must be
+  keyed using typing.KeyWithCosLatFactor namedtuple.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+
+  def __call__(
+      self,
+      inputs: Mapping[typing.KeyWithCosLatFactor, typing.Array],
+  ) -> Mapping[typing.KeyWithCosLatFactor, typing.Array]:
+    features = {}
+    for k, value in inputs.items():
+      name, cos_lat_order = k.name, k.factor_order
+      d_value_dlon, d_value_dlat = self.coords.horizontal.cos_lat_grad(value)
+      laplacian_value = self.coords.horizontal.laplacian(value)
+      dlon_key = typing.KeyWithCosLatFactor(name + '_dlon', cos_lat_order + 1)
+      dlat_key = typing.KeyWithCosLatFactor(name + '_dlat', cos_lat_order + 1)
+      del2_key = typing.KeyWithCosLatFactor(name + '_del2', cos_lat_order)
+      features[dlon_key] = d_value_dlon
+      features[dlat_key] = d_value_dlat
+      features[del2_key] = laplacian_value
+    return features
+
+
+@gin.register
+class ModalToNodalTransform(hk.Module):
+  """Transform that converts modal inputs to nodal representation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return self.coords.horizontal.to_nodal(inputs)
+
+
+@gin.register
+class NodalToModalTransform(hk.Module):
+  """Transform that converts nodal inputs to modal representation."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return self.coords.horizontal.to_modal(inputs)
+
+
+@gin.register
+class ClipTransform(hk.Module):
+  """Transform that clips highest total wavenumber in inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      wavenumbers_to_clip: int = 1,
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.coords = coords
+    self.wavenumbers_to_clip = wavenumbers_to_clip
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    return self.coords.horizontal.clip_wavenumbers(
+        inputs, self.wavenumbers_to_clip
+    )
+
+
+@gin.register
+class NondimensionalizeTransform(hk.Module):
+  """Transform that nondimensionalizes inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      input_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del coords, dt, aux_features, input_coords  # unused.
+    super().__init__(name=name)
+    self.inputs_to_units_mapping = inputs_to_units_mapping
+    self.nondimensionalize = physics_specs.nondimensionalize
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    inputs_to_units_mapping = pytree_utils.replace_with_matching_or_default(
+        inputs, self.inputs_to_units_mapping, default=None,
+        check_used_all_replace_keys=False,
+    )
+    nondim_fn = lambda x, y: self.nondimensionalize(x * typing.Quantity(y))
+    result = jax.tree_util.tree_map(nondim_fn, inputs, inputs_to_units_mapping)
+    return from_dict_fn(result)
+
+
+@gin.register
+class RedimensionalizeTransform(hk.Module):
+  """Transform that redimensionalizes inputs."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      output_coords: coordinate_systems.CoordinateSystem,
+      inputs_to_units_mapping: Dict[str, str],
+      name: Optional[str] = None,
+  ):
+    """See `time_integration.exponential_filter` for details."""
+    del coords, dt, aux_features, output_coords  # unused.
+    super().__init__(name=name)
+    self.inputs_to_units_mapping = inputs_to_units_mapping
+    self.dimensionalize = physics_specs.dimensionalize
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    inputs_to_units_mapping = pytree_utils.replace_with_matching_or_default(
+        inputs, self.inputs_to_units_mapping, default=None,
+        check_used_all_replace_keys=False,
+    )
+    dim_fn = lambda x, y: self.dimensionalize(x, typing.Quantity(y)).m
+    result = jax.tree_util.tree_map(dim_fn, inputs, inputs_to_units_mapping)
+    return from_dict_fn(result)
+
+
+@gin.register
+class SequentialTransform(hk.Module):
+  """Transform module that combines multiple transforms applied sequentially."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      transform_modules: Sequence[TransformModule],
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.transform_fns = [module(coords, dt, physics_specs, aux_features)
+                          for module in transform_modules]
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    for transform_fn in self.transform_fns:
+      inputs = transform_fn(inputs)
+    return inputs
+
+
+@gin.register
+class LevelScale(hk.Module):
+  """Transforms inputs by scaling different vertical levels."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      scales: Sequence[float],
+      keys_to_scale: Sequence[str] = tuple(),
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.scale_fn = functools.partial(
+        coordinate_systems.scale_levels_for_matching_keys,
+        scales=np.asarray(scales),
+        keys_to_scale=keys_to_scale)
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    return self.scale_fn(inputs)
+
+
+@gin.register
+class InverseLevelScale(hk.Module):
+  """Transforms inputs by inverse scaling different vertical levels."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      scales: Sequence[float],
+      keys_to_scale: Sequence[str] = tuple(),
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.scale_fn = functools.partial(
+        coordinate_systems.scale_levels_for_matching_keys,
+        scales=1/np.asarray(scales),
+        keys_to_scale=keys_to_scale)
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    return self.scale_fn(inputs)
+
+
+@gin.register
+class HardClip(hk.Module):
+  """Transforms inputs by hard clipping inputs to (-max_value, max_value)."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      max_value: float,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    super().__init__(name=name)
+    self.clip_fn = functools.partial(
+        jnp.clip, min=-max_value, max=max_value)
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    return jax.tree_util.tree_map(self.clip_fn, inputs)
+
+
+@gin.register
+class SoftClip(hk.Module):
+  """Transforms inputs by clipping values to a range with smooth boundaries.
+
+  Attributes:
+    coords: horizontal and vertical descritization.
+    dt: time step of the model.
+    physics_specs: object describing the scales and physical constants.
+    aux_features: dictionary holding static features that the model may use.
+    max_value: specifies the range (-max_value, max_value) of return values.
+    hinge_softness: controls the softness of the smoothing at the boundaries;
+      values outside of the max_value range are mapped into intervals of width
+      approximately `log(2) * hinge_softness` on the interior of each boundary.
+    name: optional name of the module.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      max_value: float,
+      hinge_softness: float = 1.0,
+      name: Optional[str] = None,
+  ):
+    del coords, dt, physics_specs, aux_features  # unused.
+    if max_value < 0 or hinge_softness < 0:
+      raise ValueError('max_value and hinge_softness must be positive, '
+                       f'{max_value=}, {hinge_softness=}')
+    super().__init__(name=name)
+    low = -max_value
+    high = max_value
+    hinge = hinge_softness
+    softplus_fn = lambda x: hinge * jax.nn.softplus(x / hinge)
+    self.clip_fn = lambda x: (  # pylint: disable=g-long-lambda
+        -softplus_fn(high - low - softplus_fn(x - low)) *
+        (high - low) / (softplus_fn(high - low)) + high)
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    return jax.tree_util.tree_map(self.clip_fn, inputs)
+
+
+@gin.register
+class ToModalDiffOperatorsWithFiltering(hk.Module):
+  """Module that returns filtered grad and laplacian features of inputs fields.
+
+  To avoid accidental accumulation of the cos(lat) factors, features must be
+  keyed using typing.KeyWithCosLatFactor namedtuple.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      filter_attenuations: Tuple[float, ...] = tuple(),
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    self.coords = coords
+    self.attenuations = filter_attenuations
+    feature_filters = []
+    for attenuation in filter_attenuations:
+      feature_filters.append(
+          filters.DataExponentialFilter(
+              coords, dt, physics_specs, aux_features,
+              order=1, attenuation=attenuation))
+    self.feature_filters = feature_filters
+
+  def __call__(
+      self,
+      inputs: Mapping[KeyWithCosLatFactor, typing.Array],
+  ) -> Mapping[KeyWithCosLatFactor, typing.Array]:
+    features = {}
+    for k, value in inputs.items():
+      name, cos_lat_order = k.name, k.factor_order
+      for filter_fn, att in zip(self.feature_filters, self.attenuations):
+        filtered_value = filter_fn(value)
+        d_value_dlon, d_value_dlat = self.coords.horizontal.cos_lat_grad(
+            filtered_value)
+        laplacian_value = self.coords.horizontal.laplacian(filtered_value)
+        # since gradient values picked up cos_lat factor we increment the
+        # corresponding key. This factor is adjusted at the caller level.
+        dlon_key = KeyWithCosLatFactor(
+            name + f'_dlon_{att}', cos_lat_order + 1, att)
+        dlat_key = KeyWithCosLatFactor(
+            name + f'_dlat_{att}', cos_lat_order + 1, att)
+        del2_key = KeyWithCosLatFactor(
+            name + f'_del2_{att}', cos_lat_order, att)
+        features[dlon_key] = d_value_dlon
+        features[dlat_key] = d_value_dlat
+        features[del2_key] = laplacian_value
+    return features
+
+
+@gin.register
+class TruncateSigmaLevels(hk.Module):
+  """Transform module that truncates vertical levels for specified variables."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      sigma_ranges: dict[str, Tuple[float, float]],
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del dt, physics_specs, aux_features  # unused.
+    self.sigma_ranges = sigma_ranges
+    self.sigma_levels = coords.vertical.centers
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    """Returns `inputs` where only specified levels are retained."""
+
+    def _slice_fn(x, sigma_range):
+      """Returns `x` sliced to include values in `sigma_range`."""
+      sigma_min_slice, sigma_max_slice = sigma_range
+      lower_index = np.argmax((self.sigma_levels - sigma_min_slice) > 0)
+      if sigma_max_slice > np.max(self.sigma_levels):
+        upper_index = len(self.sigma_levels)
+      else:
+        upper_index = np.argmin((self.sigma_levels - sigma_max_slice) < 0)
+      return x[slice(lower_index, upper_index), ...]
+
+    def recurse_and_replace(x: dict[str, Any],
+                            y: dict[str, Any],
+                            default=None) -> dict[str, Any]:
+      """Copy x, setting leaf values to `default` or value from y if keys match."""
+      return {
+          k: (
+              y.get(k, default)
+              if not isinstance(v, dict)
+              else recurse_and_replace(v, y, default)
+          )
+          for k, v in x.items()
+      }
+
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    sigma_ranges_extended = recurse_and_replace(
+        inputs, self.sigma_ranges, default=(0, 1)
+    )
+    outputs = jax.tree_util.tree_map(_slice_fn, inputs, sigma_ranges_extended)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+class TakeSurfaceAdjacentSigmaLevel(hk.Module):
+  """Transform module that retains only the vertical level nearest to Earth surface for all variables."""
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      name: Optional[str] = None,
+  ):
+    super().__init__(name=name)
+    del coords, dt, physics_specs, aux_features  # unused.
+
+  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
+    """Returns `inputs` where only last sigma level is retained."""
+
+    def _slice_fn(x):
+      return x[slice(-1, None), ...]
+
+    inputs, from_dict_fn = pytree_utils.as_dict(inputs)
+    outputs = jax.tree_util.tree_map(_slice_fn, inputs)
+    return from_dict_fn(outputs)
+
+
+@gin.register
+@dataclasses.dataclass
+class FeatureSelector:
+  """Features transform that retains items whose keys match against regex.
+
+  Attributes:
+    regex_patterns: regular expression pattern that specifies the set of keys
+      from `inputs` that will be returned by __call__ method.
+  """
+  regex_patterns: str
+
+  def __call__(
+      self,
+      inputs: Dict[str, typing.Array],
+  ) -> Dict[str, typing.Array]:
+    outputs = {}
+    for k, v in inputs.items():
+      if re.fullmatch(self.regex_patterns, k):
+        outputs[k] = v
+    return outputs
+
+
+@gin.register
+class BroadcastTransform:
+  """Features transform that broadcasts all features."""
+
+  def __init__(self, *args, **kwargs):
+    del args, kwargs  # unused.
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    leaves, tree_def = jax.tree_util.tree_flatten(inputs)
+    leaves = jnp.broadcast_arrays(*leaves)
+    return jax.tree_util.tree_unflatten(tree_def, leaves)
+
+
+@gin.register
+class SquashLevelsTransform:
+  """Transform that "squashes" values of inputs depending on their sigma level.
+
+  Multiplies inputs by the piecewise linear values used to "squash" inputs
+  by sigma level. See function χ definition at: http://screen/5V3jzU7ZFA4vVJP
+
+  The squash is paramtereizaed by low_cutoffs and high_cutoffs.
+  On Palmer 2009 (http://shortn/_56HCcQwmSS) page 4, the cutoffs for
+  perturbations are given. Below are translated to sigma levels values:
+    low_cutoffs: (100hPa, 50hPa)
+    low_cutoffs: (0.05, 0.1),
+    high_cutoffs: (1300m, 300m)
+    high_cutoffs: (0.86, 0.965)
+
+  Inputs that have a singleton or no level dimension are assumed defined at
+  the highest value of sigma ("surface level").
+
+  Attributes:
+    coords: horizontal and vertical descritization.
+    dt: time step of the model.
+    physics_specs: object describing the scales and physical constants.
+    aux_features: dictionary holding static features that the model may use.
+    low_cutoffs: σ=low_cutoffs[0] is when χ starts linearly increasing from 0.
+      σ=low_cutoffs[1] is when χ levels out at 1
+    high_cutoffs: σ=high_cutoffs[0] is when χ starts linearly decreasing from 1.
+      σ=high_cutoffs[1] is when χ reaches 0.
+  """
+
+  def __init__(
+      self,
+      coords: coordinate_systems.CoordinateSystem,
+      dt: float,
+      physics_specs: Any,
+      aux_features: typing.AuxFeatures,
+      low_cutoffs: Sequence[float] = (0.05, 0.1),
+      high_cutoffs: Sequence[float] = (0.86, 0.965),
+  ):
+    del dt, physics_specs, aux_features  # unused.
+    if not isinstance(coords.vertical, sigma_coordinates.SigmaCoordinates):
+      raise ValueError(f'Cannot apply sigma_squash on {coords.vertical=}')
+    sigma = coords.vertical.centers
+    if len(low_cutoffs) != 2:
+      raise ValueError(f'{len(low_cutoffs)=} but should have been 2.')
+    if len(high_cutoffs) != 2:
+      raise ValueError(f'{len(high_cutoffs)=} but should have been 2.')
+
+    low_func = (sigma - low_cutoffs[0]) / (low_cutoffs[1] - low_cutoffs[0])
+    high_func = (high_cutoffs[1] - sigma) / (high_cutoffs[1] - high_cutoffs[0])
+
+    # lower_bound is a function equal to the squasher between
+    # low_cutoffs[0] and high_cutoffs[1].
+    # It becomes negative outside that range.
+    lower_bound = np.minimum(1., np.minimum(low_func, high_func))
+    self._sigma_squash = np.maximum(0., lower_bound)[:, np.newaxis, np.newaxis]
+
+  def __call__(self, inputs: typing.Pytree) -> typing.Pytree:
+    def squash_per_level_only(x):
+      shape = jnp.shape(x)
+      ndim = len(shape)
+      if ndim >= 3 and shape[-3] > 1:  # If defined per-level
+        return x * self._sigma_squash
+      elif ndim in {2, 3}:
+        return x * self._sigma_squash[-1]   # If defined at surface level
+      else:
+        return x
+    return jax.tree_util.tree_map(squash_per_level_only, inputs)
+
+
+@gin.register
+def add_prefix(features: dict[str, Any], prefix: str) -> dict[str, Any]:
+  """Adds prefix to keys in features."""
+  return {prefix + k: v for k, v in features.items()}
+
+
+def straight_through(
+    f: Callable[[typing.Array], typing.Array],
+) -> Callable[[typing.Array], typing.Array]:
+  """Straight-through estimator of `func`.
+
+  The "straight-through" estimator is a trick that fools auto-diff into
+  assigning a constant gradient (≡ 1) to a function.
+  See http://shortn/_kRQjMbF2QF
+
+  Args:
+    f: Callable mapping arrays to arrays. May be non-differentiable.
+
+  Returns:
+    g: Function g such that g(x) ≡ f(x) and g'(x) ≡ 1.
+  """
+  def straight_through_f(x):
+    zero = x - jax.lax.stop_gradient(x)
+    return zero + jax.lax.stop_gradient(f(x))
+  return straight_through_f

model/reference_code/datasets.py

@@ -0,0 +1,70 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Loads datasets."""
+
+import functools
+import itertools
+import json
+import logging
+import math
+import multiprocessing
+import random
+from typing import Any, Callable, Iterator, Mapping, Optional, Tuple
+
+import jax
+import numpy as np
+import pandas as pd
+import tensorflow.compat.v2 as tf
+import xarray
+
+
+Pytree = Any
+# pylint: disable=g-bare-generic
+# pylint: disable=logging-fstring-interpolation
+
+
+def drop_static_vars(dataset: xarray.Dataset) -> xarray.Dataset:
+  """Drop fields that are static and do not vary with time."""
+  has_sample_dim = 'sample' in dataset.coords
+  vars_to_drop = []
+  for name, var in dataset.items():
+    if 'time' not in var.dims:
+      vars_to_drop.append(name)
+    elif has_sample_dim and var.dims[:2] != ('sample', 'time'):
+      raise ValueError(f'dimensions for variable {name} do not start with '
+                       f"'sample' and 'time': {var.dims}")
+    elif not has_sample_dim and var.dims[0] != 'time':
+      raise ValueError(f'dimensions for variable {name} do not start with '
+                       f"'time': {var.dims}")
+  return dataset.drop_vars(vars_to_drop)
+
+
+def attrs_from_dataset(
+    dataset: xarray.Dataset,
+    time_series_length: int,
+    subsample_rate: int = 1,
+) -> dict:
+  """Extracts attributes from `dataset`."""
+  attrs = dict(dataset.attrs)
+  attrs['trajectory_length'] = time_series_length
+  attrs['time_subsample_rate'] = subsample_rate
+  delta_t = (dataset.time[1] - dataset.time[0]).data
+  if not np.issubdtype(dataset.time.dtype, np.floating):
+    logging.info(f'converting non-float {delta_t=} to seconds')
+    delta_t = np.timedelta64(delta_t, 's') / np.timedelta64(1, 's')
+    attrs['save_dt_units'] = 's'
+  else:
+    attrs['save_dt_units'] = 'dimensionless'
+  attrs['save_dt'] = float(delta_t) * subsample_rate
+  return attrs

model/reference_code/experiment.py

@@ -0,0 +1,1400 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# pylint: disable=line-too-long
+r"""Pseudocode for training NeuralGCM models."""
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+import dataclasses
+import functools
+import logging
+import math
+from typing import Any, Callable, NamedTuple, Optional, cast
+
+from absl import app
+import datasets
+from dinosaur import typing
+from dinosaur import xarray_utils
+import gin
+from google_proprietary_code import checkpoint
+from google_proprietary_code import experiment
+from google_proprietary_code import experiment_utils
+# proprietary imports
+from google_proprietary_code import profiling_util
+from google_proprietary_code import streaming
+from google_proprietary_code import timing_util
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import jax.sharding
+import model.reference_code.metrics as metrics
+import model.reference_code.metrics_base as metrics_base
+import model.reference_code.metrics_util as metrics_util
+from ml_collections import config_dict
+from model.legacy import model_builder
+from model.legacy import model_utils
+from model.legacy import optimization
+from model.legacy import physics_specifications
+import numpy as np
+import optax
+import pandas as pd
+import model.reference_code.reader as reader
+import model.reference_code.stochastic_losses as stochastic_losses
+import tensorflow as tf
+import model.reference_code.train_utils as train_utils
+import xarray
+
+Params = typing.Params
+PyTree = Any
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+TrajectoryFn = Callable[
+    [Params, jax.Array, PyTree, PyTree],
+    tuple[TrajectoryRepresentations, TrajectoryRepresentations],
+]
+
+tree_map = jax.tree_util.tree_map
+
+# pylint: disable=logging-format-interpolation
+# pylint: disable=logging-fstring-interpolation
+
+
+@gin.configurable(allowlist=['constructor'])
+def get_loss_obj(
+    trajectory_spec: metrics_util.TrajectorySpec,
+    constructor: Callable[..., metrics_base.Loss] = gin.REQUIRED,
+) -> metrics_base.Loss:
+  """Returns configured loss_fn on first `trajectory_length` time slices."""
+  return constructor(trajectory_spec)
+
+
+EvaluatorDict = dict[str, metrics_base.Evaluator]
+TrainEvalIteratorTuple = tuple[
+    Iterator[Any],
+    train_utils.TrainStepFunction,
+    Callable[..., Any],
+    dict[str, Any],
+]
+
+
+@gin.configurable(allowlist=['constructor'])
+def get_metrics_dict(
+    trajectory_spec: metrics_util.TrajectorySpec,
+    eval_time_steps: Sequence[int],
+    loss: metrics_base.Loss,
+    constructor: Callable[..., EvaluatorDict] = metrics.default_metrics,
+    is_ensemble_data: bool = False,
+) -> EvaluatorDict:
+  """Returns configured loss_fn on first `trajectory_length` time slices."""
+  return constructor(
+      trajectory_spec, eval_time_steps, loss, is_ensemble_data=is_ensemble_data
+  )
+
+
+# start legacy configurables
+#
+# Keep these around for now (even though they are no-ops) so we can run
+# inference on old models.
+
+
+@gin.configurable
+def get_loss_fn(loss_fn):
+  raise NotImplementedError
+
+
+@gin.register
+def weighted_l2_cumulative_loss(weights, scale):
+  raise NotImplementedError
+
+
+# end legacy configurables
+
+
+def ema_params_tree(num_steps):
+  """Creates an EMAParamsTree object based on num_steps.
+
+  Args:
+    num_steps: average number of optimization steps to include in the
+      exponential moving average of model weights.
+
+  Returns:
+    Haiku module.
+  """
+  # https://en.wikipedia.org/wiki/Moving_average#Relationship_between_SMA_and_EMA
+  decay = 1 - 2 / (num_steps + 1)
+  return hk.EMAParamsTree(decay)
+
+
+def _model_inner_steps_for_data(
+    ds: xarray.Dataset,
+    model_specs: model_builder.ModelSpecs,
+    rtol: float = 1e-6,
+) -> int:
+  """Calculates model inner steps based on data time step."""
+  data_dt = xarray_utils.nondim_time_delta_from_time_axis(
+      ds.time.data, model_specs.physics_specs
+  )
+  inner_steps = round(data_dt / model_specs.dt)
+  if abs(inner_steps * model_specs.dt - data_dt) / data_dt > rtol:
+    raise RuntimeError(  # pylint: disable=g-doc-exception
+        f'{model_specs.dt=} does not divide evenly into {data_dt=}'
+    )
+  return inner_steps
+
+
+def _get_datetime_forecast_starts(
+    sample_count: int,
+    first_start: pd.Timestamp,
+    last_start: pd.Timestamp,
+) -> pd.DatetimeIndex:
+  """Get equispaced forecast start times for evaluating against ERA5."""
+  if first_start.hour != 0:
+    raise ValueError(f'dataset times must start at midnight: {first_start=}')
+  # Round-up to midnight following the last forecast day (e.g., the start of
+  # the next year).
+  stop = last_start.ceil('1D')
+  # Equally spaced from start (inclusive) to stop (exclusive).
+  start_times = pd.date_range(first_start, stop, periods=sample_count + 1)[:-1]
+  # To match ECMWF, all forecasts should be initialized at 0z or 12z. Here we
+  # alternate start times.
+  parity = np.arange(sample_count) % 2
+  return start_times.round('1D') + parity * pd.Timedelta('12H')
+
+
+P = jax.sharding.PartitionSpec
+
+
+class ExperimentState(NamedTuple):
+  opt_state: PyTree
+  params: PyTree
+  ema_params: PyTree
+
+
+class Experiment(experiment.AbstractExperiment):
+  """Training experiment based on trajectory loss minimization."""
+
+  def __init__(
+      self,
+      experiment_dir: str,
+      config: Optional[config_dict.ConfigDict] = None,
+  ):
+    """Creates an instance of a training scheme class.
+
+    Args:
+      experiment_dir: Path to experiment directory.
+      config: config struct setting up the experiment.
+    """
+    if config is None:
+      config = experiment_config.get_config()
+
+    super().__init__(
+        experiment_dir,
+        config.distributed_training,
+        writer_names=['train', 'eval', 'eval_ema'],
+    )
+    logging.info('Experiment config:\n%s', config)
+    self.config = config
+
+    self.train_ds = xarray_utils.open_dataset(config.train_dataset_path)
+    self.eval_ds = xarray_utils.open_dataset(config.eval_dataset_path)
+
+    if 'sample' in self.train_ds.dims:
+      logging.warning('only using the first sample!')
+      self.train_ds = self.train_ds.isel(sample=0, drop=True)
+      self.eval_ds = self.eval_ds.isel(sample=0, drop=True)
+
+    train_attrs = self.train_ds.attrs
+
+    # Model instantiation and trajectory unroll functions
+    # Note: we use interactive mode in experiments to split gin-configurations
+    # into separate, distinct parts provided in the config_dict.
+    gin.enter_interactive_mode()
+    self.is_nodal = self.config.is_nodal
+
+    # parse and override all the gin things
+    gin.parse_config(config.model_gin_config)
+    gin.parse_config(config.optimizer_gin_config)
+    experiment_utils.parse_config_dict(config.gin_overrides)
+
+    logging.info('Parsed gin config string:\n%s', gin.config_str())
+
+    full_model_gin_config = gin.config_str()  # do not include physics config.
+    logging.info('With overrides gin config string:\n%s', gin.config_str())
+
+    self.data_coords = model_builder.coordinate_system_from_dataset(
+        self.train_ds
+    )
+    logging.info(f'{self.model_parallel_training=}')
+
+    if self.model_parallel_training:
+      # It does not make sense to use spatial parallelism with batch size per
+      # device larger than 1. Instead, you would get better performance from
+      # using less model parallelism.
+      if self.spatial_parallelism > 1 and self.config.batch_size_per_device > 1:
+        raise NotImplementedError(
+            f'{self.config.batch_size_per_device=} is not supported for model '
+            'parallel training'
+        )
+      self.data_coords = dataclasses.replace(
+          self.data_coords, spmd_mesh=self.spmd_mesh
+      )
+
+    # try getting aux_features from dataset, if not included we rely on
+    # `model_builder.get_model_specs` to supply necessary values.
+    try:
+      data_aux_features = xarray_utils.aux_features_from_xarray(self.train_ds)
+    except KeyError:
+      data_aux_features = {}
+
+    # when available, we parse physics_config_str from metadata in train_attrs.
+    if 'physics_config_str' in train_attrs:
+      physics_config_str = train_attrs['physics_config_str']
+      experiment_utils.parse_gin_config_without_imports(physics_config_str)
+    else:
+      logging.info(
+          'physics_config_str was not provided in the dataset, '
+          'hence it is expected to be specified in model_gin_config.'
+      )
+
+    self.physics_specs = physics_specifications.get_physics_specs()
+    self.model_specs = model_builder.get_model_specs(
+        self.data_coords, self.physics_specs, data_aux_features
+    )
+    logging.info(f'{self.model_specs=}')
+
+    self.train_inner_steps = _model_inner_steps_for_data(
+        self.train_ds, self.model_specs
+    )
+    self.eval_inner_steps = _model_inner_steps_for_data(
+        self.eval_ds, self.model_specs
+    )
+    if (
+        len(self.config.train_schedule_time_steps)
+        != len(self.config.train_schedule_boundaries) + 1
+    ):
+      raise ValueError(
+          f'{self.config.train_schedule_time_steps} should be one longer '
+          f'than {self.config.train_schedule_boundaries} but was not.'
+      )
+    if any(
+        t % self.train_inner_steps
+        for t in self.config.train_schedule_time_steps
+    ):
+      raise ValueError(
+          f'{self.train_inner_steps=} does not divide '
+          f'{self.config.train_schedule_time_steps=}'
+      )
+    if any(t % self.eval_inner_steps for t in self.config.eval_time_steps):
+      raise ValueError(
+          f'{self.eval_inner_steps=} does not divide '
+          f'{self.config.eval_time_steps=}'
+      )
+    if max(self.config.eval_time_steps) < max(
+        self.config.train_schedule_time_steps
+    ):
+      raise ValueError(
+          f'Training will not work since {max(self.config.eval_time_steps)=} <'
+          f' {max(self.config.train_schedule_time_steps)=}'
+      )
+
+    self._eval_trajectory_length = (
+        max(self.config.eval_time_steps) // self.eval_inner_steps + 1
+    )
+    self._trajectory_lengths = [
+        self.config.num_init_frames + n // self.train_inner_steps
+        for n in self.config.train_schedule_time_steps
+    ]
+    self._max_trajectory_length = max(
+        self._trajectory_lengths + [self._eval_trajectory_length]
+    )
+
+    self.reference_datetime = self.model_specs.aux_features[
+        xarray_utils.REFERENCE_DATETIME_KEY
+    ]
+
+    self.whirl_model = model_builder.WhirlModel(
+        **self.model_specs,
+        input_coords=self.data_coords,
+        output_coords=self.data_coords,
+    )
+    self.from_xarray_fn = self.whirl_model.from_xarray_fn
+
+    def trajectory_fwd(x, forcing_data, model, outer_steps, inner_steps):
+      trajectory_fn = model_utils.trajectory_with_inputs_and_forcing(
+          model, config.num_init_frames, start_with_input=True
+      )
+      return trajectory_fn(x, forcing_data, outer_steps, inner_steps)
+
+    self._trajectory_fwd = trajectory_fwd
+
+    # Checkpoint items.
+    self._model_dt = self.model_specs.dt
+    self._model_gin_config = full_model_gin_config
+
+    # optimizer configuration.
+    self.optimizer = optimization.optimizer()
+
+    # exponentially moving average params tracking.
+    ema_fn = hk.without_apply_rng(
+        hk.transform_with_state(
+            lambda x: ema_params_tree(config.ema_num_steps)(x)  # pylint: disable=unnecessary-lambda
+        )
+    )
+    self._ema_init = jax.jit(ema_fn.init)
+
+    def ema_update(params, ema_state):
+      return ema_fn.apply(None, ema_state, params)
+
+    self._ema_update = jax.jit(ema_update)
+
+    logging.info('Final active config string:\n%s', gin.config_str())
+
+  #
+  # Data inputs methods.
+  #
+
+  @functools.cached_property
+  def spmd_mesh(self) -> jax.sharding.Mesh:
+    n = self.config.model_parallelism.ensemble_shards
+    z = self.config.model_parallelism.z_shards
+    x = self.config.model_parallelism.x_shards
+    y = self.config.model_parallelism.y_shards
+    global_batch = jax.device_count() // (n * z * x * y)
+    if global_batch == 0:
+      raise ValueError(
+          f'{jax.device_count()=} is insufficient for '
+          f'{self.config.model_parallelism=}'
+      )
+    return train_utils.create_spmd_mesh(
+        {'batch': global_batch, 'ensemble': n, 'z': z, 'x': x, 'y': y}
+    )
+
+  @functools.cached_property
+  def degree_of_model_parallelism(self) -> int:
+    return math.prod(v for k, v in self.spmd_mesh.shape.items() if k != 'batch')
+
+  @functools.cached_property
+  def model_parallel_training(self) -> bool:
+    return self.degree_of_model_parallelism > 1
+
+  @functools.cached_property
+  def spatial_parallelism(self) -> int:
+    return math.prod(self.spmd_mesh.shape[k] for k in 'zxy')
+
+  def to_global_array(self, pytree: PyTree, global_batch_size: int) -> PyTree:
+    """Create a pytree of global JAX arrays from a pytree of NumPy arrays."""
+    # partition arrays along batch and spatial dimensions
+    return train_utils.make_distributed_array_from_local_arrays(
+        pytree,
+        self.spmd_mesh,
+        self.data_coords.physics_partition_spec,
+        global_batch_size,
+    )
+
+  def num_eval_batches(self, large_eval: bool) -> int:
+    if large_eval:
+      return self.config.num_eval_batches[-1]
+    else:
+      return self.config.num_eval_batches[0]
+
+  def steps_between_evals(self, large_eval: bool) -> int:
+    if large_eval:
+      return self.config.steps_between_evals[-1]
+    else:
+      return self.config.steps_between_evals[0]
+
+  def eval_batch_size_per_device(self, large_eval: bool) -> int:
+    if large_eval:
+      return self.config.eval_batch_size_per_device[-1]
+    else:
+      return self.config.eval_batch_size_per_device[0]
+
+  @functools.cached_property
+  def global_batch_size(self) -> int:
+    return (
+        jax.device_count()
+        // self.degree_of_model_parallelism
+        * self.config.batch_size_per_device
+    )
+
+  def global_eval_batch_size(self, large_eval: bool) -> int:
+    return (
+        jax.device_count()
+        // self.degree_of_model_parallelism
+        * self.eval_batch_size_per_device(large_eval)
+    )
+
+  def local_eval_batch_size(self, large_eval: bool) -> int:
+    return jax.local_device_count() * self.eval_batch_size_per_device(
+        large_eval
+    )
+
+  def _to_dataset_iter(
+      self, data: tf.data.Dataset, template: xarray.Dataset
+  ) -> Callable[[], Iterable[Any]]:
+    """Convert a tf.data.Dataset into a function that makes a data iterator."""
+    leading_dims_set = {x.shape[:2] for x in data.element_spec.values()}
+    assert len(leading_dims_set) == 1, leading_dims_set
+    local_batch_size, time_series_length = leading_dims_set.pop()
+    template = (
+        template.drop_vars('time')
+        .head(time=time_series_length)
+        .pipe(xarray.zeros_like)  # replace data with zeros
+        .pipe(datasets.drop_static_vars)
+        .transpose('time', ...)
+        .expand_dims(batch=local_batch_size)
+    )
+
+    def make_iterator():
+      for example_dict in data.as_numpy_iterator():
+        yield self.from_xarray_fn(template.copy(data=example_dict))
+
+    return make_iterator
+
+  def _read_shuffled_shard(
+      self,
+      dataset: xarray.Dataset,
+      time_series_length: int,
+      min_buffer_blocks: int,
+      shard_index: int,
+      shard_count: int,
+  ) -> tf.data.Dataset:
+    sampler = reader.Windower(
+        window_size=time_series_length,
+        stride_between_windows=self.config.train_time_sample_offset,
+    )
+    local_shard_count = max(
+        self.degree_of_model_parallelism, jax.local_device_count()
+    )
+    seed = train_utils.combine_rng_seeds(
+        self.config.dataset_rng_seed, shard_index, time_series_length
+    )
+    data = reader.read_shuffled_shard(
+        dataset,
+        sampler,
+        block_size_in_bytes=self.config.block_size_in_bytes / local_shard_count,
+        buffer_size_in_bytes=(
+            self.config.shuffle_buffer_size_in_bytes / local_shard_count
+        ),
+        min_buffer_blocks=min_buffer_blocks,
+        shard_index=shard_index,
+        shard_count=shard_count,
+        seed=seed,
+    )
+    return data
+
+  def _get_train_dataset(self) -> xarray.Dataset:
+    train_dataset = xarray_utils.ds_with_sim_time(
+        self.train_ds, self.physics_specs, self.reference_datetime
+    )
+
+    if self.config.train_dataset_time_slice:
+      time_slice = slice(*self.config.train_dataset_time_slice)
+      train_dataset = train_dataset.sel(time=time_slice)
+
+    if self.config.time_subsample_rate != 1:
+      raise NotImplementedError('subsampling on the fly is not supported yet')
+    if self.config.add_noise_to_input:
+      raise NotImplementedError('add_noise_to_input not supported yet')
+
+    return train_dataset
+
+  def _spatial_addressable_indices_map(
+      self, spatial_dim_sizes: tuple[int, int, int]
+  ) -> Mapping[jax.Device, tuple[slice, slice, slice, slice]]:
+    """Get slices for indexing global arrays to local devices."""
+    spec = P('batch', *self.data_coords.physics_partition_spec)
+    sharding = jax.sharding.NamedSharding(self.spmd_mesh, spec)
+    global_shape = (self.global_batch_size,) + spatial_dim_sizes
+    indices_map = sharding.addressable_devices_indices_map(global_shape)
+    indices_map = cast(
+        Mapping[jax.Device, tuple[slice, slice, slice, slice]], indices_map
+    )
+    return indices_map
+
+  def _read_model_parallel_dataset(
+      self,
+      dataset: xarray.Dataset,
+      read_shard: Callable[[xarray.Dataset, int], tf.data.Dataset],
+  ) -> tuple[tf.data.Dataset, xarray.Dataset]:
+    """Read a shard of a training dataset into tf.data.Dataset."""
+    indices_map = self._spatial_addressable_indices_map(
+        tuple(dataset.sizes[k] for k in ['level', 'longitude', 'latitude'])
+    )
+
+    shard_data: list[tf.data.Dataset] = []
+    for device in jax.local_devices():
+      indices = indices_map[device]
+      batch_index = indices[0].start or 0
+      selection = dict(zip(['level', 'longitude', 'latitude'], indices[1:]))
+      shard_dataset = dataset.isel(selection)
+      shard_data.append(read_shard(shard_dataset, batch_index))
+
+    choices = tf.data.Dataset.range(jax.local_device_count()).repeat()
+    data = tf.data.Dataset.choose_from_datasets(shard_data, choices)
+    template = shard_dataset
+    return data, template
+
+  def _build_train_inputs(
+      self, time_series_length: int
+  ) -> tuple[Callable[[], Any], dict[str, Any]]:
+    """Loads the training dataset and returns an iterator and train_attrs."""
+    train_dataset = self._get_train_dataset()
+    local_batch_size = (
+        # Size of data needed to satisfy batch_size_per_device.
+        self.config.batch_size_per_device
+        * jax.local_device_count()
+    )
+    if self.model_parallel_training:
+
+      def read_shard(shard_dataset, batch_index):
+        return self._read_shuffled_shard(
+            shard_dataset,
+            time_series_length,
+            min_buffer_blocks=1,
+            shard_index=batch_index,
+            shard_count=self.global_batch_size,
+        )
+
+      data, template = self._read_model_parallel_dataset(
+          train_dataset, read_shard
+      )
+
+    else:
+      data = self._read_shuffled_shard(
+          train_dataset,
+          time_series_length,
+          shard_index=jax.process_index(),
+          shard_count=jax.process_count(),
+          min_buffer_blocks=local_batch_size,
+      )
+      template = train_dataset
+
+    data = data.repeat()
+    data = data.batch(local_batch_size, drop_remainder=True)
+    data = data.prefetch(tf.data.AUTOTUNE)
+    train_iter = self._to_dataset_iter(data, template)
+    data_attrs = datasets.attrs_from_dataset(train_dataset, time_series_length)
+    return train_iter, data_attrs
+
+  def build_train_and_eval_iterators(
+      self,
+      schedule_idx: int,
+      start_step: int,
+      large_eval: bool,
+  ) -> TrainEvalIteratorTuple:
+    """Build new iterators for training at schedule_idx.
+
+    Args:
+      schedule_idx: Index into the rollout schedule.
+      start_step: Step at which this training run started at. This does not
+        change unless the Borg job dies and restarts.
+      large_eval: Whether this evaluation should be done over a larger set of
+        data.
+
+    Returns:
+      TrainEvalIteratorTuple: Tuple consisting of
+        get_train_data. Iterator providing next set of training data.
+        train_step_fn. train_utils.TrainStepFunction to update weights.
+        evaluate_fn. Callable to evalate metrics and write results.
+        ckpt_kwargs. dict[str, Any] of kwargs to add to the checkpoint.
+    """
+    num_train_time_steps = self.config.train_schedule_time_steps[schedule_idx]
+    trajectory_length = self._trajectory_lengths[schedule_idx]
+
+    train_traj_spec = metrics_util.TrajectorySpec(
+        trajectory_length,
+        self._max_trajectory_length,
+        self.train_inner_steps,
+        coords=self.model_specs.coords,
+        data_coords=self.data_coords,
+    )
+
+    # These only change on the first and last rollout, but re-make them anyways.
+    get_eval_data, eval_attrs = self.build_eval_inputs(
+        self.config.eval_dataset_time_slice,
+        large_eval,
+    )
+    get_eval_on_train, _ = self.build_eval_inputs(
+        self.config.train_dataset_time_slice,
+        large_eval,
+    )
+
+    evaluate_fn = functools.partial(
+        self.evaluate,
+        eval_batch_fn=self.get_eval_batch_fn(train_traj_spec),
+        get_eval_data=get_eval_data,
+        get_train_data=get_eval_on_train,
+        large_eval=large_eval,
+    )
+    get_train_data, train_attrs = self._build_train_inputs(trajectory_length)
+    ckpt_kwargs = {
+        'train_attrs': train_attrs,
+        'eval_attrs': eval_attrs,
+    }
+
+    train_step_fn = self.get_train_step_fn(
+        num_train_time_steps, train_traj_spec
+    )
+
+    if (
+        self.config.profile_with_xprof
+        and schedule_idx == 0
+        and experiment_utils.is_coordinator()
+        and start_step == 0
+    ):
+      train_step_fn = profiling_util.Traced(
+          train_step_fn,  # only the initial train_step is profiled.
+          trace_name='train step',
+          skip_steps=2,  # avoid JIT compilation
+          num_trace_steps=3,
+          enable_python_tracer=True,
+          host_trace_level=3,
+      )
+    return get_train_data(), train_step_fn, evaluate_fn, ckpt_kwargs
+
+  def _get_eval_dataset(self) -> xarray.Dataset:
+    return xarray_utils.ds_with_sim_time(
+        self.eval_ds, self.physics_specs, self.reference_datetime
+    )
+
+  def build_eval_inputs(
+      self,
+      dataset_time_slice: tuple[str, str] | None,
+      large_eval: bool,
+  ) -> tuple[Callable[[], Any], Any]:
+    """Returns an iterable over the data and data attrs for evaluation."""
+    eval_dataset = self._get_eval_dataset()
+
+    num_eval_batches = self.num_eval_batches(large_eval)
+    eval_batch_size_per_device = self.eval_batch_size_per_device(large_eval)
+    time_series_length = (
+        self._eval_trajectory_length + self.config.num_init_frames - 1
+    )
+    local_batch_size = eval_batch_size_per_device * jax.local_device_count()
+
+    if isinstance(eval_dataset.indexes['time'], pd.DatetimeIndex):
+      # For real world training data from ERA5, carefully sample starting
+      # times to ensure they are equally spaced across the year.
+      assert 'sample' not in eval_dataset.dims
+      logging.info('Using eval data loader for build_eval_inputs')
+      sample_count = self.global_eval_batch_size(large_eval) * num_eval_batches
+      if dataset_time_slice:
+        time_source = eval_dataset.time.loc[slice(*dataset_time_slice)]
+      else:
+        time_source = eval_dataset
+      first_start = time_source.indexes['time'][0]
+      last_start = time_source.indexes['time'][-1]
+      starts = _get_datetime_forecast_starts(
+          sample_count, first_start, last_start
+      )
+      logging.info(f'determined evaluation data for {sample_count=}: {starts=}')
+      offsets = eval_dataset.indexes['time'].get_indexer(starts)
+      sampler = reader.WindowerAtOffsets(
+          window_size=time_series_length, window_offsets=offsets
+      )
+      if self.model_parallel_training:
+
+        def read_shard(shard_dataset, batch_index):
+          selector = reader.ShardSelector(batch_index, len(starts))
+          return reader.read_timeseries(shard_dataset, sampler, selector)
+
+        data, template = self._read_model_parallel_dataset(
+            eval_dataset, read_shard
+        )
+
+      else:
+        selector = reader.ShardSelector(
+            jax.process_index(), jax.process_count()
+        )
+        data = reader.read_timeseries(eval_dataset, sampler, selector)
+        template = eval_dataset
+
+      data = data.batch(local_batch_size, drop_remainder=True)
+      data = data.cache()
+    else:
+      # For synthetic datasets (e.g., from Held-Suarez), use the same shuffling
+      # we use for reading training data.
+      assert not self.model_parallel_training
+      if dataset_time_slice:
+        eval_dataset = eval_dataset.sel(time=slice(*dataset_time_slice))
+      data = self._read_shuffled_shard(
+          eval_dataset,
+          time_series_length,
+          shard_index=jax.process_index(),
+          shard_count=jax.process_count(),
+          min_buffer_blocks=local_batch_size * num_eval_batches,
+      )
+      data = data.batch(local_batch_size, drop_remainder=True)
+      data = data.take(num_eval_batches)
+      template = eval_dataset
+
+    logging.info(f'created eval data: {data}')
+    eval_iter = self._to_dataset_iter(data, template)
+    data_attrs = datasets.attrs_from_dataset(eval_dataset, time_series_length)
+    return eval_iter, data_attrs
+
+  #
+  # Training and evaluation methods.
+  #
+
+  def _make_initial_experiment_state(
+      self,
+      rng: typing.PRNGKeyArray,
+      init_example,
+      init_forcing_data: typing.ForcingData,
+      init_params: Optional[typing.Params] = None,
+  ) -> ExperimentState:
+    """Makes initial parameters (via hk.Module.init)."""
+    if self.eval_inner_steps != self.train_inner_steps:
+      raise ValueError(
+          'KroneckerCorrelatedL2LossModule stddev will be ill defined since '
+          f'{self.eval_inner_steps=} != {self.train_inner_steps=}'
+      )
+    trajectory_length = list(init_example.values())[0].shape[0]
+
+    @jax.jit
+    def init(rng, init_example, init_forcing_data):
+      outer_steps = (trajectory_length - self.config.num_init_frames) + 1
+      # We need an "ensemble" dimension for stochastic losses, but parameters
+      # are fully replicated across the ensemble.
+      if init_params is None:
+        init_fn = jax.vmap(
+            self._make_unbatched_trajectory_fn(outer_steps).init,
+            in_axes=None,
+            out_axes=0,
+            spmd_axis_name='ensemble',
+            axis_size=1,
+        )
+        unsqueezd_params = init_fn(rng, init_example, init_forcing_data)
+        params = tree_map(lambda x: jnp.squeeze(x, axis=0), unsqueezd_params)
+      else:
+        params = init_params
+      opt_state = self.optimizer.init(params)
+      _, ema_state = self._ema_init(None, params)
+      experiment_state = ExperimentState(opt_state, params, ema_state)
+      experiment_state = train_utils.ensure_replicated(
+          experiment_state, mesh=self.spmd_mesh
+      )
+      return experiment_state
+
+    return init(rng, init_example, init_forcing_data)
+
+  def _make_unbatched_trajectory_fn(self, outer_steps: int):
+    """Haiku transformation of func giving (prediction, target) trajectories.
+
+    Args:
+      outer_steps: Number of outer steps the trajectory should take.
+
+    Returns:
+      hk transformed object. The .apply member maps
+        (params, rng, target, forcing_data) --> (prediction, target)
+    """
+    if self.train_inner_steps != self.eval_inner_steps:
+      # We share a trajectory for train/eval...so the spacing better be equal.
+      raise ValueError(f'{self.train_inner_steps=} != {self.eval_inner_steps=}')
+
+    @hk.transform
+    def unbatched_trajectory_fn(target, forcing_data):
+      """Compute Fwd(target[0]) on one single batch/device."""
+      # Shapes(target) ~ (n_t, n_z, n_m, n_l)
+      model = self.whirl_model.model_cls()
+      _, predicted_trajectory = self._trajectory_fwd(
+          x=target,
+          forcing_data=forcing_data,
+          model=model,
+          outer_steps=outer_steps,
+          inner_steps=self.train_inner_steps,
+      )
+      prediction, target = (
+          model_utils.compute_prediction_and_target_representations(
+              predicted_trajectory, target, forcing_data, model
+          )
+      )
+      return prediction, target
+
+    return unbatched_trajectory_fn
+
+  def _make_batch_trajectory_fn(
+      self,
+      outer_steps: int,
+  ) -> TrajectoryFn:
+    """Target, prediction representations with shape (batch, ensemble, ...)."""
+
+    ensembled_fn = jax.vmap(
+        # (params, rng, target, forcing_data) --> (prediction, target)
+        self._make_unbatched_trajectory_fn(outer_steps).apply,
+        in_axes=(None, 0, None, None),
+        spmd_axis_name='ensemble',
+    )
+
+    batch_ensembled_fn = jax.vmap(
+        # (params, rng, target, forcing_data) --> (prediction, target)
+        # Input shapes are:
+        #   params: (...)
+        #   rng: (batch, ensemble, ...)
+        #   target: (batch, time, ...)
+        #   forcing_data: (batch, time, ...)
+        ensembled_fn,
+        in_axes=(None, 0, 0, 0),
+        spmd_axis_name='batch',
+    )
+    return batch_ensembled_fn
+
+  def get_train_step_fn(
+      self,
+      num_train_time_steps: int,
+      traj_spec: metrics_util.TrajectorySpec,
+  ) -> train_utils.TrainStepFunction:
+    """Makes a function to update weights via gradient descent.
+
+    This function makes use of on-device batching. Multiple devices are combined
+    via an all reduce step whereby the average (across devices) gradient is
+    applied to each (on device) params.
+
+    Args:
+      num_train_time_steps: Number of time steps for the trajectory in this
+        training step.
+      traj_spec: Specification of training trajectory.
+
+    Returns:
+      all_reduced_train_step: Function to exectute one training step.
+        Params are injected into self._trajectory_fwd and Loss (if Loss requires
+        Haiku params).
+    """
+
+    batch_trajectory_fn = self._make_batch_trajectory_fn(
+        # (params, rng, target, forcing_data) --> [prediction, target]
+        outer_steps=num_train_time_steps // self.train_inner_steps
+        + 1,
+    )
+
+    loss_fn = get_loss_obj(traj_spec).evaluate
+    ensembled_loss_fn = jax.vmap(
+        loss_fn,
+        axis_name='ensemble',
+        spmd_axis_name='ensemble',
+    )
+    batch_ensembled_loss_fn = jax.vmap(
+        ensembled_loss_fn,
+        axis_name='batch',
+        spmd_axis_name='batch',
+    )
+
+    def batched_parameter_loss_fn(params, rng, target, forcing_data):
+      """Mean (over on-device batch members) of loss w.r.t parameters."""
+      # Input shapes are:
+      #   params: (...)
+      #   rng: (batch, ensemble, ...)
+      #   target: (batch, time, ...)
+      #   forcing_data: (batch, time, ...)
+      prediction, target = batch_trajectory_fn(
+          # The `target` and prediction returned are TrajectoryRepresentations.
+          # So don't just re-use the arg `target`.
+          params,
+          rng,
+          target,
+          forcing_data,
+      )
+      # dimensions (batch, ensemble)
+      per_example_loss = batch_ensembled_loss_fn(prediction, target)
+      # Average over ensemble and batch dimensions (technically, we don't have
+      # to average over ensemble with our current stochastic losses, but these
+      # values are already identical and this is cleaner than using array
+      # indexing)
+      overall_loss = jnp.mean(per_example_loss, axis=(0, 1))
+      assert overall_loss.ndim == 0
+      return overall_loss
+
+    # We would use donate_argnums here to update experiment_state in-place, but
+    # that would mean we could not save the checkpoint in a separable thread.
+    # Fortunately experiment_state is usually not too big (~100 MB).
+    @train_utils.jit_once
+    def train_step(experiment_state, rng, target_trajectory, forcing_data):
+      opt_state, params, ema_state = experiment_state
+      rng = train_utils.ensure_sharded_rng_key(rng, mesh=self.spmd_mesh)
+      loss, grad = jax.value_and_grad(batched_parameter_loss_fn)(
+          params, rng, target_trajectory, forcing_data
+      )
+      updates, opt_state = self.optimizer.update(grad, opt_state, params)
+      params = optax.apply_updates(params, updates)
+      _, ema_state = self._ema_update(params, ema_state)
+      experiment_state = ExperimentState(opt_state, params, ema_state)
+      experiment_state = train_utils.ensure_replicated(
+          experiment_state, mesh=self.spmd_mesh
+      )
+      return experiment_state, loss
+
+    return train_step
+
+  def get_eval_batch_fn(
+      self,
+      train_traj_spec: metrics_util.TrajectorySpec,
+  ) -> train_utils.EvalStepFunction:
+    """Makes a function that performs a single evaluation pass.
+
+    Args:
+      train_traj_spec: TrajectorySpec for training. Used to add the "loss"
+        evaluation metrics.
+
+    Returns:
+      Function mapping (params, rng, target, forcing_data) to dictionary of
+        scalar metric values. Parameters are injected into self._trajectory_fwd
+        and Loss (if Loss requires Haiku params).
+    """
+    eval_traj_spec = metrics_util.TrajectorySpec(
+        self._eval_trajectory_length,
+        self._max_trajectory_length,
+        steps_per_save=self.eval_inner_steps,
+        coords=self.model_specs.coords,
+        data_coords=self.data_coords,
+    )
+
+    eval_time_steps = [
+        t // self.eval_inner_steps for t in self.config.eval_time_steps
+    ]
+    if any(t % self.eval_inner_steps for t in self.config.eval_time_steps):
+      raise ValueError(
+          f'cannot evaluate {self.config.eval_time_steps=} with '
+          f'{self.eval_inner_steps=}'
+      )
+
+    batch_trajectory_fn = self._make_batch_trajectory_fn(
+        # (params, rng, target, forcing_data) --> [prediction, target]
+        outer_steps=max(
+            self._eval_trajectory_length,
+            train_traj_spec.trajectory_length,
+        ),
+    )
+
+    def unbatched_eval_fn(
+        prediction: TrajectoryRepresentations, target: TrajectoryRepresentations
+    ):
+      """Evaluate(target, Fwd(target[0])) on one single batch/device."""
+      metrics_dict = get_metrics_dict(
+          eval_traj_spec,
+          eval_time_steps,
+          get_loss_obj(train_traj_spec),
+          is_ensemble_data=bool(self.config.ensemble_size),
+      )
+      return train_utils.flatten_dict({
+          k: metric.evaluate(prediction, target)
+          for k, metric in metrics_dict.items()
+      })
+
+    ensembled_fn = jax.vmap(
+        unbatched_eval_fn, axis_name='ensemble', spmd_axis_name='ensemble'
+    )
+
+    batch_ensembled_fn = jax.vmap(
+        ensembled_fn,
+        axis_name='batch',
+        spmd_axis_name='batch',
+    )
+
+    @train_utils.jit_once
+    def batch_mean_eval_fn(params, rng, target, forcing_data):
+      """Computes mean (over batch members) of evaluation."""
+      # Input shapes for this function are:
+      #   params: (...)
+      #   rng: (batch, ensemble, ...)
+      #   target: (batch, time, ...)
+      #   forcing_data: (batch, time, ...)
+      rng = train_utils.ensure_sharded_rng_key(rng, mesh=self.spmd_mesh)
+      prediction, target = batch_trajectory_fn(
+          # The `target` and prediction returned are TrajectoryRepresentations.
+          # So don't just re-use the arg `target`.
+          params,
+          rng,
+          target,
+          forcing_data,
+      )
+      batch_eval_values = batch_ensembled_fn(prediction, target)
+      return tree_map(jnp.mean, batch_eval_values)
+
+    return batch_mean_eval_fn
+
+  def run_training(self):
+    """See base class."""
+    (
+        start_step,
+        times_restarted_on_nan,
+        step_auto_restart_began_at,
+        experiment_state,
+    ) = self.initialize_experiment(
+        initial_checkpoint_path=self.config.initial_checkpoint_path
+    )
+
+    if start_step >= self.config.num_training_steps:
+      logging.warning(
+          f'Attempting to start training at {start_step=} >='
+          f' {self.config.num_training_steps=}. Will simply return'
+      )
+      return
+
+    def logging_callback(step, loss, times_restarted_on_nan):
+      loss = float(jax.device_get(loss))
+      if step % max(self.steps_between_evals(False) // 100, 1) == 0:
+        logging.info(f'{step=}, {loss=}')
+      if (
+          self.config.error_with_nan_loss
+          and times_restarted_on_nan > self.config.max_nan_restarts
+      ):
+        raise RuntimeError(
+            f'NaN loss detected at {step=}, after too many restarts since'
+            f' {times_restarted_on_nan=} >'
+            f' {self.config.max_nan_restarts=}. Aborting.'
+        )
+
+    # monitor loss using a separate thread, so it doesn't block execution
+    logging_stream = streaming.SingleThreadExecutor(logging_callback)
+    logging.info('starting training from step=%s', start_step)
+    train_step_timer = timing_util.Timer()
+
+    rng_stream = train_utils.BatchedPRNGSequence(
+        jax.random.PRNGKey(self.config.init_rng_seed),
+        batch_shape=(self.global_batch_size, self.config.ensemble_size or 1),
+    )
+
+    loss = 1.0  # setting to a non-nan value when starting an experiment.
+    schedule_idx = None
+    ckpt_kwargs = {}
+
+    step = start_step
+    while step < self.config.num_training_steps:
+      old_schedule_idx = schedule_idx
+      schedule_idx = np.sum(  # compute which leg of the schedule we are at.
+          step > np.asarray(self.config.train_schedule_boundaries)
+      )
+      large_eval = (
+          schedule_idx == len(self.config.train_schedule_time_steps) - 1
+      )
+      if schedule_idx != old_schedule_idx:
+        (
+            train_iter,
+            train_step_fn,
+            evaluate_fn,
+            ckpt_kwargs,
+        ) = self.build_train_and_eval_iterators(
+            schedule_idx=schedule_idx,
+            start_step=start_step,
+            large_eval=large_eval,
+        )
+
+      if (
+          np.isnan(loss)
+          # No sense re-initializing if we've restarted a bunch already. Also,
+          # note that if error_with_nan_loss=True, we should raise and not have
+          # worry about the times_restarted_on_nan < max_nan_restarts here.
+          and times_restarted_on_nan <= self.config.max_nan_restarts
+      ):
+        # See also logging_callback, which may raise RuntimeError for NaN loss.
+        logging.warning(
+            f'NaN loss encountered at {step=}. Re-initializing and incrementing'
+            f' times_restarted_on_nan to {times_restarted_on_nan + 1}'
+        )
+        times_restarted_on_nan += 1
+        step, _, _, experiment_state = self.initialize_experiment(
+            target_step=step
+            - times_restarted_on_nan * self.config.restart_lookback_steps
+        )
+        step_auto_restart_began_at = step_auto_restart_began_at or step
+
+      # checkpoint
+      if step % self.steps_between_evals(large_eval) == 0:
+        self.save_checkpoint(
+            step,
+            experiment_state,
+            checkpoint_buffer_size=self.config.checkpoint_buffer_size,
+            times_restarted_on_nan=times_restarted_on_nan,
+            step_auto_restart_began_at=step_auto_restart_began_at,
+            **ckpt_kwargs,
+        )
+      elif step % self.config.steps_between_checkpoints == 0:
+        max_lookback_step = (
+            step
+            - self.config.max_nan_restarts * self.config.restart_lookback_steps
+        )
+        if (
+            # If there is no chance of a restart sequence overlapping with
+            # previously used checkpoints...
+            times_restarted_on_nan
+            and max_lookback_step > step_auto_restart_began_at
+        ):
+          logging.info(
+              f'Significant progress made since {step_auto_restart_began_at=}.'
+              f' In particular, {step=} Therefore set times_restarted_on_nan'
+              ' to 0'
+          )
+          times_restarted_on_nan = 0
+          step_auto_restart_began_at = None
+        self.save_checkpoint(
+            step,
+            experiment_state,
+            update_latest_only=True,
+            checkpoint_buffer_size=self.config.checkpoint_buffer_size,
+            times_restarted_on_nan=times_restarted_on_nan,
+            step_auto_restart_began_at=step_auto_restart_began_at,
+            **ckpt_kwargs,
+        )
+
+      # evaluate
+      if (step + 1) % self.steps_between_evals(large_eval) == 0:
+        if step > start_step:
+          with train_step_timer:
+            # train_step is non-blocking, so we need to block on the output
+            # of the previous training step to reliably time it.
+            experiment_state = jax.block_until_ready(experiment_state)
+          eval_interval = self.steps_between_evals(large_eval)
+          training_time = train_step_timer.total
+          logging.info(
+              f'training for {eval_interval} steps took '
+              f'{training_time:.1f} seconds'
+          )
+          self.record_scalar(
+              'train',
+              tag='seconds_per_train_step',
+              step=step,
+              value=training_time / eval_interval,
+          )
+          train_step_timer = timing_util.Timer()  # reset
+
+          if isinstance(train_step_fn, profiling_util.Traced):
+            memory_usage = train_step_fn.tracer.memory_usage  # pytype: disable=attribute-error
+            if memory_usage is not None:
+              self.record_scalar(
+                  'train',
+                  tag='peak_memory_usage_mib',
+                  step=step,
+                  value=memory_usage,
+              )
+
+        with timing_util.Timer() as eval_timer:
+          evaluate_fn(step, experiment_state, seed=step)
+        self.record_scalar(
+            'train',
+            tag='seconds_per_evaluation',
+            step=step,
+            value=eval_timer.average,
+        )
+        logging.info('evaluation pass took %.1f seconds', eval_timer.average)
+        self.flush_writers()  # flush all writers for this training step.
+
+      # train
+      with train_step_timer:
+        # go/xprof-instrument-jax
+        with jax.profiler.StepTraceAnnotation('train', step_num=step):
+          batch, forcing_data = self.to_global_array(
+              next(train_iter), self.global_batch_size
+          )
+          # This is necessary else ValueError.
+          # See http://sponge2/960c64a7-703f-4a19-8572-2c97dd9c01f3
+          with self.spmd_mesh:
+            experiment_state, loss = train_step_fn(
+                experiment_state, next(rng_stream), batch, forcing_data
+            )
+            # If we don't device_get (or similar), asynchronous execution
+            # resultsin this timed block taking almost no time. device_get does
+            # not result in longer runs, since each loop must eventually compute
+            # the loss, one way or another.
+            loss = jax.device_get(loss)
+          logging_stream.wait()
+          logging_stream.put(step, loss, times_restarted_on_nan)
+      step += 1
+    # End of while step < self.config.num_training_steps:
+
+    evaluate_fn(self.config.num_training_steps, experiment_state)
+    self.finalize_training(
+        self.config.num_training_steps,
+        experiment_state,
+        times_restarted_on_nan=times_restarted_on_nan,
+        step_auto_restart_began_at=step_auto_restart_began_at,
+        **ckpt_kwargs,
+    )
+
+  def make_dummy_inputs(self) -> tuple[Any, Any]:
+    train_dataset = xarray_utils.ds_with_sim_time(
+        self.train_ds, self.physics_specs, self.reference_datetime
+    )
+    dummy_ds = (
+        train_dataset.drop_vars('time')
+        .head(time=self.config.num_init_frames)
+        .pipe(xarray.zeros_like)  # replace data with zeros
+        .pipe(datasets.drop_static_vars)
+        .transpose('time', ...)
+    )
+    return self.from_xarray_fn(dummy_ds)
+
+  def initialize_experiment(
+      self,
+      initial_checkpoint_path: Optional[str] = None,
+      target_step: Optional[int] = None,
+  ) -> tuple[int, int, int | None, ExperimentState]:
+    """Returns training step and experiment state from checkpoint or init."""
+
+    if target_step:
+      ckpt = self.load_buffered_checkpoint(target_step=target_step)
+      if ckpt is None:
+        logging.info(
+            'No acceptable buffered checkpoint found for {target_step=}'
+        )
+      else:
+        logging.info(
+            f'Using buffered checkpoint, which has step={ckpt.step}. Ideally '
+            f'would have used {target_step=}'
+        )
+    else:
+      ckpt = self.load_latest_checkpoint()
+      if ckpt is None:
+        logging.info('No latest checkpoint found')
+      else:
+        logging.info(f'Using latest checkpoint, which has step={ckpt.step}')
+
+    init_params = None
+    if ckpt is None and initial_checkpoint_path is not None:
+      ckpt = checkpoint.load_checkpoint(initial_checkpoint_path)
+      if self.config.reset_initial_optimizer_state:
+        init_params = ckpt.eval_params
+        ckpt = None  # if resetting optimizer, carry over only init_params.
+
+    if ckpt is not None:
+      start_step = ckpt.step
+      logging.info(f'resuming from checkpoint at step={start_step}')
+      times_restarted_on_nan = getattr(ckpt, 'times_restarted_on_nan', 0)
+      step_auto_restart_began_at = getattr(
+          ckpt, 'step_auto_restart_began_at', None
+      )
+      experiment_state = ExperimentState(
+          ckpt.opt_state, ckpt.train_params, ckpt.ema_state
+      )
+    else:
+      logging.info('starting training with new weights')
+      start_step = 0
+      times_restarted_on_nan = 0
+      step_auto_restart_began_at = None
+      rng = jax.random.PRNGKey(self.config.init_rng_seed)
+      init_example, init_forcing_data = self.make_dummy_inputs()
+      experiment_state = self._make_initial_experiment_state(
+          rng, init_example, init_forcing_data, init_params=init_params
+      )
+
+    return (
+        start_step,
+        times_restarted_on_nan,
+        step_auto_restart_began_at,
+        experiment_state,
+    )
+
+  def _checkpoint_state(
+      self,
+      step: int,
+      experiment_state: ExperimentState,
+      times_restarted_on_nan: int,
+      step_auto_restart_began_at: int,
+      train_attrs: dict[str, Any],
+      eval_attrs: dict[str, Any],
+  ) -> checkpoint.CheckpointState:
+    """Returns a checkpoint state for a given experiment_state."""
+    opt_state, params, ema_state = experiment_state
+    ema_params, _ = self._ema_update(params, ema_state)
+    ckpt_state = checkpoint.CheckpointState(
+        train_params=params,
+        eval_params=ema_params,
+        opt_state=opt_state,
+        ema_state=ema_state,
+        step=step,
+        model_time_step=self._model_dt,
+        model_config_str=self._model_gin_config,
+        train_dataset_path=self.config.train_dataset_path,
+        eval_dataset_path=self.config.eval_dataset_path,
+        times_restarted_on_nan=times_restarted_on_nan,
+        step_auto_restart_began_at=step_auto_restart_began_at,
+        train_attrs=train_attrs,
+        eval_attrs=eval_attrs,
+    )
+    return ckpt_state
+
+  def evaluate(
+      self,
+      step,
+      experiment_state,
+      eval_batch_fn,
+      get_eval_data,
+      get_train_data,
+      large_eval: bool,
+      seed=0,
+  ):
+    """Evaluates the model on train and eval data and writes summaries.
+
+    Args:
+      step: global training step.
+      experiment_state: tuple of replicated step, optimizer state and EMA
+        (exponentially moving average) state for model parameters.
+      eval_batch_fn: function that, given parameters; rng; batch of data,
+        computes evaluation metric of interest on the given samples.
+      get_eval_data: callable that returns an iterator over evaluation data that
+        is used for produce summaries on unseen evaluation data.
+      get_train_data: callable that returns an iterator over training data that
+        is used for produce summaries on training data.
+      large_eval: Whether this evaluation is on the larger size eval data.
+      seed: seed for the random number generator to be used for evaluation.
+    """
+    num_eval_batches = self.num_eval_batches(large_eval)
+    if num_eval_batches == 0:
+      logging.warning(f'skipping evaluation: {num_eval_batches=}')
+      return
+
+    _, params, ema_state = experiment_state
+    ema_params, _ = self._ema_update(params, ema_state)
+
+    global_batch_size = self.global_eval_batch_size(large_eval)
+    rng_stream = train_utils.BatchedPRNGSequence(
+        jax.random.PRNGKey(seed),
+        batch_shape=(global_batch_size, self.config.ensemble_size or 1),
+    )
+    to_global_array = functools.partial(
+        self.to_global_array, global_batch_size=global_batch_size
+    )
+
+    # In theory, the mesh context manager should not be necessary because we use
+    # jit with sharded arrays (rather than xmap or pjit), but it seems to be
+    # required to avoid triggering bugs in JAX.
+    with self.spmd_mesh:
+      logging.info('evaluating on train dataset')
+      metrics_ = train_utils.streaming_mean(
+          rng_stream,
+          map(to_global_array, get_train_data()),
+          functools.partial(eval_batch_fn, params),
+      )
+      for tag, value in metrics_.items():
+        self.record_scalar('train', tag=tag, value=value, step=step)
+
+      logging.info('evaluating on test dataset')
+      metrics_ = train_utils.streaming_mean(
+          rng_stream,
+          map(to_global_array, get_eval_data()),
+          functools.partial(eval_batch_fn, params),
+      )
+      for tag, value in metrics_.items():
+        self.record_scalar('eval', tag=tag, value=value, step=step)
+
+      logging.info('evaluating EMA model on test dataset')
+      metrics_ = train_utils.streaming_mean(
+          rng_stream,
+          map(to_global_array, get_eval_data()),
+          functools.partial(eval_batch_fn, ema_params),
+      )
+      for tag, value in metrics_.items():
+        self.record_scalar('eval_ema', tag=tag, value=value, step=step)
+
+
+if __name__ == '__main__':
+  app.run(functools.partial(run_training.main, Experiment))

model/reference_code/linear_transforms.py

@@ -0,0 +1,342 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""LinearTransforms for use in Metrics."""
+import dataclasses
+import functools
+from typing import Callable, Mapping, Optional, Sequence
+from dinosaur import coordinate_systems
+from dinosaur import filtering
+from dinosaur import horizontal_interpolation
+from dinosaur import pytree_utils
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+import gin
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import model.reference_code.metrics_util as metrics_util
+
+
+Pytree = typing.Pytree
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+
+tree_leaves = jax.tree_util.tree_leaves
+tree_map = jax.tree_util.tree_map
+
+
+@dataclasses.dataclass
+class LinearTransform:
+  """A linear transformation, for TransformedL2Loss."""
+
+  trajectory_spec: metrics_util.TrajectorySpec
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    raise NotImplementedError
+
+
+LinearTransformConstructor = Callable[
+    [metrics_util.TrajectorySpec], LinearTransform
+]
+
+
+@dataclasses.dataclass
+class ComposedTransformForLoss(LinearTransform):
+  """Composition of multiple linear transformations for computation of loss.
+
+  Attributes:
+    components: components[i](self.trajectory_spec) initializes the i + 1 member
+      of self.transforms.
+    transforms: errors are transformed as error --> transforms[0](error) -->
+      transforms[1](error) --> ⋯. The 0th transform is inserted by this class as
+      TruncateToTrajectoryLength.
+  """
+
+  components: Sequence[LinearTransformConstructor]
+  transforms: Sequence[LinearTransform] = dataclasses.field(init=False)
+
+  def __post_init__(self):
+    # Insert TruncateToTrajectoryLength first in all cases. It's okay if it was
+    # already inserted... it is idempotent. This ensures that
+    #   len(self.transforms) = len(self.components) + 1
+    # in all cases.
+    components = [TruncateToTrajectoryLength] + list(self.components)
+    self.transforms = [
+        constructor(self.trajectory_spec) for constructor in components
+    ]
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    for transform in self.transforms:
+      errors = transform(errors, targets)
+    return errors
+
+
+@gin.register
+@dataclasses.dataclass
+class LegacyTimeRescaling(LinearTransform):
+  """Time scaling from WeightedL2CumulativeLoss."""
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    n = self.trajectory_spec.trajectory_length
+    steps_per_save = self.trajectory_spec.steps_per_save
+    scale = 1 if n == 1 else 1 / np.sqrt((n - 1) * steps_per_save)
+    return tree_map(lambda x: x * scale, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class TimeRescaling(LinearTransform):
+  """Time scaling that assumes error grows like a random walk.
+
+  This rescales errors like
+    errors --> errors / σ(T),
+    σ(T) := sqrt( sum(variance) / variance(T) )
+  where variance(T) is the assumed variance. A random walk has variance ∝ T.
+  This function uses similar scaling.
+
+  See also:
+  * Climatology vs. ENS CRPS values indicate skill difficult after 240 hrs
+    http://screen/8sVodqThEk6o693
+  * Plotting this function for various parameter values
+    http://screen/AubXNomsgm7g92o and http://gpaste/6727081386835968
+
+  Attributes:
+    base_squared_error_in_hours: Number of hours before assumed variance starts
+      growing (almost) linearly.
+    asymptotic_squared_error_in_hours: Number of hours before assumed variance
+      slows its growth. Set to None (the default) if variance grows indefinitely
+  """
+
+  base_squared_error_in_hours: float
+  asymptotic_squared_error_in_hours: Optional[float] = None
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    time_sizes = np.unique([x.shape[0] for x in tree_leaves(errors)])
+    if time_sizes.size != 1:
+      raise ValueError(f'Expected unique time dimension size. {time_sizes=}')
+    time_size = time_sizes[0]
+    if self.trajectory_spec.trajectory_length != time_size:
+      logging.info(
+          f'errors has {time_size=} !='
+          f' {self.trajectory_spec.trajectory_length=}. This is probably due to'
+          ' the Loss slicing via the time_step kwarg. Will use {time_size=}'
+          ' to compute scaling.'
+      )
+
+    steps_per_save = self.trajectory_spec.steps_per_save
+    t = np.arange(time_size) * steps_per_save
+    if self.asymptotic_squared_error_in_hours is not None:
+      # Rescale "time" `t`, so it stops growing when
+      #   t >> asymptotic_squared_error_in_hours.
+      t = t / (1 + t / self.asymptotic_squared_error_in_hours)
+
+    inv_variance = 1 / (1 + t / self.base_squared_error_in_hours)
+    scale = np.sqrt(inv_variance / inv_variance.sum())
+    scale = scale.reshape(-1, 1, 1, 1)
+
+    return tree_map(lambda x: x * scale, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class CustomTimeRescaling(LinearTransform):
+  """Custom time scaling that uses pre-specified values."""
+
+  scaling_weights: Sequence[float]
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    n = self.trajectory_spec.trajectory_length
+    scale = np.asarray(self.scaling_weights)[:n].reshape(-1, 1, 1, 1)
+    return tree_map(lambda x: x * scale, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class DelayedTimeRescaling(LinearTransform):
+  """Time scaling with smooth delay that transitions into hyperbolic decay."""
+
+  base_squared_error_in_hours: float
+  delay_power: float = 1.0
+  decay_power: float = 1.0
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    n = self.trajectory_spec.trajectory_length
+    steps_per_save = self.trajectory_spec.steps_per_save
+    t = np.arange(n) * steps_per_save
+
+    a = 1 / self.base_squared_error_in_hours
+    inv_variance = 1 / (
+        (1 + (a * t) ** self.delay_power) ** (1/self.decay_power))
+    scale = np.sqrt(inv_variance / inv_variance.sum())
+    scale = scale.reshape(-1, 1, 1, 1)
+
+    return tree_map(lambda x: x * scale, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class TruncateToTrajectoryLength(LinearTransform):
+  """Truncate errors to self.trajectory_spec.trajectory_length.
+
+  To ensure loss is computed over the correct trajectory length, this transform
+  should be used as the first step in any ComposedTransformForLoss.
+  """
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    n = self.trajectory_spec.trajectory_length
+    return metrics_util.extract_time_slice(errors, slice(0, n))
+
+
+@gin.register
+@dataclasses.dataclass
+class TotalWavenumberMasking(LinearTransform):
+  """Transform that masks out wavenumbers greater than `max_wavenumber`."""
+
+  max_wavenumber: int
+  is_encoded: bool = False
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    if self.is_encoded:
+      grid = self.trajectory_spec.coords.horizontal
+    else:
+      grid = self.trajectory_spec.data_coords.horizontal
+
+    modal_shape = grid.modal_shape
+    mask = np.arange(modal_shape[-1]) < self.max_wavenumber
+    mask = mask.astype(float)
+    return tree_map(lambda x: x * mask, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class ConservativeRegridder(LinearTransform):
+  """Linear transform that regrids."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      target_grid: spherical_harmonic.Grid,
+  ):
+    super().__init__(trajectory_spec=trajectory_spec)
+    self.regridder = horizontal_interpolation.ConservativeRegridder(
+        source_grid=trajectory_spec.coords.horizontal, target_grid=target_grid
+    )
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # Unused
+    return tree_map(self.regridder, errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class PerVariableRescaling(LinearTransform):
+  """Transform that reweights contribution per variable."""
+  weights: Pytree
+  scale: float = 1.0
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    weights = self.weights
+    if weights is None:
+      weights = tree_map(lambda x: 1.0, errors)
+    else:
+      weights = pytree_utils.replace_with_matching_or_default(
+          errors, weights, default=None,
+          check_used_all_replace_keys=True,
+      )
+    root_weights = tree_map(lambda w: np.sqrt(w * self.scale), weights)
+    return tree_map(jnp.multiply, errors, root_weights)
+
+
+@gin.register
+class ExponentialFilteringByLeadtime(LinearTransform):
+  """Applied leadtime dependent exponential filters to errors."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      filter_attenuations: typing.Pytree,
+      filter_orders: typing.Pytree,
+      is_encoded: bool = False,
+  ):
+    super().__init__(trajectory_spec=trajectory_spec)
+    n = trajectory_spec.trajectory_length
+    if is_encoded:
+      grid = trajectory_spec.coords.horizontal
+    else:
+      grid = trajectory_spec.data_coords.horizontal
+    # expand dims for `level, lon, total wavenumbers` so that filter parameters
+    # are applied to different time values.
+    to_array_fn = lambda x: np.expand_dims(np.array(x)[:n], axis=(1, 2, 3))
+    is_leaf = lambda x: isinstance(x, Sequence)
+    attenuations = tree_map(to_array_fn, filter_attenuations, is_leaf=is_leaf)
+    orders = tree_map(to_array_fn, filter_orders, is_leaf=is_leaf)
+    self.filter_fns = tree_map(
+        lambda a, p: filtering.exponential_filter(grid, a, p),
+        attenuations,
+        orders,
+    )
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    filter_fns = pytree_utils.replace_with_matching_or_default(
+        errors, self.filter_fns, default=None, check_used_all_replace_keys=True)
+    return tree_map(lambda fn, err: fn(err), filter_fns, errors)
+
+
+@gin.register
+class LevelRescaling(LinearTransform):
+  """Linear transform that scales values with vertical levels."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      scale: Sequence[float],
+      keys_to_scale: Sequence[str] = tuple(),
+  ):
+    super().__init__(trajectory_spec)
+    self.scale_fn = functools.partial(
+        coordinate_systems.scale_levels_for_matching_keys,
+        scales=np.asarray(scale),
+        keys_to_scale=keys_to_scale,
+    )
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    return self.scale_fn(errors)
+
+
+@gin.register
+class LevelRemoval(LinearTransform):
+  """Linear transform that removes vertical levels."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      keep_levels: Sequence[float],
+  ):
+    super().__init__(trajectory_spec)
+    n_levels = trajectory_spec.data_coords.vertical.layers
+    indices = jnp.array([i for i in range(n_levels) if keep_levels[i]])
+    self.take_arr = lambda x: jnp.take(x, indices, axis=metrics_util.LEVEL_AXIS)
+
+  def __call__(self, errors: Pytree, targets: Pytree) -> Pytree:
+    del targets  # unused.
+    return tree_map(self.take_arr, errors)

model/reference_code/metrics.py

@@ -0,0 +1,696 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Metrics and loss functions for NeuralGCM."""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+from typing import Callable, Optional, Sequence
+
+from dinosaur import coordinate_systems
+from dinosaur import horizontal_interpolation
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+from dinosaur import vertical_interpolation
+import gin
+import jax
+import jax.numpy as jnp
+import model.reference_code.linear_transforms as linear_transforms
+import model.reference_code.metrics_base as metrics_base
+import model.reference_code.metrics_util as metrics_util
+from model.legacy import model_utils
+import numpy as np
+import model.reference_code.train_utils as train_utils
+
+
+Pytree = typing.Pytree
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+
+tree_leaves = jax.tree_util.tree_leaves
+tree_map = jax.tree_util.tree_map
+
+
+def _compute_spectral_norm(
+    x: typing.Array, coords: coordinate_systems.CoordinateSystem
+) -> typing.Array:
+  """Computes spectral norm of nodal inputs `x`."""
+  x = coordinate_systems.maybe_to_modal(x, coords)
+  # axis = -2 corresponds to the longitudinal wavenumber.
+  return model_utils.safe_sqrt(
+      jnp.sum((x * x.conj()).real, axis=-2, keepdims=True)
+  )
+
+
+@gin.register
+def _spectral_amplitude(
+    x: typing.Array, coords: coordinate_systems.CoordinateSystem
+) -> typing.Array:
+  """Computes spectral amplitude ."""
+  x = coordinate_systems.maybe_to_modal(x, coords)
+  return jnp.abs(x)
+
+
+@gin.register
+@dataclasses.dataclass
+class TransformedL2Loss(metrics_base.Loss):
+  """L2 loss on linearly transformed errors."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      components: Sequence[linear_transforms.LinearTransformConstructor],
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      time_step: Optional[int | slice] = None,
+  ):
+    super().__init__(
+        trajectory_spec,
+        is_nodal=is_nodal,
+        is_encoded=is_encoded,
+        time_step=time_step,
+    )
+    self.components = components
+    self.getter = getter
+    self.transform = linear_transforms.ComposedTransformForLoss(
+        trajectory_spec, components
+    )
+
+  def evaluate_per_variable(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    prediction = self.get_representation(prediction)
+    target = self.get_representation(target)
+    trajectory = self.getter(prediction)
+    target = self.getter(target)
+    errors = tree_map(jnp.subtract, trajectory, target)
+    transformed_errors = self.transform(errors, target)
+    squared_transformed_errors = tree_map(jnp.square, transformed_errors)
+    return self.mean_per_variable(squared_transformed_errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class TransformedL2SpectrumLoss(metrics_base.Loss):
+  """L2 loss on linearly transformed errors of spectal norms.
+
+  Here we define spectrum norm at a given total wavenumber as the length of the
+  vector formed by longitude wavenumbers. i.e. for a field `x` with indices
+  `{z, m, l}` corresponding to level, longitude wavenumber, total wavenumber
+  we have:
+
+  spectrum_norm(x)_{z, l} = ||x_{z, :, l}||₂
+
+  The loss is then computed as MSE(spectrum_norm(x), spectrum_norm(y)) where
+  `x` and `y` are predicted and target signals in modal representation.
+  """
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      components: Sequence[linear_transforms.LinearTransformConstructor],
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      time_step: Optional[int | slice] = None,
+  ):
+    super().__init__(
+        trajectory_spec,
+        is_nodal=is_nodal,
+        is_encoded=is_encoded,
+        time_step=time_step,
+    )
+    if self.is_encoded:
+      coords = trajectory_spec.coords
+    else:
+      coords = trajectory_spec.data_coords
+    spectrum_fn = lambda x: _compute_spectral_norm(x, coords)
+    self.components = components
+    self.getter = getter
+    self.spectrum_fn = lambda tree: tree_map(spectrum_fn, tree)
+    self.transform = linear_transforms.ComposedTransformForLoss(
+        trajectory_spec, components
+    )
+
+  def mean_per_variable(self, trajectory: Pytree) -> Pytree:
+    return tree_map(jnp.mean, trajectory)
+
+  def evaluate_per_variable(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    prediction = self.get_representation(prediction)
+    target = self.get_representation(target)
+    trajectory_spectrum = self.spectrum_fn(self.getter(prediction))
+    target_spectrum = self.spectrum_fn(self.getter(target))
+    errors = tree_map(jnp.subtract, trajectory_spectrum, target_spectrum)
+    transformed_errors = self.transform(errors, target)
+    squared_transformed_errors = tree_map(jnp.square, transformed_errors)
+    return self.mean_per_variable(squared_transformed_errors)
+
+
+@gin.register
+@dataclasses.dataclass
+class SumLoss(metrics_base.Loss):
+  """Loss that consists of a sum of separate losses."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      terms: Sequence[Callable[..., metrics_base.Loss]],
+      labels: Optional[Sequence[str]] = None,
+      time_step: Optional[int | slice] = None,
+  ):
+    super().__init__(trajectory_spec)
+    self.losses = [term(trajectory_spec, time_step=time_step) for term in terms]
+    if labels is not None:
+      if len(labels) != len(self.losses):
+        raise ValueError(f'Not all losses are labeled: {labels}, {len(terms)=}')
+      self.labels = labels
+    else:
+      self.labels = [''] * len(self.losses)
+
+  def evaluate_per_variable(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    all_per_variable_losses = [
+        loss.evaluate_per_variable(prediction, target) for loss in self.losses
+    ]
+    output = {}
+    for per_variable_loss, prefix in zip(all_per_variable_losses, self.labels):
+      for k, v in per_variable_loss.items():
+        if isinstance(v, dict):
+          current_values = output.get(prefix + k, {})
+          for ik, iv in v.items():
+            current_values[ik] = current_values.get(ik, 0) + iv
+          output[prefix + k] = current_values
+        else:
+          output[prefix + k] = output.get(prefix + k, 0) + v
+    return output
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    return sum(loss.evaluate(prediction, target) for loss in self.losses)
+
+  def debug_loss_terms_instance(self) -> metrics_base.EvaluateFunctionWrapper:
+    """Returns class that evaluates relative loss per variable."""
+
+    def evaluate_fn(
+        prediction: TrajectoryRepresentations,
+        target: TrajectoryRepresentations,
+    ) -> Pytree:
+      return train_utils.flatten_dict({
+          label: loss.debug_loss_terms_instance().evaluate(prediction, target)
+          for label, loss in zip(self.labels, self.losses)
+      })
+
+    return metrics_base.EvaluateFunctionWrapper(evaluate_fn)
+
+
+@gin.register
+def WeightedL2CumulativeLoss(  # pylint: disable=invalid-name
+    trajectory_spec: metrics_util.TrajectorySpec,
+    weights: Pytree = None,
+    scale: float = 1.0,
+) -> TransformedL2Loss:
+  """Legacy wrapper for TransformedL2Loss with weighted cumulative error."""
+  components = [
+      linear_transforms.LegacyTimeRescaling,
+      functools.partial(
+          linear_transforms.PerVariableRescaling, weights=weights, scale=scale
+      ),
+  ]
+  return TransformedL2Loss(trajectory_spec, components)
+
+
+@gin.register
+class RMSE(metrics_base.ScalarMetric):
+  """Root mean squared error."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      time_step: int,
+      level: Optional[int] = None,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      is_ensemble_data: bool = False,
+  ):
+    super().__init__(trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded)
+    self.time_step = time_step
+    self.level = level
+    self.getter = getter
+    self.is_ensemble_data = is_ensemble_data
+
+  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
+    """Prepares target or predictions."""
+    trajectory = metrics_util.extract_variable(
+        trajectory,
+        self.trajectory_spec,
+        self.time_step,
+        self.level,
+        self.getter,
+        self.is_nodal,
+        self.is_encoded,
+    )
+    if self.is_ensemble_data:
+      # Evaluate RMSE vs. the ensemble mean.
+      trajectory = jax.lax.pmean(trajectory, axis_name='ensemble')
+    return trajectory
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> jnp.ndarray:
+    """Evaluates RMSE between prediction and target."""
+    prediction = self._prepare(prediction)
+    target = self._prepare(target)
+    squared_error = tree_map(lambda x, y: (x - y) ** 2, prediction, target)
+    mse_per_variable = self.mean_per_variable(squared_error)
+    return jnp.sqrt(sum(tree_leaves(mse_per_variable)))
+
+
+@gin.register
+class SpatialBiasRMSE(metrics_base.ScalarMetric):
+  """Root mean squared error of spatial bias.
+
+  This is given by the formula:
+
+    RMSE(batch_average(prediction - target))
+
+  where `batch_average()` denotes an average over distinct weather forecasts
+  (initialization times or valid times) and ensemble members (if relevant).
+  """
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      time_step: int,
+      level: Optional[int] = None,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      is_batch_data: bool = True,
+      is_ensemble_data: bool = False,
+  ):
+    super().__init__(trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded)
+    self.time_step = time_step
+    self.level = level
+    self.getter = getter
+    self.is_ensemble_data = is_ensemble_data
+    self.is_batch_data = is_batch_data
+
+  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
+    """Prepares target or predictions."""
+    trajectory = metrics_util.extract_variable(
+        trajectory,
+        self.trajectory_spec,
+        time_step=self.time_step,
+        level=self.level,
+        getter=self.getter,
+        is_nodal=self.is_nodal,
+        is_encoded=self.is_encoded,
+    )
+    if self.is_batch_data:
+      trajectory = jax.lax.pmean(trajectory, axis_name='batch')
+    if self.is_ensemble_data:
+      trajectory = jax.lax.pmean(trajectory, axis_name='ensemble')
+    return trajectory
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> jnp.ndarray:
+    """Evaluates RMSE between prediction and target."""
+    prediction = self._prepare(prediction)
+    target = self._prepare(target)
+    squared_error = tree_map(lambda x, y: (x - y) ** 2, prediction, target)
+    mse_per_variable = self.mean_per_variable(squared_error)
+    return jnp.sqrt(sum(tree_leaves(mse_per_variable)))
+
+
+@gin.register
+class BatchMeanSquaredBias(metrics_base.Loss):
+  """Mean squared error for a chosen metric.
+
+  This is given by the formula:
+
+    MSE(rollout_average(batch_average(prediction - target)))
+
+  where `batch_average()` denotes an average over distinct weather forecasts
+  (initialization times or valid times) or ensemble members (whichever is
+  vmapped first) and 'rollout_average()' denotes an average over all predicted
+  times. The MSE is taken over all nodal/modal points.
+  """
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      components: Sequence[linear_transforms.LinearTransformConstructor] = (),
+      observation_fn=_spectral_amplitude,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      is_nodal: bool = False,
+      is_encoded: bool = False,
+      time_step: Optional[int | slice] = None,
+  ):
+    super().__init__(
+        trajectory_spec,
+        is_nodal=is_nodal,
+        is_encoded=is_encoded,
+        time_step=time_step,
+    )
+    if self.is_encoded:
+      coords = trajectory_spec.coords
+    else:
+      coords = trajectory_spec.data_coords
+    metric_fn = lambda x: observation_fn(x, coords)
+    self.components = components
+    self.getter = getter
+    self.metric_fn = lambda tree: tree_map(metric_fn, tree)
+    self.transform = linear_transforms.ComposedTransformForLoss(
+        trajectory_spec, components
+    )
+
+  def evaluate_per_variable(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    """Evaluates the squere bias of a chosen metric between prediction and target.
+
+    Note: this method is only valid when vmapped.
+
+    Args:
+      prediction: a TrajectoryRepresentations of prediction
+      target: a TrajectoryRepresentations of ground truth
+
+    Returns:
+      Pytree of MSE
+    """
+    prediction = self.get_representation(prediction)
+    target = self.get_representation(target)
+    # because this function applies average over time axis, we apply
+    # `TruncateToTrajectoryLength` prior to computing
+    truncate_transform = self.transform.transforms[0]
+    assert isinstance(
+        truncate_transform, linear_transforms.TruncateToTrajectoryLength
+    )
+    getter_fn = lambda x: self.getter(truncate_transform(x, None))
+    trajectory_calc = self.metric_fn(getter_fn(prediction))
+    target_calc = self.metric_fn(getter_fn(target))
+    # Batch mean over "ensemble" and "batch" dimensions
+    trajectory_calc = tree_map(metrics_util.pmean_all_axes, trajectory_calc)
+    target_calc = tree_map(metrics_util.pmean_all_axes, target_calc)
+    # Time mean:
+    trajectory_calc = tree_map(
+        lambda x,: jnp.mean(x, axis=0, keepdims=True), trajectory_calc
+    )
+    target_calc = tree_map(
+        lambda x,: jnp.mean(x, axis=0, keepdims=True), target_calc
+    )
+    errors = tree_map(jnp.subtract, trajectory_calc, target_calc)
+    transformed_errors = self.transform(errors, target)
+    squared_transformed_errors = tree_map(jnp.square, transformed_errors)
+    mse_per_variable = tree_map(jnp.mean, squared_transformed_errors)
+    return mse_per_variable
+
+
+@gin.register
+class MAE(metrics_base.ScalarMetric):
+  """Mean absolute error."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      time_step: int,
+      level: Optional[int] = None,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+  ):
+    super().__init__(trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded)
+    self.time_step = time_step
+    self.level = level
+    self.getter = getter
+
+  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
+    return metrics_util.extract_variable(
+        trajectory,
+        self.trajectory_spec,
+        self.time_step,
+        self.level,
+        self.getter,
+        self.is_nodal,
+        self.is_encoded,
+    )
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> jnp.ndarray:
+    prediction = self._prepare(prediction)
+    target = self._prepare(target)
+    abs_error = tree_map(lambda x, y: abs(x - y), prediction, target)
+    mse_per_variable = self.mean_per_variable(abs_error)
+    flat_mse = tree_leaves(mse_per_variable)
+    return sum(flat_mse) / len(flat_mse)
+
+
+@jax.jit
+def weighted_quantile(
+    data: jax.Array, quantile: jax.Array, weights: jax.Array
+) -> jax.Array:
+  """Calculate a weighted quantile."""
+  if data.shape != weights.shape:
+    raise ValueError(f'incompatible shapes: {data.shape=} != {weights.shape=}')
+  data = data.ravel()
+  weights = weights.ravel() / weights.sum()
+  indices = jnp.argsort(data)
+  cum_weights = weights[indices].cumsum()
+  return jnp.interp(quantile, cum_weights, data[indices])
+
+
+@dataclasses.dataclass
+class AbsErrorQuantile(metrics_base.ScalarMetric):
+  """Quantile of absolute error."""
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      quantile: float,
+      time_step: int,
+      level: Optional[int] = None,
+      getter: Callable[[Pytree], Pytree] = metrics_util.filter_sim_time,
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      is_ensemble_data: bool = False,
+  ):
+    super().__init__(trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded)
+    self.quantile = quantile
+    self.time_step = time_step
+    self.level = level
+    self.getter = getter
+    self.is_ensemble_data = is_ensemble_data
+
+  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
+    return metrics_util.extract_variable(
+        trajectory,
+        self.trajectory_spec,
+        self.time_step,
+        self.level,
+        self.getter,
+        self.is_nodal,
+        self.is_encoded,
+    )
+
+  def _get_weights(self) -> np.ndarray:
+    if self.is_encoded:
+      coords = self.trajectory_spec.coords
+    else:
+      coords = self.trajectory_spec.data_coords
+    if self.is_nodal:
+      weights = coords.horizontal.quadrature_weights
+    else:
+      weights = coords.horizontal.mask
+    return weights
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> jnp.ndarray:
+    prediction = self._prepare(prediction)
+    target = self._prepare(target)
+    abs_error = tree_map(lambda x, y: abs(x - y), prediction, target)
+    weights = jnp.broadcast_to(self._get_weights(), target.shape)
+    result = tree_map(
+        lambda e: weighted_quantile(e, self.quantile, weights), abs_error
+    )
+    if self.is_ensemble_data:
+      # metrics must be consistent across the ensmble dimension.
+      result = jax.lax.pmean(result, axis_name='ensemble')
+    return result
+
+
+def weatherbench2_rmse_metrics(
+    trajectory_spec: metrics_util.TrajectorySpec,
+    time_steps: Sequence[int],
+    is_ensemble_data: bool = False,
+    extra_metric_grids: Optional[dict[str, spherical_harmonic.Grid]] = None,
+) -> dict[str, metrics_base.Metric]:
+  """RMSE based metrics for WeatherBench2."""
+  metric_grids = {} if extra_metric_grids is None else extra_metric_grids.copy()
+  trajectory_grid = trajectory_spec.coords.horizontal
+  if trajectory_grid not in metric_grids.values():
+    metric_grids['Traj'] = trajectory_grid
+
+  def get_and_regrid(tree, regrid_fn, getter):
+    return tree_map(regrid_fn, getter(tree))
+
+  metrics = {}
+  for name, grid in metric_grids.items():
+    if grid == trajectory_grid:
+      regrid = lambda tree: tree
+      rmse_traj_spec = trajectory_spec
+    else:
+      regrid = horizontal_interpolation.ConservativeRegridder(
+          source_grid=trajectory_spec.coords.horizontal, target_grid=grid
+      )
+      rmse_traj_spec = dataclasses.replace(
+          trajectory_spec,
+          # Only data_coords needs to be replaced since RMSE.is_encoded=False.
+          data_coords=dataclasses.replace(
+              trajectory_spec.data_coords,
+              horizontal=grid,
+          ),
+      )
+    for time_step in time_steps:
+      for var, level, getter in [
+          ('T', 850, lambda x: x['t']),
+          ('Z', 500, lambda x: x['z']),
+          ('UV', 700, lambda x: (x['u'], x['v'])),
+          ('Q', 700, lambda x: 1000 * x['tracers']['specific_humidity']),
+      ]:
+        t = time_step * trajectory_spec.steps_per_save
+        key = f'RMSE[{name}]_{var}{level}_{t:03d}_hours'
+        metrics[key] = RMSE(
+            rmse_traj_spec,
+            is_encoded=False,  # To make this (default) clear.
+            time_step=time_step,
+            level=level,
+            getter=functools.partial(
+                get_and_regrid, regrid_fn=regrid, getter=getter
+            ),
+            is_ensemble_data=is_ensemble_data,
+        )
+  return metrics
+
+
+def default_metrics(
+    trajectory_spec: metrics_util.TrajectorySpec,
+    eval_time_steps: Sequence[int],
+    train_loss: metrics_base.Loss,
+    is_batch_data: bool = True,
+    is_ensemble_data: bool = False,
+) -> dict[str, metrics_base.Evaluator]:
+  """Default evaluation metrics for Whirl models."""
+  metrics_dict = {
+      'training_loss': train_loss,
+      'debug': train_loss.debug_loss_terms_instance(),
+  }
+
+  if isinstance(
+      trajectory_spec.data_coords.vertical,
+      vertical_interpolation.PressureCoordinates,
+  ):
+    tl31_grid = dataclasses.replace(
+        spherical_harmonic.Grid.TL31(),
+        spherical_harmonics_impl=trajectory_spec.data_coords.horizontal.spherical_harmonics_impl,
+    )
+    metrics_dict.update(
+        weatherbench2_rmse_metrics(
+            trajectory_spec,
+            eval_time_steps,
+            is_ensemble_data=is_ensemble_data,
+            extra_metric_grids={'TL31': tl31_grid},
+        )
+    )
+
+    for time_step in eval_time_steps:
+      t = time_step * trajectory_spec.steps_per_save
+
+      for var, getter in [
+          ('T', lambda x: x['t']),
+          ('Z', lambda x: x['z']),
+          ('UV', lambda x: (x['u'], x['v'])),
+          ('Q', lambda x: 1000 * x['tracers']['specific_humidity']),
+      ]:
+        key = f'rmse_{var}_all_levels_{t:03d}_hours'
+        metrics_dict[key] = RMSE(
+            trajectory_spec,
+            time_step=time_step,
+            level=None,
+            getter=getter,
+            is_ensemble_data=is_ensemble_data,
+        )
+
+      for var, level, getter in [
+          ('T', 850, lambda x: x['t']),
+          ('Z', 500, lambda x: x['z']),
+          ('U', 700, lambda x: x['u']),
+          ('V', 700, lambda x: x['v']),
+          ('Q', 700, lambda x: 1000 * x['tracers']['specific_humidity']),
+      ]:
+        key = f'spatial_bias_rmse_{var}{level}_{t:03d}_hours'
+        metrics_dict[key] = SpatialBiasRMSE(
+            trajectory_spec,
+            time_step=time_step,
+            level=level,
+            getter=getter,
+            is_batch_data=is_batch_data,
+            is_ensemble_data=is_ensemble_data,
+        )
+
+      for var, level, getter in [
+          ('T', 850, lambda x: x['t']),
+          ('Z', 500, lambda x: x['z']),
+          ('U', 700, lambda x: x['u']),
+          ('V', 700, lambda x: x['v']),
+          ('Q', 700, lambda x: 1000 * x['tracers']['specific_humidity']),
+      ]:
+        for q in [0.99]:
+          key = f'abs_error_q{q}_{var}{level}_{t:03d}_hours'
+          metrics_dict[key] = AbsErrorQuantile(
+              trajectory_spec,
+              quantile=q,
+              time_step=time_step,
+              level=level,
+              getter=getter,
+              is_ensemble_data=is_ensemble_data,
+          )
+
+  return metrics_dict

model/reference_code/metrics_base.py

@@ -0,0 +1,150 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base classes for Metrics."""
+import dataclasses
+from typing import Callable
+from dinosaur import typing
+import jax
+import jax.numpy as jnp
+import model.reference_code.metrics_util as metrics_util
+
+
+Pytree = typing.Pytree
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+
+tree_leaves = jax.tree_util.tree_leaves
+tree_map = jax.tree_util.tree_map
+
+
+@dataclasses.dataclass
+class Evaluator:
+  """Class that evaluates on (prediction, trajectory) returning Pytree."""
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    """Evaluates giving values of interest."""
+    raise NotImplementedError()
+
+
+@dataclasses.dataclass
+class EvaluateFunctionWrapper(Evaluator):
+  """Wraps `evaluate_fn` function to be used as an Evaluator."""
+
+  def __init__(
+      self,
+      evaluate_fn: Callable[
+          [TrajectoryRepresentations, TrajectoryRepresentations], Pytree
+      ],
+  ):
+    self._evaluate_fn = evaluate_fn
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    return self._evaluate_fn(prediction, target)
+
+
+class MetricRuntimeError(Exception):
+  """Generic error for Metrics to raise in place of generic RuntimeError."""
+
+
+@dataclasses.dataclass
+class Metric(Evaluator):
+  """An Evaluator that derives information from a TrajectorySpec."""
+
+  trajectory_spec: metrics_util.TrajectorySpec
+  is_nodal: bool = dataclasses.field(default=True, kw_only=True)
+  is_encoded: bool = dataclasses.field(default=False, kw_only=True)
+
+  def get_representation(self, x: TrajectoryRepresentations) -> Pytree:
+    x_rep = x.get_representation(
+        is_nodal=self.is_nodal, is_encoded=self.is_encoded
+    )
+    if x_rep is None:
+      raise MetricRuntimeError(
+          'Desired representation of `x` was None. '
+          f'{self.is_nodal=}, {self.is_encoded=}'
+      )
+    return x_rep
+
+  def surface_mean(self, trajectory: Pytree) -> Pytree:
+    if self.is_encoded:
+      coords = self.trajectory_spec.coords
+    else:
+      coords = self.trajectory_spec.data_coords
+    if self.is_nodal:
+      # Mean over lat/lon. Converts shapes
+      #   (n_time, n_level, n_lon, n_lat) --> (n_time, n_level)
+      fn = lambda x: metrics_util.nodal_surface_mean(x, coords)
+    else:
+      fn = lambda x: metrics_util.modal_surface_mean(x, coords)
+    return tree_map(fn, trajectory)
+
+  def mean_per_variable(self, trajectory: Pytree) -> Pytree:
+    # In practice this is used to reduce shape (n_time, n_level) --> ()
+    return tree_map(jnp.mean, self.surface_mean(trajectory))
+
+
+class ScalarMetric(Metric):
+  """Metric that compute scalar quantities."""
+
+
+@dataclasses.dataclass
+class Loss(ScalarMetric):
+  """Metric that can be used as a loss."""
+
+  trajectory_spec: metrics_util.TrajectorySpec
+  is_nodal: bool = dataclasses.field(default=True, kw_only=True)
+  is_encoded: bool = dataclasses.field(default=False, kw_only=True)
+  time_step: int | slice | None = dataclasses.field(default=None, kw_only=True)
+
+  def evaluate_per_variable(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    raise NotImplementedError()
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> jnp.ndarray:
+    error_per_variable = self.evaluate_per_variable(prediction, target)
+    return sum(tree_leaves(error_per_variable))
+
+  def debug_loss_terms_instance(self) -> EvaluateFunctionWrapper:
+    """Returns class that evaluates relative loss per variable."""
+
+    def evaluate_fn(
+        prediction: TrajectoryRepresentations,
+        target: TrajectoryRepresentations,
+    ) -> Pytree:
+      # self.loss.evaluate takes ensemble mean (to evaluate on ensemble mean) if
+      # needed.
+      loss_per_variable = self.evaluate_per_variable(prediction, target)
+      # here we reduce terms by summation to expose relative contributions,
+      # even though the actual total_loss might be different.
+      sum_of_all_terms = sum(tree_leaves(loss_per_variable))
+      relative_loss = tree_map(
+          lambda x: x / sum_of_all_terms, loss_per_variable
+      )
+      return {'relative_loss': relative_loss}
+
+    return EvaluateFunctionWrapper(evaluate_fn)

model/reference_code/metrics_util.py

@@ -0,0 +1,453 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Shared utilities and classes for metrics and related modules."""
+from __future__ import annotations
+import dataclasses
+from types import MappingProxyType
+from typing import Callable, Optional, Sequence
+
+from dinosaur import coordinate_systems
+from dinosaur import horizontal_interpolation
+from dinosaur import pytree_utils
+from dinosaur import spherical_harmonic
+from dinosaur import typing
+import gin
+import jax
+import jax.extend as jex
+from jax.interpreters import ad
+from jax.interpreters import batching
+from jax.interpreters import mlir
+import jax.numpy as jnp
+import numpy as np
+
+
+tree_map = jax.tree_util.tree_map
+tree_leaves = jax.tree_util.tree_leaves
+Pytree = typing.Pytree
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+
+# Number of state variables in the model. t/z/u/v/specific_humidity.
+N_VARS = 5
+
+
+# Axis names.
+TIME = 'time'
+LEVEL = 'level'
+LONGITUDINAL_WAVENUMBER = 'longitudinal_wavenumber'
+TOTAL_WAVENUMBER = 'total_wavenumber'
+LONGITUDINAL = 'longitudinal'
+LATITUDINAL = 'latitudinal'
+
+# SPATIAL_AXES is negatively indexed because it is used in a place where there
+# are variable number of leading axis.
+SPATIAL_AXES = (-2, -1)
+TIME_AXIS = 0
+LEVEL_AXIS = 1
+ALL_AXES = (TIME_AXIS, LEVEL_AXIS) + SPATIAL_AXES
+
+
+MODAL_AXIS_INDICES = MappingProxyType({
+    TIME: TIME_AXIS,
+    LEVEL: LEVEL_AXIS,
+    LONGITUDINAL_WAVENUMBER: 2,
+    TOTAL_WAVENUMBER: 3,
+})
+
+
+NODAL_AXIS_INDICES = MappingProxyType({
+    TIME: TIME_AXIS,
+    LEVEL: LEVEL_AXIS,
+    LONGITUDINAL: 2,
+    LATITUDINAL: 3,
+})
+
+
+class ShapeError(Exception):
+  """Raised when an unexpected shape is encountered."""
+
+
+@dataclasses.dataclass
+class TrajectorySpec:
+  """Specification of a saved model trajectory."""
+
+  trajectory_length: int  # i.e., max "outer steps"
+  max_trajectory_length: int  # Maximum length for any stage of an Experiment.
+  steps_per_save: int  # Number of (1 hr) inner steps between each outer step.
+  coords: coordinate_systems.CoordinateSystem  # i.e., model coords
+  data_coords: coordinate_systems.CoordinateSystem  # i.e., data coords
+
+  def __post_init__(self):
+    if self.trajectory_length > self.max_trajectory_length:
+      raise ValueError(
+          f'{self.trajectory_length=} > {self.max_trajectory_length=}.'
+      )
+
+
+@dataclasses.dataclass
+class TrajectoryShape:
+  """Specifies shape of trajectory after LinearTransforms are applied."""
+
+  n_times: int
+  n_levels: int
+  n_longitudinal_wavenumbers: int
+  n_total_wavenumbers: int
+  n_longitude_nodes: int
+  n_latitude_nodes: int
+
+  def assert_compliant(self, trajectory: typing.Pytree, is_nodal: bool) -> None:
+    """Asserts `trajectory` is compliant with this `TrajectoryShape`.
+
+    Args:
+      trajectory: A trajectory, after LinearTransforms have been applied.
+      is_nodal: Whether the trajectory is presumed nodal (vs. modal).
+
+    Raises:
+      ShapeError: If the shape is not compliant.
+    """
+    if is_nodal:
+      expected_shape = (
+          self.n_times,
+          self.n_levels,
+          self.n_longitude_nodes,
+          self.n_latitude_nodes,
+      )
+    else:
+      expected_shape = (
+          self.n_times,
+          self.n_levels,
+          self.n_longitudinal_wavenumbers,
+          self.n_total_wavenumbers,
+      )
+
+    is_compliant = tree_map(lambda x: np.shape(x) == expected_shape, trajectory)
+    if not all(tree_leaves(is_compliant)):
+      shapes = tree_map(np.shape, trajectory)
+      raise ShapeError(
+          f'Some `trajectory` shapes were non-compliant ({is_nodal=}). '
+          f'{expected_shape=}. Found {shapes=}. '
+          f'This TrajectoryShape is {self}.'
+      )
+
+
+def nodal_surface_mean(
+    x: typing.Array, coords: coordinate_systems.CoordinateSystem
+) -> typing.Array:
+  """Integrates x over the surface of a sphere, normalized by surface area."""
+  if x.shape[-2:] != coords.horizontal.nodal_shape[-2:]:
+    raise ValueError(f'Input to nodal_surface_mean: {x.shape=}, while expected '
+                     f'spatial shape is {coords.horizontal.nodal_shape=}.')
+  surface_area = 4 * jnp.pi * coords.horizontal.radius**2
+  # Changes shape (n_t, n_z, n_lon, n_lat) --> (n_t, n_z)
+  return coords.horizontal.integrate(x) / surface_area
+
+
+def modal_surface_mean(
+    x: typing.Array, coords: coordinate_systems.CoordinateSystem
+) -> typing.Array:
+  """Integrates Σxₖφₖ² over a sphere, normalized by surface area."""
+  if x.shape[-2:] != coords.horizontal.modal_shape[-2:]:
+    raise ValueError(f'Input to modal_surface_mean: {x.shape=}, while expected '
+                     f'modal shape is {coords.horizontal.modal_shape=}.')
+  # This is equivalent to computing ||f||² / SurfaceArea, where
+  #   f = Σₖsqrt(x)ₖφₖ
+  surface_area = 4 * jnp.pi * coords.horizontal.radius**2
+
+  # Changes shape (n_t, n_z, m, l) --> (n_t, n_z)
+  return jnp.sum(x, axis=SPATIAL_AXES) / surface_area
+
+
+def extract_time_slice(trajectory: Pytree, time_slice: slice) -> Pytree:
+  return pytree_utils.slice_along_axis(trajectory, TIME_AXIS, time_slice)
+
+
+def extract_time_step(trajectory: Pytree, time_step: int) -> Pytree:
+  return extract_time_slice(trajectory, slice(time_step, time_step + 1))
+
+
+def extract_vertical_slice(
+    trajectory: Pytree,
+    coords: coordinate_systems.CoordinateSystem,
+    level: int,
+) -> Pytree:
+  i = coords.vertical.centers.tolist().index(level)
+  index = slice(i, i + 1)
+  trajectory = pytree_utils.slice_along_axis(trajectory, LEVEL_AXIS, index)
+  return trajectory
+
+
+def filter_sim_time(trajectory: Pytree) -> Pytree:
+  if isinstance(trajectory, dict):
+    trajectory = dict(trajectory)
+    trajectory.pop('sim_time', None)
+  return trajectory
+
+
+def filter_sim_time_and_diagnostics(trajectory: Pytree) -> Pytree:
+  if isinstance(trajectory, dict):
+    trajectory = dict(trajectory)
+    trajectory.pop('sim_time', None)
+    trajectory.pop('diagnostics', None)
+  return trajectory
+
+
+def extract_variable(
+    trajectory: TrajectoryRepresentations,
+    trajectory_spec: TrajectorySpec,
+    time_step: int | slice | None = None,
+    level: int | None = None,
+    getter: Callable[[Pytree], Pytree] = filter_sim_time,
+    is_nodal: bool = True,
+    is_encoded: bool = False,
+) -> Pytree:
+  """Extract a variable from a trajectory."""
+  if is_encoded:
+    coords = trajectory_spec.coords
+  else:
+    coords = trajectory_spec.data_coords
+  trajectory = trajectory.get_representation(
+      is_nodal=is_nodal, is_encoded=is_encoded
+  )
+  trajectory = getter(trajectory)
+  if time_step is not None:
+    if isinstance(time_step, slice):
+      trajectory = extract_time_slice(trajectory, time_step)
+    else:
+      trajectory = extract_time_step(trajectory, time_step)
+  if level is not None:
+    trajectory = extract_vertical_slice(trajectory, coords, level)
+  return trajectory
+
+
+def replace_with_linear_trucation(
+    trajectory_spec: TrajectorySpec,
+) -> TrajectorySpec:
+  """Replaces TrajectorySpec with a TL* version of it."""
+  grid = trajectory_spec.data_coords.horizontal
+  max_wavenumber = grid.longitude_wavenumbers - 1
+  assert max_wavenumber + 2 == grid.total_wavenumbers
+  gaussian_nodes = grid.longitude_nodes // 4
+  assert gaussian_nodes == grid.latitude_nodes // 2
+
+  # pytype: disable=attribute-error
+  new_horizontal = spherical_harmonic.Grid.construct(
+      max_wavenumber=2 * gaussian_nodes - 1,  # Larger in TL version
+      gaussian_nodes=gaussian_nodes,  # Same in T and TL versions
+      latitude_spacing=grid.latitude_spacing,
+      radius=grid.radius,
+  )
+  # pytype: enable=attribute-error
+
+  return dataclasses.replace(
+      trajectory_spec,
+      data_coords=dataclasses.replace(
+          trajectory_spec.data_coords,
+          horizontal=new_horizontal,
+      ),
+  )
+
+
+def trajectory_4d_shape(
+    trajectory_spec: TrajectorySpec,
+    keep_levels: Optional[Sequence[float]] = None,
+) -> TrajectoryShape:
+  """Returns the shape of the trajectory leaf values in data representation."""
+  if keep_levels is None:
+    n_levels = trajectory_spec.data_coords.vertical.layers
+  else:
+    n_levels = sum(bool(i) for i in keep_levels)
+    if n_levels > trajectory_spec.data_coords.vertical.layers:
+      raise ValueError(
+          f'{n_levels=} implied by `keep_levels` was greater than '
+          f'{trajectory_spec.data_coords.vertical.layers=}'
+      )
+  grid = trajectory_spec.data_coords.horizontal
+  n_m, n_l = grid.modal_shape
+  return TrajectoryShape(
+      n_times=trajectory_spec.trajectory_length,
+      n_levels=n_levels,
+      n_longitudinal_wavenumbers=n_m,
+      n_total_wavenumbers=n_l,
+      n_longitude_nodes=grid.longitude_nodes,
+      n_latitude_nodes=grid.latitude_nodes,
+  )
+
+
+def pmean_all_axes(x: jax.Array) -> jax.Array:
+  """Average over all vmapped axes."""
+  return _pmean_all_axes_p.bind(x)
+
+
+def _pmean_all_axes_impl(x):
+  return x
+
+
+def _pmean_all_axes_batch(args, batch_axes):
+  (x,) = args
+  (batch_axis,) = batch_axes
+  y = jnp.broadcast_to(x.mean(axis=batch_axes, keepdims=True), x.shape)
+  return _pmean_all_axes_p.bind(y), batch_axis
+
+
+_pmean_all_axes_p = jex.core.Primitive('pmean_all_axes')
+_pmean_all_axes_p.def_impl(_pmean_all_axes_impl)
+_pmean_all_axes_p.def_abstract_eval(_pmean_all_axes_impl)
+batching.primitive_batchers[_pmean_all_axes_p] = _pmean_all_axes_batch
+ad.deflinear(_pmean_all_axes_p, lambda cotangent: [pmean_all_axes(cotangent)])
+mlir.register_lowering(
+    _pmean_all_axes_p,
+    mlir.lower_fun(_pmean_all_axes_impl, multiple_results=False),
+)
+
+
+@dataclasses.dataclass
+class AggregationTransform:
+  """A transformation that aggregates spatial or temporal groups in inputs.
+
+  These transformations are useful for (1) coarsening of error observations and
+  (2) aggregation of error norms to compute L2^2 distance between two vectors.
+  The former case does not strictly impose any restrictions on the coarsening
+  transformation, although in most cases we would expect it to be a form of a
+  linear, non-invertible transformation. The latter requires that the result of
+  aggregation of non-negative values is non-negative.
+  """
+
+  trajectory_spec: TrajectorySpec
+  out_trajectory_spec: TrajectorySpec
+  is_nodal: bool
+  is_encoded: bool
+
+  def __call__(self, inputs: Pytree) -> Pytree:
+    raise NotImplementedError
+
+
+AggregationTransformConstructor = Callable[..., AggregationTransform]
+
+
+@gin.register
+class AggregateIdentity(AggregationTransform):
+
+  def __init__(
+      self,
+      trajectory_spec: TrajectorySpec,
+      is_nodal: bool,
+      is_encoded: bool,
+  ):
+    super().__init__(trajectory_spec, trajectory_spec, is_nodal, is_encoded)
+
+  def __call__(self, inputs: Pytree) -> Pytree:
+    return inputs
+
+
+@gin.register
+class SumVariables(AggregationTransform):
+  """Transform that adds sums all variables aka pytree leaves of inputs."""
+
+  def __init__(
+      self,
+      trajectory_spec: TrajectorySpec,
+      is_nodal: bool,
+      is_encoded: bool,
+  ):
+    super().__init__(trajectory_spec, trajectory_spec, is_nodal, is_encoded)
+
+  def __call__(self, inputs: Pytree) -> Pytree:
+    return sum(jax.tree_util.tree_leaves(inputs))
+
+
+@gin.register
+class RegriddingAggregation(AggregationTransform):
+  """Transform that aggregates horizontal cells via regridding.
+
+  To perform aggregation over a few nearby lon/lat cells this transform performs
+  regridding to a coarser `target_grid`. By default, the aggregated value
+  contains a regridded (i.e. mean) value of the inputs. Setting `scale_by_area`
+  to `True` multiplies outputs by an area which is close to area-weighted
+  aggregation.
+  """
+
+  def __init__(
+      self,
+      trajectory_spec: TrajectorySpec,
+      is_nodal: bool,
+      is_encoded: bool,
+      target_grid: coordinate_systems.CoordinateSystem,
+      scale_by_area: bool = False,
+  ):
+    if not is_nodal:
+      raise ValueError('AggregateHorizontal is only supported on nodal data')
+    if is_encoded:
+      source_coords = trajectory_spec.coords
+      coords = dataclasses.replace(source_coords, horizontal=target_grid)  # pytype: disable=wrong-arg-types  # dataclasses-replace-types
+      out_trajectory_spec = dataclasses.replace(trajectory_spec, coords=coords)
+    else:
+      source_coords = trajectory_spec.data_coords
+      coords = dataclasses.replace(source_coords, horizontal=target_grid)  # pytype: disable=wrong-arg-types  # dataclasses-replace-types
+      out_trajectory_spec = dataclasses.replace(
+          trajectory_spec, data_coords=coords)
+    super().__init__(trajectory_spec, out_trajectory_spec, is_nodal, is_encoded)
+    self.regrid_fn = horizontal_interpolation.ConservativeRegridder(
+        source_coords.horizontal, coords.horizontal)
+    # conservative regridding computes weighted averages rather than aggregation
+    # so we reweight the results by area.
+    lower_lon_boundaries = horizontal_interpolation._periodic_lower_bounds(
+        coords.horizontal.longitudes, 2 * np.pi)
+    upper_lon_boundaries = horizontal_interpolation._periodic_upper_bounds(
+        coords.horizontal.longitudes, 2 * np.pi)
+    lat_boundaries = horizontal_interpolation._latitude_cell_bounds(
+        coords.horizontal.latitudes)
+    lon_weights = upper_lon_boundaries - lower_lon_boundaries
+    lat_weights = jnp.sin(lat_boundaries[1:]) - jnp.sin(lat_boundaries[:-1])
+    self.weights = lat_weights[np.newaxis, :] * lon_weights[:, np.newaxis]
+    self.scale_by_area = scale_by_area
+
+  def __call__(self, inputs: Pytree) -> Pytree:
+    if self.scale_by_area:
+      return tree_map(lambda x: self.regrid_fn(x) * self.weights, inputs)
+    else:
+      return tree_map(self.regrid_fn, inputs)
+
+
+@gin.register
+class TimeWindowSum(AggregationTransform):
+  """Transform that sums temporal blocks of `time_window_size`."""
+
+  def __init__(
+      self,
+      trajectory_spec: TrajectorySpec,
+      is_nodal: bool,
+      is_encoded: bool,
+      time_window_size: int,
+  ):
+    trajectory_length = trajectory_spec.trajectory_length
+    if trajectory_length % time_window_size != 0:
+      raise ValueError(f'Cannot aggregate {trajectory_length=} '
+                       f'into {time_window_size=} sections.')
+    new_length = trajectory_spec.trajectory_length // time_window_size
+    out_trajectory_spec = dataclasses.replace(
+        trajectory_spec,
+        trajectory_length=new_length,
+        steps_per_save=trajectory_spec.steps_per_save * time_window_size)
+    super().__init__(trajectory_spec, out_trajectory_spec, is_nodal, is_encoded)
+    eye = np.eye(trajectory_length)
+    # columns of the weight matrix have 1s in rows that are in the same window.
+    # see http://screen/8CaZoBVNPtjpwNu for a hint.
+    self.time_axis_weights = sum(
+        [np.roll(eye, i, 0) for i in range(time_window_size)]
+    )[:, ::time_window_size]
+
+  def __call__(self, inputs: Pytree) -> Pytree:
+    def _aggregate_time(x: jax.Array):
+      return jnp.einsum(
+          'tk,...thml->...khml', self.time_axis_weights, x, precision='float32')
+    return tree_map(_aggregate_time, inputs)

model/reference_code/paper_configs/deterministic_0_7_deg.gin

@@ -0,0 +1,2368 @@
+# Macros:
+# ==============================================================================
+ACTIVATION = @gelu
+BASE_SHAPE_MULTIPLE = None
+CORRECTOR_MODULE = @CustomCoordsCorrector
+CORRECTOR_SCALE = 0.01
+DATA_FILTER_ATTENUATION = 0.0
+DYCORE_FILTER_ORDER = 3
+DYCORE_GRID = @GridWithWavenumbers()
+DYCORE_INTEGRATOR = @imex_rk_sil3
+DYCORE_TAU = '120 minutes'
+GLOBAL_OUT_SCALE = 0.01
+LATENT_SIZE = 384
+LAYER_SIZE = 384
+N_CNN_FEATURES = 32
+N_INNER_DYCORE_STEPS = 8
+N_SIGMA_LAYERS = 32
+N_TO_CLIP = 1
+NUM_BLOCKS = 5
+NUM_SUBSTEPS = 2
+PARAMETERIZATION_FILTER = @ml/SequentialStepFilter
+POSITIONAL_LATENT_SIZE = 32
+REVERSE_EINSUM_ARG_ORDER = None
+STABILITY_TAU = '4 minutes'
+SURFACE_MODEL_LATENT_SIZE = 8
+SURFACE_MODEL_LAYER_SIZE = 8
+SURFACE_MODEL_OUTPUT_SIZE = 8
+
+# Parameters for decode/ColumnTower:
+# ==============================================================================
+decode/ColumnTower.checkpoint_tower = False
+decode/ColumnTower.column_net_factory = @decode/MlpUniform
+decode/ColumnTower.name = 'decode_tower'
+
+# Parameters for encode/ColumnTower:
+# ==============================================================================
+encode/ColumnTower.checkpoint_tower = False
+encode/ColumnTower.column_net_factory = @encode/MlpUniform
+encode/ColumnTower.name = 'encode_tower'
+
+# Parameters for process/ColumnTower:
+# ==============================================================================
+process/ColumnTower.checkpoint_tower = False
+process/ColumnTower.column_net_factory = @process/MlpUniform
+process/ColumnTower.name = 'process_tower'
+
+# Parameters for surface_model_decode/ColumnTower:
+# ==============================================================================
+surface_model_decode/ColumnTower.checkpoint_tower = False
+surface_model_decode/ColumnTower.column_net_factory = \
+    @surface_model_decode/MlpUniform
+surface_model_decode/ColumnTower.name = 'surface_model_decode_tower'
+
+# Parameters for surface_model_encode/ColumnTower:
+# ==============================================================================
+surface_model_encode/ColumnTower.checkpoint_tower = False
+surface_model_encode/ColumnTower.column_net_factory = \
+    @surface_model_encode/MlpUniform
+surface_model_encode/ColumnTower.name = 'surface_model_encode_tower'
+
+# Parameters for surface_model_process/ColumnTower:
+# ==============================================================================
+surface_model_process/ColumnTower.checkpoint_tower = False
+surface_model_process/ColumnTower.column_net_factory = \
+    @surface_model_process/MlpUniform
+surface_model_process/ColumnTower.name = 'surface_model_process_tower'
+
+# Parameters for advance/CombinedFeatures:
+# ==============================================================================
+advance/CombinedFeatures.feature_module_names_to_exclude = ()
+advance/CombinedFeatures.feature_modules = \
+    (@EmbeddingSurfaceFeatures,
+     @EmbeddingVolumeFeatures,
+     @PressureFeatures,
+     @RadiationFeatures,
+     @LatitudeFeatures,
+     @advance/VelocityAndPrognostics,
+     @MemoryVelocityAndValues,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+advance/CombinedFeatures.features_to_exclude = ()
+advance/CombinedFeatures.features_transform_module = @advance/SequentialTransform
+advance/CombinedFeatures.name = None
+
+# Parameters for decoder_model/CombinedFeatures:
+# ==============================================================================
+decoder_model/CombinedFeatures.feature_module_names_to_exclude = ()
+decoder_model/CombinedFeatures.feature_modules = \
+    (@RadiationFeatures,
+     @LatitudeFeatures,
+     @model/VelocityAndPrognostics,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+decoder_model/CombinedFeatures.features_to_exclude = ()
+decoder_model/CombinedFeatures.features_transform_module = \
+    @decoder_model/SequentialTransform
+decoder_model/CombinedFeatures.name = None
+
+# Parameters for embedding_model/CombinedFeatures:
+# ==============================================================================
+embedding_model/CombinedFeatures.feature_module_names_to_exclude = ()
+embedding_model/CombinedFeatures.feature_modules = \
+    (@embedding_model/VelocityAndPrognostics, @PressureFeatures)
+embedding_model/CombinedFeatures.features_to_exclude = ()
+embedding_model/CombinedFeatures.features_transform_module = \
+    @embedding_model/ShiftAndNormalize
+embedding_model/CombinedFeatures.name = None
+
+# Parameters for encoder_data/CombinedFeatures:
+# ==============================================================================
+encoder_data/CombinedFeatures.feature_module_names_to_exclude = ()
+encoder_data/CombinedFeatures.feature_modules = \
+    (@RadiationFeatures,
+     @LatitudeFeatures,
+     @encoder_data/VelocityAndPrognostics,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+encoder_data/CombinedFeatures.features_to_exclude = ()
+encoder_data/CombinedFeatures.features_transform_module = \
+    @encoder_data/SequentialTransform
+encoder_data/CombinedFeatures.name = None
+
+# Parameters for sea_model/CombinedFeatures:
+# ==============================================================================
+sea_model/CombinedFeatures.feature_module_names_to_exclude = ()
+sea_model/CombinedFeatures.feature_modules = (@sea_model/ForcingFeatures,)
+sea_model/CombinedFeatures.features_to_exclude = ()
+sea_model/CombinedFeatures.features_transform_module = @sea_model/ShiftAndNormalize
+sea_model/CombinedFeatures.name = None
+
+# Parameters for coordinate_system_from_dataset:
+# ==============================================================================
+coordinate_system_from_dataset.spherical_harmonics_impl = \
+    @RealSphericalHarmonicsWithZeroImag
+coordinate_system_from_dataset.truncation = 'LINEAR'
+
+# Parameters for CoordinateSystem:
+# ==============================================================================
+CoordinateSystem.horizontal = @GridTL255()
+CoordinateSystem.vertical = @SigmaCoordinatesEquidistant()
+
+# Parameters for custom_corrds/CoordinateSystem:
+# ==============================================================================
+custom_corrds/CoordinateSystem.horizontal = %DYCORE_GRID
+custom_corrds/CoordinateSystem.vertical = @SigmaCoordinatesEquidistant()
+
+# Parameters for CustomCoordsCorrector:
+# ==============================================================================
+CustomCoordsCorrector.corrector_module = @DycoreWithPhysicsCorrector
+CustomCoordsCorrector.custom_coords = @custom_corrds/CoordinateSystem()
+CustomCoordsCorrector.name = None
+
+# Parameters for data_to_xarray_with_renaming:
+# ==============================================================================
+data_to_xarray_with_renaming.additional_coords = None
+data_to_xarray_with_renaming.attrs = None
+data_to_xarray_with_renaming.renaming_dict = \
+    {'geopotential': 'z',
+     'latitude': 'lat',
+     'longitude': 'lon',
+     'temperature': 't',
+     'u_component_of_wind': 'u',
+     'v_component_of_wind': 'v'}
+data_to_xarray_with_renaming.sample_ids = None
+data_to_xarray_with_renaming.to_xarray_fn = @primitive_eq_to_xarray
+
+# Parameters for divergence/DataExponentialFilter:
+# ==============================================================================
+divergence/DataExponentialFilter.attenuation = 14.0
+divergence/DataExponentialFilter.cutoff = 0.62
+divergence/DataExponentialFilter.name = None
+divergence/DataExponentialFilter.order = 2
+
+# Parameters for lsp/DataExponentialFilter:
+# ==============================================================================
+lsp/DataExponentialFilter.attenuation = 8
+lsp/DataExponentialFilter.cutoff = 0.82
+lsp/DataExponentialFilter.name = None
+lsp/DataExponentialFilter.order = 1
+
+# Parameters for orography/DataExponentialFilter:
+# ==============================================================================
+orography/DataExponentialFilter.attenuation = %DATA_FILTER_ATTENUATION
+orography/DataExponentialFilter.cutoff = 0
+orography/DataExponentialFilter.name = None
+orography/DataExponentialFilter.order = 1
+
+# Parameters for q/DataExponentialFilter:
+# ==============================================================================
+q/DataExponentialFilter.attenuation = 17
+q/DataExponentialFilter.cutoff = 0.62
+q/DataExponentialFilter.name = None
+q/DataExponentialFilter.order = 2
+
+# Parameters for temperature_variation/DataExponentialFilter:
+# ==============================================================================
+temperature_variation/DataExponentialFilter.attenuation = 13.0
+temperature_variation/DataExponentialFilter.cutoff = 0.45
+temperature_variation/DataExponentialFilter.name = None
+temperature_variation/DataExponentialFilter.order = 3
+
+# Parameters for vorticity/DataExponentialFilter:
+# ==============================================================================
+vorticity/DataExponentialFilter.attenuation = 14
+vorticity/DataExponentialFilter.cutoff = 0.62
+vorticity/DataExponentialFilter.name = None
+vorticity/DataExponentialFilter.order = 2
+
+# Parameters for DataNoFilter:
+# ==============================================================================
+DataNoFilter.name = None
+
+# Parameters for DimensionalLearnedPrimitiveToWeatherbenchDecoder:
+# ==============================================================================
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.correction_transform_module = \
+    @decoder/SequentialTransform
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.inputs_to_units_mapping = \
+    {'sim_time': 'dimensionless',
+     't': 'kelvin',
+     'tracers': {'specific_cloud_ice_water_content': 'dimensionless',
+                 'specific_cloud_liquid_water_content': 'dimensionless',
+                 'specific_humidity': 'dimensionless'},
+     'u': 'meter / second',
+     'v': 'meter / second',
+     'z': 'm**2 s**-2'}
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.modal_to_nodal_data_features_module = \
+    @NullFeatures
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.modal_to_nodal_model_features_module = \
+    @decoder_model/CombinedFeatures
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.name = None
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.nodal_mapping_module = \
+    @NodalMapping
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.orography_module = \
+    @LearnedOrography
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.prediction_mask = \
+    {'sim_time': False,
+     't': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'u': True,
+     'v': True,
+     'z': True}
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.time_axis = 0
+
+# Parameters for DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder:
+# ==============================================================================
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.correction_transform_module = \
+    @encode/SequentialTransform
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.inputs_to_units_mapping = \
+    {'sim_time': 'dimensionless',
+     't': 'kelvin',
+     'tracers': {'specific_cloud_ice_water_content': 'dimensionless',
+                 'specific_cloud_liquid_water_content': 'dimensionless',
+                 'specific_humidity': 'dimensionless'},
+     'u': 'meter / second',
+     'v': 'meter / second',
+     'z': 'm**2 s**-2'}
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.modal_to_nodal_data_features_module = \
+    @encoder_data/CombinedFeatures
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.modal_to_nodal_model_features_module = \
+    @NullFeatures
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.name = None
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.nodal_mapping_module = \
+    @NodalMapping
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.orography_module = \
+    @LearnedOrography
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.prediction_mask = \
+    {'divergence': True,
+     'log_surface_pressure': True,
+     'sim_time': False,
+     'temperature_variation': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'vorticity': True}
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.time_axis = 0
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.transform_module = \
+    @EncoderCombinedTransform
+
+# Parameters for DivCurlNeuralParameterization:
+# ==============================================================================
+DivCurlNeuralParameterization.filter_module = %PARAMETERIZATION_FILTER
+DivCurlNeuralParameterization.modal_to_nodal_features_module = \
+    @advance/CombinedFeatures
+DivCurlNeuralParameterization.name = None
+DivCurlNeuralParameterization.nodal_mapping_module = @NodalMapping
+DivCurlNeuralParameterization.prediction_mask = \
+    {'divergence': True,
+     'log_surface_pressure': False,
+     'sim_time': False,
+     'temperature_variation': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'vorticity': True}
+DivCurlNeuralParameterization.tendency_transform_module = \
+    @div_curl_tendency_outputs/SequentialTransform
+
+# Parameters for DycoreWithPhysicsCorrector:
+# ==============================================================================
+DycoreWithPhysicsCorrector.checkpoint_explicit_terms = True
+DycoreWithPhysicsCorrector.dycore_equation_module = \
+    @MoistPrimitiveEquationsWithCloudMoisture
+DycoreWithPhysicsCorrector.dycore_substeps = %N_INNER_DYCORE_STEPS
+DycoreWithPhysicsCorrector.filter_module = @dycore/SequentialStepFilter
+DycoreWithPhysicsCorrector.name = None
+DycoreWithPhysicsCorrector.time_integrator = %DYCORE_INTEGRATOR
+
+# Parameters for DynamicDataForcing:
+# ==============================================================================
+DynamicDataForcing.check_sim_time_errors = False
+DynamicDataForcing.data_time_step = '6 hours'
+DynamicDataForcing.dt_tolerance = '1 year'
+DynamicDataForcing.inputs_to_units_mapping = \
+    {'sea_ice_cover': 'dimensionless',
+     'sea_surface_temperature': 'kelvin',
+     'sim_time': 'dimensionless'}
+DynamicDataForcing.name = None
+DynamicDataForcing.time_axis = 0
+
+# Parameters for advance/EmbeddingSurfaceFeatures:
+# ==============================================================================
+advance/EmbeddingSurfaceFeatures.embedding_module = @NodalLandSeaIceEmbedding
+advance/EmbeddingSurfaceFeatures.feature_name = 'surface_embedding'
+advance/EmbeddingSurfaceFeatures.name = None
+advance/EmbeddingSurfaceFeatures.output_size = %SURFACE_MODEL_OUTPUT_SIZE
+
+# Parameters for advance/EmbeddingVolumeFeatures:
+# ==============================================================================
+advance/EmbeddingVolumeFeatures.embedding_module = @ModalToNodalEmbedding
+advance/EmbeddingVolumeFeatures.feature_name = 'CNN1D'
+advance/EmbeddingVolumeFeatures.name = None
+advance/EmbeddingVolumeFeatures.output_size = %N_CNN_FEATURES
+
+# Parameters for EncoderCombinedTransform:
+# ==============================================================================
+EncoderCombinedTransform.name = None
+EncoderCombinedTransform.transforms = \
+    (@InputClipTransform, @EncoderFilterTransform)
+
+# Parameters for EncoderFilterTransform:
+# ==============================================================================
+EncoderFilterTransform.filter_modules = (@PerVariableDataFilter,)
+EncoderFilterTransform.name = None
+
+# Parameters for EpdTower:
+# ==============================================================================
+EpdTower.decode_tower_factory = @decode/ColumnTower
+EpdTower.encode_tower_factory = @encode/ColumnTower
+EpdTower.final_activation = None
+EpdTower.latent_size = %LATENT_SIZE
+EpdTower.name = None
+EpdTower.num_process_blocks = %NUM_BLOCKS
+EpdTower.post_encode_activation = None
+EpdTower.pre_decode_activation = None
+EpdTower.process_tower_factory = @process/ColumnTower
+
+# Parameters for surface_model/EpdTower:
+# ==============================================================================
+surface_model/EpdTower.decode_tower_factory = @surface_model_decode/ColumnTower
+surface_model/EpdTower.encode_tower_factory = @surface_model_encode/ColumnTower
+surface_model/EpdTower.final_activation = None
+surface_model/EpdTower.latent_size = %SURFACE_MODEL_LATENT_SIZE
+surface_model/EpdTower.name = None
+surface_model/EpdTower.num_process_blocks = 1
+surface_model/EpdTower.post_encode_activation = None
+surface_model/EpdTower.pre_decode_activation = None
+surface_model/EpdTower.process_tower_factory = @surface_model_process/ColumnTower
+
+# Parameters for dycore/ExponentialFilter:
+# ==============================================================================
+dycore/ExponentialFilter.cutoff = 0
+dycore/ExponentialFilter.name = None
+dycore/ExponentialFilter.order = %DYCORE_FILTER_ORDER
+dycore/ExponentialFilter.tau = %DYCORE_TAU
+
+# Parameters for stability/ExponentialFilter:
+# ==============================================================================
+stability/ExponentialFilter.cutoff = 0.4
+stability/ExponentialFilter.name = None
+stability/ExponentialFilter.order = 6
+stability/ExponentialFilter.tau = %STABILITY_TAU
+
+# Parameters for FilteredCustomOrography:
+# ==============================================================================
+FilteredCustomOrography.filter_modules = (@orography/DataExponentialFilter,)
+FilteredCustomOrography.name = None
+FilteredCustomOrography.orography_data_path = None
+FilteredCustomOrography.renaming_dict = {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for with_grads/FloatDataFeatures:
+# ==============================================================================
+with_grads/FloatDataFeatures.compute_gradients_module = @ToModalDiffOperators
+with_grads/FloatDataFeatures.covariate_data_path = None
+with_grads/FloatDataFeatures.covariate_keys = ('geopotential_at_surface',)
+with_grads/FloatDataFeatures.name = None
+with_grads/FloatDataFeatures.renaming_dict = \
+    {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for without_grads/FloatDataFeatures:
+# ==============================================================================
+without_grads/FloatDataFeatures.covariate_data_path = None
+without_grads/FloatDataFeatures.covariate_keys = ('land_sea_mask',)
+without_grads/FloatDataFeatures.name = None
+without_grads/FloatDataFeatures.renaming_dict = \
+    {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for sea_model/ForcingFeatures:
+# ==============================================================================
+sea_model/ForcingFeatures.forcing_to_include = ('sea_surface_temperature',)
+sea_model/ForcingFeatures.name = None
+
+# Parameters for gelu:
+# ==============================================================================
+gelu.approximate = True
+
+# Parameters for GET_ATMOSPHERIC_SCALE:
+# ==============================================================================
+# None.
+
+# Parameters for get_model_specs:
+# ==============================================================================
+get_model_specs.custom_coords = @CoordinateSystem()
+get_model_specs.model_time_step = '1 hour'
+get_model_specs.reference_datetime_str = None
+get_model_specs.reference_temperature = \
+    [215.58614815,
+     211.47405876,
+     205.87815406,
+     206.40755302,
+     210.43452345,
+     214.5683887,
+     218.75303863,
+     223.23145107,
+     227.9710687,
+     232.85381503,
+     237.53588735,
+     242.05068293,
+     246.29986585,
+     250.14294113,
+     253.74839535,
+     256.98024283,
+     259.94441031,
+     262.7041158,
+     265.21752838,
+     267.62333985,
+     269.94462121,
+     272.10056439,
+     274.12518288,
+     275.99833711,
+     277.72759392,
+     279.3292128,
+     280.79178708,
+     282.13507065,
+     283.41832023,
+     284.7682506,
+     286.33945487,
+     288.06707666]
+
+# Parameters for get_physics_specs:
+# ==============================================================================
+get_physics_specs.construct_fn = @primitive_eq_specs_constructor
+
+# Parameters for GridTL255:
+# ==============================================================================
+GridTL255.spherical_harmonics_impl = @RealSphericalHarmonicsWithZeroImag
+
+# Parameters for DYCORE_GRID/GridWithWavenumbers:
+# ==============================================================================
+DYCORE_GRID/GridWithWavenumbers.dealiasing = 'quadratic'
+DYCORE_GRID/GridWithWavenumbers.latitude_spacing = 'gauss'
+DYCORE_GRID/GridWithWavenumbers.longitude_offset = 0.0
+DYCORE_GRID/GridWithWavenumbers.longitude_wavenumbers = 254
+DYCORE_GRID/GridWithWavenumbers.radius = None
+DYCORE_GRID/GridWithWavenumbers.spherical_harmonics_impl = \
+    @RealSphericalHarmonicsWithZeroImag
+
+# Parameters for advance/IdentityTransform:
+# ==============================================================================
+advance/IdentityTransform.name = None
+
+# Parameters for sea_model/IdentityTransform:
+# ==============================================================================
+sea_model/IdentityTransform.name = None
+
+# Parameters for imex_rk_sil3:
+# ==============================================================================
+# None.
+
+# Parameters for InputClipTransform:
+# ==============================================================================
+InputClipTransform.name = None
+InputClipTransform.wavenumbers_to_clip = %N_TO_CLIP
+
+# Parameters for advance/InverseLevelScale:
+# ==============================================================================
+advance/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+advance/InverseLevelScale.name = None
+advance/InverseLevelScale.scales = \
+    [8.822e-05,
+     7.126e-05,
+     0.0001047,
+     0.0001858,
+     0.0007601,
+     0.002642,
+     0.007242,
+     0.01568,
+     0.02907,
+     0.04654,
+     0.07084,
+     0.09971,
+     0.1355,
+     0.1762,
+     0.2243,
+     0.2822,
+     0.3459,
+     0.4172,
+     0.471,
+     0.5286,
+     0.5932,
+     0.6682,
+     0.7546,
+     0.8532,
+     0.9553,
+     1.058,
+     1.164,
+     1.278,
+     1.401,
+     1.553,
+     1.709,
+     1.791]
+
+# Parameters for decoder_model/InverseLevelScale:
+# ==============================================================================
+decoder_model/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+decoder_model/InverseLevelScale.name = None
+decoder_model/InverseLevelScale.scales = \
+    [8.822e-05,
+     7.126e-05,
+     0.0001047,
+     0.0001858,
+     0.0007601,
+     0.002642,
+     0.007242,
+     0.01568,
+     0.02907,
+     0.04654,
+     0.07084,
+     0.09971,
+     0.1355,
+     0.1762,
+     0.2243,
+     0.2822,
+     0.3459,
+     0.4172,
+     0.471,
+     0.5286,
+     0.5932,
+     0.6682,
+     0.7546,
+     0.8532,
+     0.9553,
+     1.058,
+     1.164,
+     1.278,
+     1.401,
+     1.553,
+     1.709,
+     1.791]
+
+# Parameters for encoder_data/InverseLevelScale:
+# ==============================================================================
+encoder_data/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+encoder_data/InverseLevelScale.name = None
+encoder_data/InverseLevelScale.scales = \
+    [3.823e-05,
+     5.985e-05,
+     7.103e-05,
+     8.198e-05,
+     8.836e-05,
+     9.113e-05,
+     7.745e-05,
+     7.234e-05,
+     8.308e-05,
+     9.885e-05,
+     0.0001443,
+     0.0002999,
+     0.001019,
+     0.002773,
+     0.006194,
+     0.01195,
+     0.02037,
+     0.04634,
+     0.08648,
+     0.1406,
+     0.2104,
+     0.3017,
+     0.4097,
+     0.4949,
+     0.5891,
+     0.716,
+     0.8645,
+     0.9418,
+     1.019,
+     1.098,
+     1.178,
+     1.262,
+     1.35,
+     1.454,
+     1.581,
+     1.677,
+     1.713]
+
+# Parameters for decoder/InverseShiftAndNormalize:
+# ==============================================================================
+decoder/InverseShiftAndNormalize.global_scale = %CORRECTOR_SCALE
+decoder/InverseShiftAndNormalize.name = None
+decoder/InverseShiftAndNormalize.scales = \
+    {'sim_time': 0.0,
+     't': 27.99,
+     'tracers': {'specific_cloud_ice_water_content': 8.255e-06,
+                 'specific_cloud_liquid_water_content': 2.182e-05,
+                 'specific_humidity': 0.003493},
+     'u': 0.01935,
+     'v': 0.01038,
+     'z': 0.1496}
+decoder/InverseShiftAndNormalize.shifts = \
+    {'sim_time': 0.0,
+     't': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'u': 0.0,
+     'v': 0.0,
+     'z': 0.0}
+
+# Parameters for div_curl_tendency_outputs/InverseShiftAndNormalize:
+# ==============================================================================
+div_curl_tendency_outputs/InverseShiftAndNormalize.global_scale = %GLOBAL_OUT_SCALE
+div_curl_tendency_outputs/InverseShiftAndNormalize.name = None
+div_curl_tendency_outputs/InverseShiftAndNormalize.scales = \
+    {'log_surface_pressure': 0.05021,
+     'sim_time': 0.0,
+     'temperature_variation': 33.94,
+     'tracers': {'specific_cloud_ice_water_content': 4.85e-05,
+                 'specific_cloud_liquid_water_content': 9.693e-05,
+                 'specific_humidity': 0.006182},
+     'u': 0.05863,
+     'v': 0.0516}
+div_curl_tendency_outputs/InverseShiftAndNormalize.shifts = \
+    {'log_surface_pressure': 0.0,
+     'sim_time': 0.0,
+     'temperature_variation': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'u': 0.0,
+     'v': 0.0}
+
+# Parameters for encoder/InverseShiftAndNormalize:
+# ==============================================================================
+encoder/InverseShiftAndNormalize.global_scale = %CORRECTOR_SCALE
+encoder/InverseShiftAndNormalize.name = None
+encoder/InverseShiftAndNormalize.scales = \
+    {'divergence': 0.1078,
+     'log_surface_pressure': 0.1134,
+     'sim_time': 0.0,
+     'temperature_variation': 15.04,
+     'tracers': {'specific_cloud_ice_water_content': 9.812e-06,
+                 'specific_cloud_liquid_water_content': 2.039e-05,
+                 'specific_humidity': 0.003305},
+     'vorticity': 0.2831}
+encoder/InverseShiftAndNormalize.shifts = \
+    {'divergence': 0.0,
+     'log_surface_pressure': 0.0,
+     'sim_time': 0.0,
+     'temperature_variation': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'vorticity': 0.0}
+
+# Parameters for advance/LatitudeFeatures:
+# ==============================================================================
+advance/LatitudeFeatures.name = None
+
+# Parameters for decoder_model/LatitudeFeatures:
+# ==============================================================================
+decoder_model/LatitudeFeatures.name = None
+
+# Parameters for encoder_data/LatitudeFeatures:
+# ==============================================================================
+encoder_data/LatitudeFeatures.name = None
+
+# Parameters for LearnedOrography:
+# ==============================================================================
+LearnedOrography.base_orography_module = @FilteredCustomOrography
+LearnedOrography.correction_scale = 2e-06
+LearnedOrography.name = None
+
+# Parameters for advance/LearnedPositionalFeatures:
+# ==============================================================================
+advance/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+advance/LearnedPositionalFeatures.name = None
+advance/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for decoder_model/LearnedPositionalFeatures:
+# ==============================================================================
+decoder_model/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+decoder_model/LearnedPositionalFeatures.name = None
+decoder_model/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for encoder_data/LearnedPositionalFeatures:
+# ==============================================================================
+encoder_data/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+encoder_data/LearnedPositionalFeatures.name = None
+encoder_data/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for decoder/LevelScale:
+# ==============================================================================
+decoder/LevelScale.keys_to_scale = ['specific_humidity']
+decoder/LevelScale.name = None
+decoder/LevelScale.scales = \
+    [3.823e-05,
+     5.985e-05,
+     7.103e-05,
+     8.198e-05,
+     8.836e-05,
+     9.113e-05,
+     7.745e-05,
+     7.234e-05,
+     8.308e-05,
+     9.885e-05,
+     0.0001443,
+     0.0002999,
+     0.001019,
+     0.002773,
+     0.006194,
+     0.01195,
+     0.02037,
+     0.04634,
+     0.08648,
+     0.1406,
+     0.2104,
+     0.3017,
+     0.4097,
+     0.4949,
+     0.5891,
+     0.716,
+     0.8645,
+     0.9418,
+     1.019,
+     1.098,
+     1.178,
+     1.262,
+     1.35,
+     1.454,
+     1.581,
+     1.677,
+     1.713]
+
+# Parameters for div_curl_tendency_outputs/LevelScale:
+# ==============================================================================
+div_curl_tendency_outputs/LevelScale.keys_to_scale = ['specific_humidity']
+div_curl_tendency_outputs/LevelScale.name = None
+div_curl_tendency_outputs/LevelScale.scales = \
+    [0.0001303,
+     0.0002029,
+     0.0002682,
+     0.0004815,
+     0.001437,
+     0.004719,
+     0.01328,
+     0.03024,
+     0.05782,
+     0.09671,
+     0.1485,
+     0.2107,
+     0.2892,
+     0.3723,
+     0.4705,
+     0.5785,
+     0.6885,
+     0.8155,
+     0.9028,
+     0.9944,
+     1.103,
+     1.197,
+     1.31,
+     1.431,
+     1.526,
+     1.58,
+     1.61,
+     1.624,
+     1.634,
+     1.649,
+     1.729,
+     1.824]
+
+# Parameters for encode/LevelScale:
+# ==============================================================================
+encode/LevelScale.keys_to_scale = ['specific_humidity']
+encode/LevelScale.name = None
+encode/LevelScale.scales = \
+    [8.822e-05,
+     7.126e-05,
+     0.0001047,
+     0.0001858,
+     0.0007601,
+     0.002642,
+     0.007242,
+     0.01568,
+     0.02907,
+     0.04654,
+     0.07084,
+     0.09971,
+     0.1355,
+     0.1762,
+     0.2243,
+     0.2822,
+     0.3459,
+     0.4172,
+     0.471,
+     0.5286,
+     0.5932,
+     0.6682,
+     0.7546,
+     0.8532,
+     0.9553,
+     1.058,
+     1.164,
+     1.278,
+     1.401,
+     1.553,
+     1.709,
+     1.791]
+
+# Parameters for advance/MemoryVelocityAndValues:
+# ==============================================================================
+advance/MemoryVelocityAndValues.fields_to_include = None
+advance/MemoryVelocityAndValues.name = None
+
+# Parameters for decode/MlpUniform:
+# ==============================================================================
+decode/MlpUniform.activate_final = False
+decode/MlpUniform.activation = %ACTIVATION
+decode/MlpUniform.b_init = None
+decode/MlpUniform.b_init_final = None
+decode/MlpUniform.name = None
+decode/MlpUniform.num_hidden_layers = 0
+decode/MlpUniform.num_hidden_units = %LAYER_SIZE
+decode/MlpUniform.w_init = None
+decode/MlpUniform.w_init_final = None
+decode/MlpUniform.with_bias = False
+
+# Parameters for encode/MlpUniform:
+# ==============================================================================
+encode/MlpUniform.activate_final = False
+encode/MlpUniform.activation = %ACTIVATION
+encode/MlpUniform.b_init = None
+encode/MlpUniform.b_init_final = None
+encode/MlpUniform.name = None
+encode/MlpUniform.num_hidden_layers = 0
+encode/MlpUniform.num_hidden_units = 0
+encode/MlpUniform.w_init = None
+encode/MlpUniform.w_init_final = None
+encode/MlpUniform.with_bias = True
+
+# Parameters for process/MlpUniform:
+# ==============================================================================
+process/MlpUniform.activate_final = False
+process/MlpUniform.activation = %ACTIVATION
+process/MlpUniform.b_init = None
+process/MlpUniform.b_init_final = None
+process/MlpUniform.name = None
+process/MlpUniform.num_hidden_layers = 3
+process/MlpUniform.num_hidden_units = %LAYER_SIZE
+process/MlpUniform.w_init = None
+process/MlpUniform.w_init_final = None
+process/MlpUniform.with_bias = True
+
+# Parameters for surface_model_decode/MlpUniform:
+# ==============================================================================
+surface_model_decode/MlpUniform.activate_final = False
+surface_model_decode/MlpUniform.activation = %ACTIVATION
+surface_model_decode/MlpUniform.b_init = None
+surface_model_decode/MlpUniform.b_init_final = None
+surface_model_decode/MlpUniform.name = None
+surface_model_decode/MlpUniform.num_hidden_layers = 1
+surface_model_decode/MlpUniform.num_hidden_units = %SURFACE_MODEL_LAYER_SIZE
+surface_model_decode/MlpUniform.w_init = None
+surface_model_decode/MlpUniform.w_init_final = None
+surface_model_decode/MlpUniform.with_bias = False
+
+# Parameters for surface_model_encode/MlpUniform:
+# ==============================================================================
+surface_model_encode/MlpUniform.activate_final = False
+surface_model_encode/MlpUniform.activation = %ACTIVATION
+surface_model_encode/MlpUniform.b_init = None
+surface_model_encode/MlpUniform.b_init_final = None
+surface_model_encode/MlpUniform.name = None
+surface_model_encode/MlpUniform.num_hidden_layers = 0
+surface_model_encode/MlpUniform.num_hidden_units = 0
+surface_model_encode/MlpUniform.w_init = None
+surface_model_encode/MlpUniform.w_init_final = None
+surface_model_encode/MlpUniform.with_bias = True
+
+# Parameters for surface_model_process/MlpUniform:
+# ==============================================================================
+surface_model_process/MlpUniform.activate_final = False
+surface_model_process/MlpUniform.activation = %ACTIVATION
+surface_model_process/MlpUniform.b_init = None
+surface_model_process/MlpUniform.b_init_final = None
+surface_model_process/MlpUniform.name = None
+surface_model_process/MlpUniform.num_hidden_layers = 3
+surface_model_process/MlpUniform.num_hidden_units = %SURFACE_MODEL_LAYER_SIZE
+surface_model_process/MlpUniform.w_init = None
+surface_model_process/MlpUniform.w_init_final = None
+surface_model_process/MlpUniform.with_bias = True
+
+# Parameters for advance/ModalToNodalEmbedding:
+# ==============================================================================
+advance/ModalToNodalEmbedding.modal_to_nodal_features_module = \
+    @embedding_model/CombinedFeatures
+advance/ModalToNodalEmbedding.name = None
+advance/ModalToNodalEmbedding.nodal_mapping_module = @NodalVolumeMapping
+advance/ModalToNodalEmbedding.output_transform_module = @IdentityTransform
+
+# Parameters for sea_model/ModalToNodalEmbedding:
+# ==============================================================================
+sea_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \
+    @sea_model/CombinedFeatures
+sea_model/ModalToNodalEmbedding.name = None
+sea_model/ModalToNodalEmbedding.nodal_mapping_module = @sea_model/NodalMapping
+sea_model/ModalToNodalEmbedding.output_transform_module = @IdentityTransform
+
+# Parameters for MoistPrimitiveEquationsWithCloudMoisture:
+# ==============================================================================
+MoistPrimitiveEquationsWithCloudMoisture.include_vertical_advection = True
+MoistPrimitiveEquationsWithCloudMoisture.name = None
+MoistPrimitiveEquationsWithCloudMoisture.orography_module = @LearnedOrography
+
+# Parameters for advance/NodalLandSeaIceEmbedding:
+# ==============================================================================
+advance/NodalLandSeaIceEmbedding.land_embedding = None
+advance/NodalLandSeaIceEmbedding.name = None
+advance/NodalLandSeaIceEmbedding.sea_embedding = @sea_model/ModalToNodalEmbedding
+advance/NodalLandSeaIceEmbedding.sea_ice_embedding = None
+advance/NodalLandSeaIceEmbedding.static_vars_ds_path = None
+
+# Parameters for NodalMapping:
+# ==============================================================================
+NodalMapping.name = None
+NodalMapping.tower_factory = @EpdTower
+
+# Parameters for sea_model/NodalMapping:
+# ==============================================================================
+sea_model/NodalMapping.name = None
+sea_model/NodalMapping.tower_factory = @surface_model/EpdTower
+
+# Parameters for NodalVolumeMapping:
+# ==============================================================================
+NodalVolumeMapping.name = None
+NodalVolumeMapping.tower_factory = @VerticalConvTower
+
+# Parameters for NullFeatures:
+# ==============================================================================
+NullFeatures.name = None
+
+# Parameters for PerVariableDataFilter:
+# ==============================================================================
+PerVariableDataFilter.name = None
+PerVariableDataFilter.per_variable_filters = \
+    {'divergence': @divergence/DataExponentialFilter,
+     'log_surface_pressure': @lsp/DataExponentialFilter,
+     'sim_time': @DataNoFilter,
+     'temperature_variation': @temperature_variation/DataExponentialFilter,
+     'tracers': {'specific_cloud_ice_water_content': @q/DataExponentialFilter,
+                 'specific_cloud_liquid_water_content': @q/DataExponentialFilter,
+                 'specific_humidity': @q/DataExponentialFilter},
+     'vorticity': @vorticity/DataExponentialFilter}
+
+# Parameters for advance/PressureFeatures:
+# ==============================================================================
+advance/PressureFeatures.name = None
+
+# Parameters for embedding_model/PressureFeatures:
+# ==============================================================================
+embedding_model/PressureFeatures.name = None
+
+# Parameters for primitive_eq_specs_constructor:
+# ==============================================================================
+primitive_eq_specs_constructor.scale = @GET_ATMOSPHERIC_SCALE()
+
+# Parameters for primitive_eq_to_xarray:
+# ==============================================================================
+# None.
+
+# Parameters for PrimitiveToWeatherbenchDecoder:
+# ==============================================================================
+# None.
+
+# Parameters for advance/RadiationFeatures:
+# ==============================================================================
+advance/RadiationFeatures.name = None
+
+# Parameters for decoder_model/RadiationFeatures:
+# ==============================================================================
+decoder_model/RadiationFeatures.name = None
+
+# Parameters for encoder_data/RadiationFeatures:
+# ==============================================================================
+encoder_data/RadiationFeatures.name = None
+
+# Parameters for RealSphericalHarmonicsWithZeroImag:
+# ==============================================================================
+RealSphericalHarmonicsWithZeroImag.base_shape_multiple = %BASE_SHAPE_MULTIPLE
+RealSphericalHarmonicsWithZeroImag.reverse_einsum_arg_order = \
+    %REVERSE_EINSUM_ARG_ORDER
+RealSphericalHarmonicsWithZeroImag.stacked_fourier_transforms = None
+RealSphericalHarmonicsWithZeroImag.transform_precision = 'tensorfloat32'
+
+# Parameters for orography/RealSphericalHarmonicsWithZeroImag:
+# ==============================================================================
+orography/RealSphericalHarmonicsWithZeroImag.base_shape_multiple = \
+    %BASE_SHAPE_MULTIPLE
+orography/RealSphericalHarmonicsWithZeroImag.reverse_einsum_arg_order = \
+    %REVERSE_EINSUM_ARG_ORDER
+orography/RealSphericalHarmonicsWithZeroImag.stacked_fourier_transforms = None
+orography/RealSphericalHarmonicsWithZeroImag.transform_precision = 'tensorfloat32'
+
+# Parameters for dycore/SequentialStepFilter:
+# ==============================================================================
+dycore/SequentialStepFilter.filter_modules = \
+    (@dycore/ExponentialFilter, @stability/ExponentialFilter)
+dycore/SequentialStepFilter.name = None
+
+# Parameters for ml/SequentialStepFilter:
+# ==============================================================================
+ml/SequentialStepFilter.filter_modules = (@stability/ExponentialFilter,)
+ml/SequentialStepFilter.name = None
+
+# Parameters for advance/SequentialTransform:
+# ==============================================================================
+advance/SequentialTransform.name = None
+advance/SequentialTransform.transform_modules = \
+    (@advance/ShiftAndNormalize,
+     @advance/InverseLevelScale,
+     @advance/TruncateSigmaLevels,
+     @SoftClip)
+
+# Parameters for decoder/SequentialTransform:
+# ==============================================================================
+decoder/SequentialTransform.name = None
+decoder/SequentialTransform.transform_modules = \
+    (@decoder/InverseShiftAndNormalize, @decoder/LevelScale)
+
+# Parameters for decoder_model/SequentialTransform:
+# ==============================================================================
+decoder_model/SequentialTransform.name = None
+decoder_model/SequentialTransform.transform_modules = \
+    (@decoder_model/ShiftAndNormalize,
+     @decoder_model/InverseLevelScale,
+     @decoder_model/TruncateSigmaLevels)
+
+# Parameters for div_curl_tendency_outputs/SequentialTransform:
+# ==============================================================================
+div_curl_tendency_outputs/SequentialTransform.name = None
+div_curl_tendency_outputs/SequentialTransform.transform_modules = \
+    (@div_curl_tendency_outputs/InverseShiftAndNormalize,
+     @div_curl_tendency_outputs/LevelScale)
+
+# Parameters for encode/SequentialTransform:
+# ==============================================================================
+encode/SequentialTransform.name = None
+encode/SequentialTransform.transform_modules = \
+    (@encoder/InverseShiftAndNormalize, @encode/LevelScale)
+
+# Parameters for encoder_data/SequentialTransform:
+# ==============================================================================
+encoder_data/SequentialTransform.name = None
+encoder_data/SequentialTransform.transform_modules = \
+    (@encoder_data/ShiftAndNormalize, @encoder_data/InverseLevelScale)
+
+# Parameters for advance/ShiftAndNormalize:
+# ==============================================================================
+advance/ShiftAndNormalize.features_to_exclude = ()
+advance/ShiftAndNormalize.global_scale = None
+advance/ShiftAndNormalize.name = None
+advance/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3068,
+     'divergence': 0.1078,
+     'divergence_del2': 1825.0,
+     'divergence_dlat': 8.5,
+     'divergence_dlon': 8.842,
+     'geopotential_at_surface': 0.009579,
+     'geopotential_at_surface_del2': 28.66,
+     'geopotential_at_surface_dlat': 0.1882,
+     'geopotential_at_surface_dlon': 0.1564,
+     'land_sea_mask': 0.456,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1134,
+     'log_surface_pressure_del2': 264.6,
+     'log_surface_pressure_dlat': 1.656,
+     'log_surface_pressure_dlon': 1.645,
+     'memory_divergence': 0.1078,
+     'memory_log_surface_pressure': 0.1134,
+     'memory_specific_cloud_ice_water_content': 9.812e-06,
+     'memory_specific_cloud_liquid_water_content': 2.039e-05,
+     'memory_specific_humidity': 0.003305,
+     'memory_temperature_variation': 15.04,
+     'memory_u': 0.01489,
+     'memory_v': 0.01022,
+     'memory_vorticity': 0.2831,
+     'pressure': 1.644,
+     'radiation': 0.2862,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7064,
+     'specific_cloud_ice_water_content': 9.812e-06,
+     'specific_cloud_ice_water_content_del2': 0.08441,
+     'specific_cloud_ice_water_content_dlat': 0.0004684,
+     'specific_cloud_ice_water_content_dlon': 0.0004646,
+     'specific_cloud_liquid_water_content': 2.039e-05,
+     'specific_cloud_liquid_water_content_del2': 0.1802,
+     'specific_cloud_liquid_water_content_dlat': 0.0009877,
+     'specific_cloud_liquid_water_content_dlon': 0.0009565,
+     'specific_humidity': 0.003305,
+     'specific_humidity_del2': 4.442,
+     'specific_humidity_dlat': 0.0319,
+     'specific_humidity_dlon': 0.02678,
+     'surface_embedding': 1.0,
+     'temperature_variation': 15.04,
+     'temperature_variation_del2': 10460.0,
+     'temperature_variation_dlat': 82.59,
+     'temperature_variation_dlon': 77.07,
+     'u': 0.01489,
+     'u_del2': 18.46,
+     'u_dlat': 0.2197,
+     'u_dlon': 0.1677,
+     'v': 0.01022,
+     'v_del2': 17.2,
+     'v_dlat': 0.169,
+     'v_dlon': 0.2083,
+     'vorticity': 0.2831,
+     'vorticity_del2': 2831.0,
+     'vorticity_dlat': 14.96,
+     'vorticity_dlon': 14.91}
+advance/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.638,
+     'divergence': -0.0,
+     'divergence_del2': -0.015,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.048,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': -0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.715,
+     'log_surface_pressure_del2': 0.59,
+     'log_surface_pressure_dlat': 0.126,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.715,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.0,
+     'memory_temperature_variation': -4.98,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.213,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': -0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.98,
+     'temperature_variation_del2': 43.608,
+     'temperature_variation_dlat': 4.843,
+     'temperature_variation_dlon': 0.0,
+     'u': 0.007,
+     'u_del2': 0.242,
+     'u_dlat': 0.002,
+     'u_dlon': 0.0,
+     'v': -0.0,
+     'v_del2': -0.0,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.365,
+     'vorticity_dlat': 0.041,
+     'vorticity_dlon': 0.0}
+
+# Parameters for decoder_model/ShiftAndNormalize:
+# ==============================================================================
+decoder_model/ShiftAndNormalize.features_to_exclude = ()
+decoder_model/ShiftAndNormalize.global_scale = None
+decoder_model/ShiftAndNormalize.name = None
+decoder_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3068,
+     'divergence': 0.1078,
+     'divergence_del2': 1825.0,
+     'divergence_dlat': 8.5,
+     'divergence_dlon': 8.842,
+     'geopotential_at_surface': 0.009579,
+     'geopotential_at_surface_del2': 28.66,
+     'geopotential_at_surface_dlat': 0.1882,
+     'geopotential_at_surface_dlon': 0.1564,
+     'land_sea_mask': 0.456,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1134,
+     'log_surface_pressure_del2': 264.6,
+     'log_surface_pressure_dlat': 1.656,
+     'log_surface_pressure_dlon': 1.645,
+     'memory_divergence': 0.1078,
+     'memory_log_surface_pressure': 0.1134,
+     'memory_specific_cloud_ice_water_content': 9.812e-06,
+     'memory_specific_cloud_liquid_water_content': 2.039e-05,
+     'memory_specific_humidity': 0.003305,
+     'memory_temperature_variation': 15.04,
+     'memory_u': 0.01489,
+     'memory_v': 0.01022,
+     'memory_vorticity': 0.2831,
+     'pressure': 1.644,
+     'radiation': 0.2862,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7064,
+     'specific_cloud_ice_water_content': 9.812e-06,
+     'specific_cloud_ice_water_content_del2': 0.08441,
+     'specific_cloud_ice_water_content_dlat': 0.0004684,
+     'specific_cloud_ice_water_content_dlon': 0.0004646,
+     'specific_cloud_liquid_water_content': 2.039e-05,
+     'specific_cloud_liquid_water_content_del2': 0.1802,
+     'specific_cloud_liquid_water_content_dlat': 0.0009877,
+     'specific_cloud_liquid_water_content_dlon': 0.0009565,
+     'specific_humidity': 0.003305,
+     'specific_humidity_del2': 4.442,
+     'specific_humidity_dlat': 0.0319,
+     'specific_humidity_dlon': 0.02678,
+     'surface_embedding': 1.0,
+     'temperature_variation': 15.04,
+     'temperature_variation_del2': 10460.0,
+     'temperature_variation_dlat': 82.59,
+     'temperature_variation_dlon': 77.07,
+     'u': 0.01489,
+     'u_del2': 18.46,
+     'u_dlat': 0.2197,
+     'u_dlon': 0.1677,
+     'v': 0.01022,
+     'v_del2': 17.2,
+     'v_dlat': 0.169,
+     'v_dlon': 0.2083,
+     'vorticity': 0.2831,
+     'vorticity_del2': 2831.0,
+     'vorticity_dlat': 14.96,
+     'vorticity_dlon': 14.91}
+decoder_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.638,
+     'divergence': -0.0,
+     'divergence_del2': -0.015,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.048,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': -0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.715,
+     'log_surface_pressure_del2': 0.59,
+     'log_surface_pressure_dlat': 0.126,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.715,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.0,
+     'memory_temperature_variation': -4.98,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.213,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': -0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.98,
+     'temperature_variation_del2': 43.608,
+     'temperature_variation_dlat': 4.843,
+     'temperature_variation_dlon': 0.0,
+     'u': 0.007,
+     'u_del2': 0.242,
+     'u_dlat': 0.002,
+     'u_dlon': 0.0,
+     'v': -0.0,
+     'v_del2': -0.0,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.365,
+     'vorticity_dlat': 0.041,
+     'vorticity_dlon': 0.0}
+
+# Parameters for embedding_model/ShiftAndNormalize:
+# ==============================================================================
+embedding_model/ShiftAndNormalize.features_to_exclude = ()
+embedding_model/ShiftAndNormalize.global_scale = None
+embedding_model/ShiftAndNormalize.name = None
+embedding_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3068,
+     'divergence': 0.1078,
+     'divergence_del2': 1825.0,
+     'divergence_dlat': 8.5,
+     'divergence_dlon': 8.842,
+     'geopotential_at_surface': 0.009579,
+     'geopotential_at_surface_del2': 28.66,
+     'geopotential_at_surface_dlat': 0.1882,
+     'geopotential_at_surface_dlon': 0.1564,
+     'land_sea_mask': 0.456,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1134,
+     'log_surface_pressure_del2': 264.6,
+     'log_surface_pressure_dlat': 1.656,
+     'log_surface_pressure_dlon': 1.645,
+     'memory_divergence': 0.1078,
+     'memory_log_surface_pressure': 0.1134,
+     'memory_specific_cloud_ice_water_content': 9.812e-06,
+     'memory_specific_cloud_liquid_water_content': 2.039e-05,
+     'memory_specific_humidity': 0.003305,
+     'memory_temperature_variation': 15.04,
+     'memory_u': 0.01489,
+     'memory_v': 0.01022,
+     'memory_vorticity': 0.2831,
+     'pressure': 1.644,
+     'radiation': 0.2862,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7064,
+     'specific_cloud_ice_water_content': 9.812e-06,
+     'specific_cloud_ice_water_content_del2': 0.08441,
+     'specific_cloud_ice_water_content_dlat': 0.0004684,
+     'specific_cloud_ice_water_content_dlon': 0.0004646,
+     'specific_cloud_liquid_water_content': 2.039e-05,
+     'specific_cloud_liquid_water_content_del2': 0.1802,
+     'specific_cloud_liquid_water_content_dlat': 0.0009877,
+     'specific_cloud_liquid_water_content_dlon': 0.0009565,
+     'specific_humidity': 0.003305,
+     'specific_humidity_del2': 4.442,
+     'specific_humidity_dlat': 0.0319,
+     'specific_humidity_dlon': 0.02678,
+     'surface_embedding': 1.0,
+     'temperature_variation': 15.04,
+     'temperature_variation_del2': 10460.0,
+     'temperature_variation_dlat': 82.59,
+     'temperature_variation_dlon': 77.07,
+     'u': 0.01489,
+     'u_del2': 18.46,
+     'u_dlat': 0.2197,
+     'u_dlon': 0.1677,
+     'v': 0.01022,
+     'v_del2': 17.2,
+     'v_dlat': 0.169,
+     'v_dlon': 0.2083,
+     'vorticity': 0.2831,
+     'vorticity_del2': 2831.0,
+     'vorticity_dlat': 14.96,
+     'vorticity_dlon': 14.91}
+embedding_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.638,
+     'divergence': -0.0,
+     'divergence_del2': -0.015,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.048,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': -0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.715,
+     'log_surface_pressure_del2': 0.59,
+     'log_surface_pressure_dlat': 0.126,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.715,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.0,
+     'memory_temperature_variation': -4.98,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.213,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': -0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.98,
+     'temperature_variation_del2': 43.608,
+     'temperature_variation_dlat': 4.843,
+     'temperature_variation_dlon': 0.0,
+     'u': 0.007,
+     'u_del2': 0.242,
+     'u_dlat': 0.002,
+     'u_dlon': 0.0,
+     'v': -0.0,
+     'v_del2': -0.0,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.365,
+     'vorticity_dlat': 0.041,
+     'vorticity_dlon': 0.0}
+
+# Parameters for encoder_data/ShiftAndNormalize:
+# ==============================================================================
+encoder_data/ShiftAndNormalize.features_to_exclude = ()
+encoder_data/ShiftAndNormalize.global_scale = None
+encoder_data/ShiftAndNormalize.name = None
+encoder_data/ShiftAndNormalize.scales = \
+    {'cos_latitude': 0.3068,
+     'geopotential_at_surface': 0.009579,
+     'geopotential_at_surface_del2': 28.66,
+     'geopotential_at_surface_dlat': 0.1882,
+     'geopotential_at_surface_dlon': 0.1564,
+     'land_sea_mask': 0.456,
+     'learned_positional_features': 1.0,
+     'radiation': 0.2862,
+     'sin_latitude': 0.7064,
+     'specific_cloud_ice_water_content': 8.255e-06,
+     'specific_cloud_ice_water_content_del2': 0.09943,
+     'specific_cloud_ice_water_content_dlat': 0.0005741,
+     'specific_cloud_ice_water_content_dlon': 0.0004452,
+     'specific_cloud_liquid_water_content': 2.182e-05,
+     'specific_cloud_liquid_water_content_del2': 0.2881,
+     'specific_cloud_liquid_water_content_dlat': 0.001674,
+     'specific_cloud_liquid_water_content_dlon': 0.001224,
+     'specific_humidity': 0.003493,
+     'specific_humidity_del2': 5.716,
+     'specific_humidity_dlat': 0.0396,
+     'specific_humidity_dlon': 0.02857,
+     't': 27.99,
+     't_del2': 9788.0,
+     't_dlat': 71.14,
+     't_dlon': 62.14,
+     'u': 0.01935,
+     'u_del2': 27.4,
+     'u_dlat': 0.1876,
+     'u_dlon': 0.1712,
+     'v': 0.01038,
+     'v_del2': 27.11,
+     'v_dlat': 0.1326,
+     'v_dlon': 0.1998,
+     'z': 0.1496,
+     'z_del2': 0.6108,
+     'z_dlat': 0.01606,
+     'z_dlon': 0.009139}
+encoder_data/ShiftAndNormalize.shifts = \
+    {'cos_latitude': 0.638,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.048,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': -0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'radiation': 0.213,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': -0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': 0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': 0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     't': 247.045,
+     't_del2': 37.572,
+     't_dlat': 1.755,
+     't_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.056,
+     'u_dlat': 0.0,
+     'u_dlon': 0.0,
+     'v': 0.0,
+     'v_del2': 0.008,
+     'v_dlat': 0.0,
+     'v_dlon': -0.0,
+     'z': 0.145,
+     'z_del2': 0.011,
+     'z_dlat': 0.001,
+     'z_dlon': 0.0}
+
+# Parameters for sea_model/ShiftAndNormalize:
+# ==============================================================================
+sea_model/ShiftAndNormalize.features_to_exclude = ()
+sea_model/ShiftAndNormalize.global_scale = None
+sea_model/ShiftAndNormalize.name = None
+sea_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3068,
+     'divergence': 0.1078,
+     'divergence_del2': 1825.0,
+     'divergence_dlat': 8.5,
+     'divergence_dlon': 8.842,
+     'geopotential_at_surface': 0.009579,
+     'geopotential_at_surface_del2': 28.66,
+     'geopotential_at_surface_dlat': 0.1882,
+     'geopotential_at_surface_dlon': 0.1564,
+     'land_sea_mask': 0.456,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1134,
+     'log_surface_pressure_del2': 264.6,
+     'log_surface_pressure_dlat': 1.656,
+     'log_surface_pressure_dlon': 1.645,
+     'memory_divergence': 0.1078,
+     'memory_log_surface_pressure': 0.1134,
+     'memory_specific_cloud_ice_water_content': 9.812e-06,
+     'memory_specific_cloud_liquid_water_content': 2.039e-05,
+     'memory_specific_humidity': 0.003305,
+     'memory_temperature_variation': 15.04,
+     'memory_u': 0.01489,
+     'memory_v': 0.01022,
+     'memory_vorticity': 0.2831,
+     'pressure': 1.644,
+     'radiation': 0.2862,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7064,
+     'specific_cloud_ice_water_content': 9.812e-06,
+     'specific_cloud_ice_water_content_del2': 0.08441,
+     'specific_cloud_ice_water_content_dlat': 0.0004684,
+     'specific_cloud_ice_water_content_dlon': 0.0004646,
+     'specific_cloud_liquid_water_content': 2.039e-05,
+     'specific_cloud_liquid_water_content_del2': 0.1802,
+     'specific_cloud_liquid_water_content_dlat': 0.0009877,
+     'specific_cloud_liquid_water_content_dlon': 0.0009565,
+     'specific_humidity': 0.003305,
+     'specific_humidity_del2': 4.442,
+     'specific_humidity_dlat': 0.0319,
+     'specific_humidity_dlon': 0.02678,
+     'surface_embedding': 1.0,
+     'temperature_variation': 15.04,
+     'temperature_variation_del2': 10460.0,
+     'temperature_variation_dlat': 82.59,
+     'temperature_variation_dlon': 77.07,
+     'u': 0.01489,
+     'u_del2': 18.46,
+     'u_dlat': 0.2197,
+     'u_dlon': 0.1677,
+     'v': 0.01022,
+     'v_del2': 17.2,
+     'v_dlat': 0.169,
+     'v_dlon': 0.2083,
+     'vorticity': 0.2831,
+     'vorticity_del2': 2831.0,
+     'vorticity_dlat': 14.96,
+     'vorticity_dlon': 14.91}
+sea_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.638,
+     'divergence': -0.0,
+     'divergence_del2': -0.015,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.048,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': -0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.715,
+     'log_surface_pressure_del2': 0.59,
+     'log_surface_pressure_dlat': 0.126,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.715,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.0,
+     'memory_temperature_variation': -4.98,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.213,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': -0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.98,
+     'temperature_variation_del2': 43.608,
+     'temperature_variation_dlat': 4.843,
+     'temperature_variation_dlon': 0.0,
+     'u': 0.007,
+     'u_del2': 0.242,
+     'u_dlat': 0.002,
+     'u_dlon': 0.0,
+     'v': -0.0,
+     'v_del2': -0.0,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.365,
+     'vorticity_dlat': 0.041,
+     'vorticity_dlon': 0.0}
+
+# Parameters for SigmaCoordinatesEquidistant:
+# ==============================================================================
+SigmaCoordinatesEquidistant.layers = %N_SIGMA_LAYERS
+
+# Parameters for custom_corrds/SigmaCoordinatesEquidistant:
+# ==============================================================================
+custom_corrds/SigmaCoordinatesEquidistant.layers = %N_SIGMA_LAYERS
+
+# Parameters for advance/SoftClip:
+# ==============================================================================
+advance/SoftClip.hinge_softness = 1.0
+advance/SoftClip.max_value = 16
+advance/SoftClip.name = None
+
+# Parameters for StochasticModularStepModel:
+# ==============================================================================
+StochasticModularStepModel.advance_module = @StochasticPhysicsParameterizationStep
+StochasticModularStepModel.decoder_module = \
+    @DimensionalLearnedPrimitiveToWeatherbenchDecoder
+StochasticModularStepModel.encoder_module = \
+    @DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder
+StochasticModularStepModel.forcing_module = @DynamicDataForcing
+StochasticModularStepModel.name = None
+
+# Parameters for StochasticPhysicsParameterizationStep:
+# ==============================================================================
+StochasticPhysicsParameterizationStep.checkpoint_substep = False
+StochasticPhysicsParameterizationStep.corrector_module = %CORRECTOR_MODULE
+StochasticPhysicsParameterizationStep.name = None
+StochasticPhysicsParameterizationStep.num_substeps = %NUM_SUBSTEPS
+StochasticPhysicsParameterizationStep.physics_parameterization_module = \
+    @DivCurlNeuralParameterization
+StochasticPhysicsParameterizationStep.randomness_module = @ZerosRandomField
+
+# Parameters for advance/ToModalDiffOperators:
+# ==============================================================================
+advance/ToModalDiffOperators.name = None
+
+# Parameters for encoder_data/ToModalDiffOperators:
+# ==============================================================================
+encoder_data/ToModalDiffOperators.name = None
+
+# Parameters for with_grads/ToModalDiffOperators:
+# ==============================================================================
+with_grads/ToModalDiffOperators.name = None
+
+# Parameters for trajectory_from_step:
+# ==============================================================================
+trajectory_from_step.checkpoint_multistep = False
+trajectory_from_step.checkpoint_post_process = True
+trajectory_from_step.checkpoint_step = True
+
+# Parameters for advance/TruncateSigmaLevels:
+# ==============================================================================
+advance/TruncateSigmaLevels.name = None
+advance/TruncateSigmaLevels.sigma_ranges = \
+    {'divergence': (0.0, 1),
+     'geopotential': (0.0, 1),
+     'specific_cloud_ice_water_content': (0.0, 1),
+     'specific_cloud_liquid_water_content': (0.0, 1),
+     'specific_humidity': (0.0, 1),
+     't': (0.0, 1),
+     'temperature_variation': (0.0, 1),
+     'u': (0.0, 1),
+     'u_component_of_wind': (0.0, 1),
+     'v': (0.0, 1),
+     'v_component_of_wind': (0.0, 1),
+     'vorticity': (0.0, 1),
+     'z': (0.0, 1)}
+
+# Parameters for decoder_model/TruncateSigmaLevels:
+# ==============================================================================
+decoder_model/TruncateSigmaLevels.name = None
+decoder_model/TruncateSigmaLevels.sigma_ranges = \
+    {'divergence': (0.0, 1),
+     'geopotential': (0.0, 1),
+     'specific_cloud_ice_water_content': (0.0, 1),
+     'specific_cloud_liquid_water_content': (0.0, 1),
+     'specific_humidity': (0.0, 1),
+     't': (0.0, 1),
+     'temperature_variation': (0.0, 1),
+     'u': (0.0, 1),
+     'u_component_of_wind': (0.0, 1),
+     'v': (0.0, 1),
+     'v_component_of_wind': (0.0, 1),
+     'vorticity': (0.0, 1),
+     'z': (0.0, 1)}
+
+# Parameters for advance/VelocityAndPrognostics:
+# ==============================================================================
+advance/VelocityAndPrognostics.compute_gradients_module = @ToModalDiffOperators
+advance/VelocityAndPrognostics.fields_to_include = \
+    ['divergence',
+     'vorticity',
+     'u',
+     'v',
+     'temperature_variation',
+     'log_surface_pressure',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+advance/VelocityAndPrognostics.name = None
+
+# Parameters for embedding_model/VelocityAndPrognostics:
+# ==============================================================================
+embedding_model/VelocityAndPrognostics.fields_to_include = \
+    ['divergence',
+     'vorticity',
+     'u',
+     'v',
+     'temperature_variation',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+embedding_model/VelocityAndPrognostics.name = None
+
+# Parameters for encoder_data/VelocityAndPrognostics:
+# ==============================================================================
+encoder_data/VelocityAndPrognostics.compute_gradients_module = \
+    @ToModalDiffOperators
+encoder_data/VelocityAndPrognostics.fields_to_include = \
+    ['u',
+     'v',
+     't',
+     'z',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+encoder_data/VelocityAndPrognostics.name = None
+
+# Parameters for model/VelocityAndPrognostics:
+# ==============================================================================
+model/VelocityAndPrognostics.fields_to_include = None
+model/VelocityAndPrognostics.name = None
+
+# Parameters for VerticalConvTower:
+# ==============================================================================
+VerticalConvTower.activate_final = False
+VerticalConvTower.activation = %ACTIVATION
+VerticalConvTower.channels = [64, 64, 64, 64]
+VerticalConvTower.checkpoint_tower = True
+VerticalConvTower.kernel_shape = 5
+VerticalConvTower.name = None
+VerticalConvTower.with_bias = True
+
+# Parameters for WhirlModel:
+# ==============================================================================
+WhirlModel.from_xarray_fn = @xarray_to_state_and_dynamic_covariate_data
+WhirlModel.model_cls = @StochasticModularStepModel
+WhirlModel.to_xarray_fn = @data_to_xarray_with_renaming
+
+# Parameters for xarray_to_data_with_renaming:
+# ==============================================================================
+xarray_to_data_with_renaming.renaming_dict = \
+    {'geopotential': 'z',
+     'latitude': 'lat',
+     'longitude': 'lon',
+     'temperature': 't',
+     'u_component_of_wind': 'u',
+     'v_component_of_wind': 'v'}
+xarray_to_data_with_renaming.xarray_to_data_fn = @xarray_to_weatherbench_data
+
+# Parameters for xarray_to_dynamic_covariate_data:
+# ==============================================================================
+xarray_to_dynamic_covariate_data.covariates_to_include = \
+    ('sea_ice_cover', 'sea_surface_temperature')
+
+# Parameters for xarray_to_state_and_dynamic_covariate_data:
+# ==============================================================================
+xarray_to_state_and_dynamic_covariate_data.values = 'values'
+xarray_to_state_and_dynamic_covariate_data.xarray_to_dynamic_covariate_data_fn = \
+    @xarray_to_dynamic_covariate_data
+xarray_to_state_and_dynamic_covariate_data.xarray_to_state_data_fn = \
+    @xarray_to_data_with_renaming
+
+# Parameters for xarray_to_weatherbench_data:
+# ==============================================================================
+xarray_to_weatherbench_data.diagnostics_to_include = ()
+xarray_to_weatherbench_data.tracers_to_include = \
+    ('specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content')
+
+# Parameters for ZerosRandomField:
+# ==============================================================================
+ZerosRandomField.prefer_nodal = True

model/reference_code/paper_configs/deterministic_1_4_deg.gin

@@ -0,0 +1,2371 @@
+# Macros:
+# ==============================================================================
+ACTIVATION = @gelu
+CORRECTOR_MODULE = @CustomCoordsCorrector
+CORRECTOR_SCALE = 0.01
+DATA_FILTER_ATTENUATION = 0.0
+DYCORE_FILTER_ORDER = 3
+DYCORE_GRID = @GridWithWavenumbers()
+DYCORE_INTEGRATOR = @imex_rk_sil3
+DYCORE_TAU = '120 minutes'
+GLOBAL_OUT_SCALE = 0.02
+LATENT_SIZE = 384
+LAYER_SIZE = 384
+N_CNN_FEATURES = 32
+N_INNER_DYCORE_STEPS = 5
+N_SIGMA_LAYERS = 32
+N_TO_CLIP = 1
+NUM_BLOCKS = 5
+NUM_SUBSTEPS = 2
+PARAMETERIZATION_FILTER = @ml/SequentialStepFilter
+POSITIONAL_LATENT_SIZE = 32
+STABILITY_TAU = '8 minutes'
+SURFACE_MODEL_LATENT_SIZE = 8
+SURFACE_MODEL_LAYER_SIZE = 8
+SURFACE_MODEL_OUTPUT_SIZE = 8
+
+# Parameters for decode/ColumnTower:
+# ==============================================================================
+decode/ColumnTower.checkpoint_tower = False
+decode/ColumnTower.column_net_factory = @decode/MlpUniform
+decode/ColumnTower.name = 'decode_tower'
+
+# Parameters for encode/ColumnTower:
+# ==============================================================================
+encode/ColumnTower.checkpoint_tower = False
+encode/ColumnTower.column_net_factory = @encode/MlpUniform
+encode/ColumnTower.name = 'encode_tower'
+
+# Parameters for process/ColumnTower:
+# ==============================================================================
+process/ColumnTower.checkpoint_tower = False
+process/ColumnTower.column_net_factory = @process/MlpUniform
+process/ColumnTower.name = 'process_tower'
+
+# Parameters for surface_model_decode/ColumnTower:
+# ==============================================================================
+surface_model_decode/ColumnTower.checkpoint_tower = False
+surface_model_decode/ColumnTower.column_net_factory = \
+    @surface_model_decode/MlpUniform
+surface_model_decode/ColumnTower.name = 'surface_model_decode_tower'
+
+# Parameters for surface_model_encode/ColumnTower:
+# ==============================================================================
+surface_model_encode/ColumnTower.checkpoint_tower = False
+surface_model_encode/ColumnTower.column_net_factory = \
+    @surface_model_encode/MlpUniform
+surface_model_encode/ColumnTower.name = 'surface_model_encode_tower'
+
+# Parameters for surface_model_process/ColumnTower:
+# ==============================================================================
+surface_model_process/ColumnTower.checkpoint_tower = False
+surface_model_process/ColumnTower.column_net_factory = \
+    @surface_model_process/MlpUniform
+surface_model_process/ColumnTower.name = 'surface_model_process_tower'
+
+# Parameters for advance/CombinedFeatures:
+# ==============================================================================
+advance/CombinedFeatures.feature_module_names_to_exclude = ()
+advance/CombinedFeatures.feature_modules = \
+    (@EmbeddingSurfaceFeatures,
+     @EmbeddingVolumeFeatures,
+     @PressureFeatures,
+     @RadiationFeatures,
+     @LatitudeFeatures,
+     @advance/VelocityAndPrognostics,
+     @MemoryVelocityAndValues,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+advance/CombinedFeatures.features_to_exclude = ()
+advance/CombinedFeatures.features_transform_module = @advance/SequentialTransform
+advance/CombinedFeatures.name = None
+
+# Parameters for decoder_model/CombinedFeatures:
+# ==============================================================================
+decoder_model/CombinedFeatures.feature_module_names_to_exclude = ()
+decoder_model/CombinedFeatures.feature_modules = \
+    (@RadiationFeatures,
+     @LatitudeFeatures,
+     @model/VelocityAndPrognostics,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+decoder_model/CombinedFeatures.features_to_exclude = ()
+decoder_model/CombinedFeatures.features_transform_module = \
+    @decoder_model/SequentialTransform
+decoder_model/CombinedFeatures.name = None
+
+# Parameters for embedding_model/CombinedFeatures:
+# ==============================================================================
+embedding_model/CombinedFeatures.feature_module_names_to_exclude = ()
+embedding_model/CombinedFeatures.feature_modules = \
+    (@embedding_model/VelocityAndPrognostics, @PressureFeatures)
+embedding_model/CombinedFeatures.features_to_exclude = ()
+embedding_model/CombinedFeatures.features_transform_module = \
+    @embedding_model/ShiftAndNormalize
+embedding_model/CombinedFeatures.name = None
+
+# Parameters for encoder_data/CombinedFeatures:
+# ==============================================================================
+encoder_data/CombinedFeatures.feature_module_names_to_exclude = ()
+encoder_data/CombinedFeatures.feature_modules = \
+    (@RadiationFeatures,
+     @LatitudeFeatures,
+     @encoder_data/VelocityAndPrognostics,
+     @with_grads/FloatDataFeatures,
+     @without_grads/FloatDataFeatures,
+     @LearnedPositionalFeatures)
+encoder_data/CombinedFeatures.features_to_exclude = ()
+encoder_data/CombinedFeatures.features_transform_module = \
+    @encoder_data/SequentialTransform
+encoder_data/CombinedFeatures.name = None
+
+# Parameters for sea_model/CombinedFeatures:
+# ==============================================================================
+sea_model/CombinedFeatures.feature_module_names_to_exclude = ()
+sea_model/CombinedFeatures.feature_modules = (@sea_model/ForcingFeatures,)
+sea_model/CombinedFeatures.features_to_exclude = ()
+sea_model/CombinedFeatures.features_transform_module = @sea_model/ShiftAndNormalize
+sea_model/CombinedFeatures.name = None
+
+# Parameters for coordinate_system_from_dataset:
+# ==============================================================================
+coordinate_system_from_dataset.spherical_harmonics_impl = \
+    @RealSphericalHarmonicsWithZeroImag
+coordinate_system_from_dataset.truncation = 'LINEAR'
+
+# Parameters for CoordinateSystem:
+# ==============================================================================
+CoordinateSystem.horizontal = @GridTL127()
+CoordinateSystem.vertical = @SigmaCoordinatesEquidistant()
+
+# Parameters for custom_corrds/CoordinateSystem:
+# ==============================================================================
+custom_corrds/CoordinateSystem.horizontal = %DYCORE_GRID
+custom_corrds/CoordinateSystem.vertical = @SigmaCoordinatesEquidistant()
+
+# Parameters for CustomCoordsCorrector:
+# ==============================================================================
+CustomCoordsCorrector.corrector_module = @DycoreWithPhysicsCorrector
+CustomCoordsCorrector.custom_coords = @custom_corrds/CoordinateSystem()
+CustomCoordsCorrector.name = None
+
+# Parameters for data_to_xarray_with_renaming:
+# ==============================================================================
+data_to_xarray_with_renaming.additional_coords = None
+data_to_xarray_with_renaming.attrs = None
+data_to_xarray_with_renaming.renaming_dict = \
+    {'geopotential': 'z',
+     'latitude': 'lat',
+     'longitude': 'lon',
+     'temperature': 't',
+     'u_component_of_wind': 'u',
+     'v_component_of_wind': 'v'}
+data_to_xarray_with_renaming.sample_ids = None
+data_to_xarray_with_renaming.to_xarray_fn = @primitive_eq_to_xarray
+
+# Parameters for divergence/DataExponentialFilter:
+# ==============================================================================
+divergence/DataExponentialFilter.attenuation = 14.0
+divergence/DataExponentialFilter.cutoff = 0.62
+divergence/DataExponentialFilter.name = None
+divergence/DataExponentialFilter.order = 2
+
+# Parameters for lsp/DataExponentialFilter:
+# ==============================================================================
+lsp/DataExponentialFilter.attenuation = 8
+lsp/DataExponentialFilter.cutoff = 0.82
+lsp/DataExponentialFilter.name = None
+lsp/DataExponentialFilter.order = 1
+
+# Parameters for orography/DataExponentialFilter:
+# ==============================================================================
+orography/DataExponentialFilter.attenuation = %DATA_FILTER_ATTENUATION
+orography/DataExponentialFilter.cutoff = 0
+orography/DataExponentialFilter.name = None
+orography/DataExponentialFilter.order = 1
+
+# Parameters for q/DataExponentialFilter:
+# ==============================================================================
+q/DataExponentialFilter.attenuation = 14
+q/DataExponentialFilter.cutoff = 0.62
+q/DataExponentialFilter.name = None
+q/DataExponentialFilter.order = 2
+
+# Parameters for temperature_variation/DataExponentialFilter:
+# ==============================================================================
+temperature_variation/DataExponentialFilter.attenuation = 13.0
+temperature_variation/DataExponentialFilter.cutoff = 0.45
+temperature_variation/DataExponentialFilter.name = None
+temperature_variation/DataExponentialFilter.order = 3
+
+# Parameters for vorticity/DataExponentialFilter:
+# ==============================================================================
+vorticity/DataExponentialFilter.attenuation = 14
+vorticity/DataExponentialFilter.cutoff = 0.62
+vorticity/DataExponentialFilter.name = None
+vorticity/DataExponentialFilter.order = 2
+
+# Parameters for DataNoFilter:
+# ==============================================================================
+DataNoFilter.name = None
+
+# Parameters for DimensionalLearnedPrimitiveToWeatherbenchDecoder:
+# ==============================================================================
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.correction_transform_module = \
+    @decoder/SequentialTransform
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.inputs_to_units_mapping = \
+    {'sim_time': 'dimensionless',
+     't': 'kelvin',
+     'tracers': {'specific_cloud_ice_water_content': 'dimensionless',
+                 'specific_cloud_liquid_water_content': 'dimensionless',
+                 'specific_humidity': 'dimensionless'},
+     'u': 'meter / second',
+     'v': 'meter / second',
+     'z': 'm**2 s**-2'}
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.modal_to_nodal_data_features_module = \
+    @NullFeatures
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.modal_to_nodal_model_features_module = \
+    @decoder_model/CombinedFeatures
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.name = None
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.nodal_mapping_module = \
+    @NodalMapping
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.orography_module = \
+    @LearnedOrography
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.prediction_mask = \
+    {'sim_time': False,
+     't': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'u': True,
+     'v': True,
+     'z': True}
+DimensionalLearnedPrimitiveToWeatherbenchDecoder.time_axis = 0
+
+# Parameters for DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder:
+# ==============================================================================
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.correction_transform_module = \
+    @encode/SequentialTransform
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.inputs_to_units_mapping = \
+    {'sim_time': 'dimensionless',
+     't': 'kelvin',
+     'tracers': {'specific_cloud_ice_water_content': 'dimensionless',
+                 'specific_cloud_liquid_water_content': 'dimensionless',
+                 'specific_humidity': 'dimensionless'},
+     'u': 'meter / second',
+     'v': 'meter / second',
+     'z': 'm**2 s**-2'}
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.modal_to_nodal_data_features_module = \
+    @encoder_data/CombinedFeatures
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.modal_to_nodal_model_features_module = \
+    @NullFeatures
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.name = None
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.nodal_mapping_module = \
+    @NodalMapping
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.orography_module = \
+    @LearnedOrography
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.prediction_mask = \
+    {'divergence': True,
+     'log_surface_pressure': True,
+     'sim_time': False,
+     'temperature_variation': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'vorticity': True}
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.time_axis = 0
+DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder.transform_module = \
+    @EncoderCombinedTransform
+
+# Parameters for DivCurlNeuralParameterization:
+# ==============================================================================
+DivCurlNeuralParameterization.filter_module = %PARAMETERIZATION_FILTER
+DivCurlNeuralParameterization.modal_to_nodal_features_module = \
+    @advance/CombinedFeatures
+DivCurlNeuralParameterization.name = None
+DivCurlNeuralParameterization.nodal_mapping_module = @NodalMapping
+DivCurlNeuralParameterization.prediction_mask = \
+    {'divergence': True,
+     'log_surface_pressure': False,
+     'sim_time': False,
+     'temperature_variation': True,
+     'tracers': {'specific_cloud_ice_water_content': True,
+                 'specific_cloud_liquid_water_content': True,
+                 'specific_humidity': True},
+     'vorticity': True}
+DivCurlNeuralParameterization.tendency_transform_module = \
+    @div_curl_tendency_outputs/SequentialTransform
+
+# Parameters for DycoreWithPhysicsCorrector:
+# ==============================================================================
+DycoreWithPhysicsCorrector.checkpoint_explicit_terms = True
+DycoreWithPhysicsCorrector.dycore_equation_module = \
+    @MoistPrimitiveEquationsWithCloudMoisture
+DycoreWithPhysicsCorrector.dycore_substeps = %N_INNER_DYCORE_STEPS
+DycoreWithPhysicsCorrector.filter_module = @dycore/SequentialStepFilter
+DycoreWithPhysicsCorrector.name = None
+DycoreWithPhysicsCorrector.time_integrator = %DYCORE_INTEGRATOR
+
+# Parameters for DynamicDataForcing:
+# ==============================================================================
+DynamicDataForcing.check_sim_time_errors = False
+DynamicDataForcing.data_time_step = '12 hours'
+DynamicDataForcing.dt_tolerance = '1 year'
+DynamicDataForcing.inputs_to_units_mapping = \
+    {'sea_ice_cover': 'dimensionless',
+     'sea_surface_temperature': 'kelvin',
+     'sim_time': 'dimensionless'}
+DynamicDataForcing.name = None
+DynamicDataForcing.time_axis = 0
+
+# Parameters for advance/EmbeddingSurfaceFeatures:
+# ==============================================================================
+advance/EmbeddingSurfaceFeatures.embedding_module = @NodalLandSeaIceEmbedding
+advance/EmbeddingSurfaceFeatures.feature_name = 'surface_embedding'
+advance/EmbeddingSurfaceFeatures.name = None
+advance/EmbeddingSurfaceFeatures.output_size = %SURFACE_MODEL_OUTPUT_SIZE
+
+# Parameters for advance/EmbeddingVolumeFeatures:
+# ==============================================================================
+advance/EmbeddingVolumeFeatures.embedding_module = @ModalToNodalEmbedding
+advance/EmbeddingVolumeFeatures.feature_name = 'CNN1D'
+advance/EmbeddingVolumeFeatures.name = None
+advance/EmbeddingVolumeFeatures.output_size = %N_CNN_FEATURES
+
+# Parameters for EncoderCombinedTransform:
+# ==============================================================================
+EncoderCombinedTransform.name = None
+EncoderCombinedTransform.transforms = \
+    (@InputClipTransform, @EncoderFilterTransform)
+
+# Parameters for EncoderFilterTransform:
+# ==============================================================================
+EncoderFilterTransform.filter_modules = (@PerVariableDataFilter,)
+EncoderFilterTransform.name = None
+
+# Parameters for EpdTower:
+# ==============================================================================
+EpdTower.decode_tower_factory = @decode/ColumnTower
+EpdTower.encode_tower_factory = @encode/ColumnTower
+EpdTower.final_activation = None
+EpdTower.latent_size = %LATENT_SIZE
+EpdTower.name = None
+EpdTower.num_process_blocks = %NUM_BLOCKS
+EpdTower.post_encode_activation = None
+EpdTower.pre_decode_activation = None
+EpdTower.process_tower_factory = @process/ColumnTower
+
+# Parameters for surface_model/EpdTower:
+# ==============================================================================
+surface_model/EpdTower.decode_tower_factory = @surface_model_decode/ColumnTower
+surface_model/EpdTower.encode_tower_factory = @surface_model_encode/ColumnTower
+surface_model/EpdTower.final_activation = None
+surface_model/EpdTower.latent_size = %SURFACE_MODEL_LATENT_SIZE
+surface_model/EpdTower.name = None
+surface_model/EpdTower.num_process_blocks = 1
+surface_model/EpdTower.post_encode_activation = None
+surface_model/EpdTower.pre_decode_activation = None
+surface_model/EpdTower.process_tower_factory = @surface_model_process/ColumnTower
+
+# Parameters for dycore/ExponentialFilter:
+# ==============================================================================
+dycore/ExponentialFilter.cutoff = 0
+dycore/ExponentialFilter.name = None
+dycore/ExponentialFilter.order = %DYCORE_FILTER_ORDER
+dycore/ExponentialFilter.tau = %DYCORE_TAU
+
+# Parameters for stability/ExponentialFilter:
+# ==============================================================================
+stability/ExponentialFilter.cutoff = 0.4
+stability/ExponentialFilter.name = None
+stability/ExponentialFilter.order = 6
+stability/ExponentialFilter.tau = %STABILITY_TAU
+
+# Parameters for FilteredCustomOrography:
+# ==============================================================================
+FilteredCustomOrography.filter_modules = (@orography/DataExponentialFilter,)
+FilteredCustomOrography.name = None
+FilteredCustomOrography.orography_data_path = None
+FilteredCustomOrography.renaming_dict = {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for with_grads/FloatDataFeatures:
+# ==============================================================================
+with_grads/FloatDataFeatures.compute_gradients_module = @ToModalDiffOperators
+with_grads/FloatDataFeatures.covariate_data_path = None
+with_grads/FloatDataFeatures.covariate_keys = ('geopotential_at_surface',)
+with_grads/FloatDataFeatures.name = None
+with_grads/FloatDataFeatures.renaming_dict = \
+    {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for without_grads/FloatDataFeatures:
+# ==============================================================================
+without_grads/FloatDataFeatures.covariate_data_path = None
+without_grads/FloatDataFeatures.covariate_keys = ('land_sea_mask',)
+without_grads/FloatDataFeatures.name = None
+without_grads/FloatDataFeatures.renaming_dict = \
+    {'latitude': 'lat', 'longitude': 'lon'}
+
+# Parameters for sea_model/ForcingFeatures:
+# ==============================================================================
+sea_model/ForcingFeatures.forcing_to_include = ('sea_surface_temperature',)
+sea_model/ForcingFeatures.name = None
+
+# Parameters for gelu:
+# ==============================================================================
+gelu.approximate = True
+
+# Parameters for GET_ATMOSPHERIC_SCALE:
+# ==============================================================================
+# None.
+
+# Parameters for get_model_specs:
+# ==============================================================================
+get_model_specs.custom_coords = @CoordinateSystem()
+get_model_specs.model_time_step = '1 hour'
+get_model_specs.reference_datetime_str = None
+get_model_specs.reference_temperature = \
+    [215.58614815,
+     211.47405876,
+     205.87815406,
+     206.40755302,
+     210.43452345,
+     214.5683887,
+     218.75303863,
+     223.23145107,
+     227.9710687,
+     232.85381503,
+     237.53588735,
+     242.05068293,
+     246.29986585,
+     250.14294113,
+     253.74839535,
+     256.98024283,
+     259.94441031,
+     262.7041158,
+     265.21752838,
+     267.62333985,
+     269.94462121,
+     272.10056439,
+     274.12518288,
+     275.99833711,
+     277.72759392,
+     279.3292128,
+     280.79178708,
+     282.13507065,
+     283.41832023,
+     284.7682506,
+     286.33945487,
+     288.06707666]
+
+# Parameters for get_physics_specs:
+# ==============================================================================
+get_physics_specs.construct_fn = @primitive_eq_specs_constructor
+
+# Parameters for GridTL127:
+# ==============================================================================
+GridTL127.spherical_harmonics_impl = @RealSphericalHarmonicsWithZeroImag
+
+# Parameters for DYCORE_GRID/GridWithWavenumbers:
+# ==============================================================================
+DYCORE_GRID/GridWithWavenumbers.dealiasing = 'quadratic'
+DYCORE_GRID/GridWithWavenumbers.latitude_spacing = 'gauss'
+DYCORE_GRID/GridWithWavenumbers.longitude_offset = 0.0
+DYCORE_GRID/GridWithWavenumbers.longitude_wavenumbers = 126
+DYCORE_GRID/GridWithWavenumbers.radius = None
+DYCORE_GRID/GridWithWavenumbers.spherical_harmonics_impl = \
+    @RealSphericalHarmonicsWithZeroImag
+
+# Parameters for advance/IdentityTransform:
+# ==============================================================================
+advance/IdentityTransform.name = None
+
+# Parameters for sea_model/IdentityTransform:
+# ==============================================================================
+sea_model/IdentityTransform.name = None
+
+# Parameters for imex_rk_sil3:
+# ==============================================================================
+# None.
+
+# Parameters for InputClipTransform:
+# ==============================================================================
+InputClipTransform.name = None
+InputClipTransform.wavenumbers_to_clip = %N_TO_CLIP
+
+# Parameters for advance/InverseLevelScale:
+# ==============================================================================
+advance/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+advance/InverseLevelScale.name = None
+advance/InverseLevelScale.scales = \
+    [8.824e-05,
+     7.105e-05,
+     0.0001046,
+     0.0001851,
+     0.0007562,
+     0.002625,
+     0.007185,
+     0.01554,
+     0.02886,
+     0.04613,
+     0.07032,
+     0.09901,
+     0.1344,
+     0.175,
+     0.2226,
+     0.2802,
+     0.3438,
+     0.4144,
+     0.4684,
+     0.5257,
+     0.5895,
+     0.6645,
+     0.7505,
+     0.8483,
+     0.9499,
+     1.053,
+     1.159,
+     1.274,
+     1.399,
+     1.552,
+     1.709,
+     1.791]
+
+# Parameters for decoder_model/InverseLevelScale:
+# ==============================================================================
+decoder_model/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+decoder_model/InverseLevelScale.name = None
+decoder_model/InverseLevelScale.scales = \
+    [8.824e-05,
+     7.105e-05,
+     0.0001046,
+     0.0001851,
+     0.0007562,
+     0.002625,
+     0.007185,
+     0.01554,
+     0.02886,
+     0.04613,
+     0.07032,
+     0.09901,
+     0.1344,
+     0.175,
+     0.2226,
+     0.2802,
+     0.3438,
+     0.4144,
+     0.4684,
+     0.5257,
+     0.5895,
+     0.6645,
+     0.7505,
+     0.8483,
+     0.9499,
+     1.053,
+     1.159,
+     1.274,
+     1.399,
+     1.552,
+     1.709,
+     1.791]
+
+# Parameters for encoder_data/InverseLevelScale:
+# ==============================================================================
+encoder_data/InverseLevelScale.keys_to_scale = \
+    ['specific_humidity',
+     'specific_humidity_del2',
+     'specific_humidity_dlat',
+     'specific_humidity_dlon']
+encoder_data/InverseLevelScale.name = None
+encoder_data/InverseLevelScale.scales = \
+    [3.818e-05,
+     5.988e-05,
+     7.108e-05,
+     8.204e-05,
+     8.841e-05,
+     9.115e-05,
+     7.736e-05,
+     7.213e-05,
+     8.284e-05,
+     9.873e-05,
+     0.0001438,
+     0.0002976,
+     0.001013,
+     0.002753,
+     0.006139,
+     0.01182,
+     0.02014,
+     0.04587,
+     0.08568,
+     0.1393,
+     0.2085,
+     0.2993,
+     0.4066,
+     0.4914,
+     0.5848,
+     0.7112,
+     0.8588,
+     0.936,
+     1.013,
+     1.092,
+     1.173,
+     1.257,
+     1.346,
+     1.452,
+     1.58,
+     1.677,
+     1.713]
+
+# Parameters for decoder/InverseShiftAndNormalize:
+# ==============================================================================
+decoder/InverseShiftAndNormalize.global_scale = %CORRECTOR_SCALE
+decoder/InverseShiftAndNormalize.name = None
+decoder/InverseShiftAndNormalize.scales = \
+    {'sim_time': 0.0,
+     't': 28.0,
+     'tracers': {'specific_cloud_ice_water_content': 7.538e-06,
+                 'specific_cloud_liquid_water_content': 1.979e-05,
+                 'specific_humidity': 0.003487},
+     'u': 0.01933,
+     'v': 0.01033,
+     'z': 0.1496}
+decoder/InverseShiftAndNormalize.shifts = \
+    {'sim_time': 0.0,
+     't': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'u': 0.0,
+     'v': 0.0,
+     'z': 0.0}
+
+# Parameters for div_curl_tendency_outputs/InverseShiftAndNormalize:
+# ==============================================================================
+div_curl_tendency_outputs/InverseShiftAndNormalize.global_scale = %GLOBAL_OUT_SCALE
+div_curl_tendency_outputs/InverseShiftAndNormalize.name = None
+div_curl_tendency_outputs/InverseShiftAndNormalize.scales = \
+    {'log_surface_pressure': 0.05008,
+     'sim_time': 0.0,
+     'temperature_variation': 33.85,
+     'tracers': {'specific_cloud_ice_water_content': 4.471e-05,
+                 'specific_cloud_liquid_water_content': 8.884e-05,
+                 'specific_humidity': 0.00608},
+     'u': 0.05839,
+     'v': 0.05138}
+div_curl_tendency_outputs/InverseShiftAndNormalize.shifts = \
+    {'log_surface_pressure': 0.0,
+     'sim_time': 0.0,
+     'temperature_variation': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'u': 0.0,
+     'v': 0.0}
+
+# Parameters for encoder/InverseShiftAndNormalize:
+# ==============================================================================
+encoder/InverseShiftAndNormalize.global_scale = %CORRECTOR_SCALE
+encoder/InverseShiftAndNormalize.name = None
+encoder/InverseShiftAndNormalize.scales = \
+    {'divergence': 0.08243,
+     'log_surface_pressure': 0.1123,
+     'sim_time': 0.0,
+     'temperature_variation': 14.99,
+     'tracers': {'specific_cloud_ice_water_content': 9.111e-06,
+                 'specific_cloud_liquid_water_content': 1.897e-05,
+                 'specific_humidity': 0.003298},
+     'vorticity': 0.2579}
+encoder/InverseShiftAndNormalize.shifts = \
+    {'divergence': 0.0,
+     'log_surface_pressure': 0.0,
+     'sim_time': 0.0,
+     'temperature_variation': 0.0,
+     'tracers': {'specific_cloud_ice_water_content': 0.0,
+                 'specific_cloud_liquid_water_content': 0.0,
+                 'specific_humidity': 0.0},
+     'vorticity': 0.0}
+
+# Parameters for advance/LatitudeFeatures:
+# ==============================================================================
+advance/LatitudeFeatures.name = None
+
+# Parameters for decoder_model/LatitudeFeatures:
+# ==============================================================================
+decoder_model/LatitudeFeatures.name = None
+
+# Parameters for encoder_data/LatitudeFeatures:
+# ==============================================================================
+encoder_data/LatitudeFeatures.name = None
+
+# Parameters for LearnedOrography:
+# ==============================================================================
+LearnedOrography.base_orography_module = @FilteredCustomOrography
+LearnedOrography.correction_scale = 2e-06
+LearnedOrography.name = None
+
+# Parameters for advance/LearnedPositionalFeatures:
+# ==============================================================================
+advance/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+advance/LearnedPositionalFeatures.name = None
+advance/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for decoder_model/LearnedPositionalFeatures:
+# ==============================================================================
+decoder_model/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+decoder_model/LearnedPositionalFeatures.name = None
+decoder_model/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for encoder_data/LearnedPositionalFeatures:
+# ==============================================================================
+encoder_data/LearnedPositionalFeatures.latent_size = %POSITIONAL_LATENT_SIZE
+encoder_data/LearnedPositionalFeatures.name = None
+encoder_data/LearnedPositionalFeatures.scale = 1.0
+
+# Parameters for decoder/LevelScale:
+# ==============================================================================
+decoder/LevelScale.keys_to_scale = ['specific_humidity']
+decoder/LevelScale.name = None
+decoder/LevelScale.scales = \
+    [3.818e-05,
+     5.988e-05,
+     7.108e-05,
+     8.204e-05,
+     8.841e-05,
+     9.115e-05,
+     7.736e-05,
+     7.213e-05,
+     8.284e-05,
+     9.873e-05,
+     0.0001438,
+     0.0002976,
+     0.001013,
+     0.002753,
+     0.006139,
+     0.01182,
+     0.02014,
+     0.04587,
+     0.08568,
+     0.1393,
+     0.2085,
+     0.2993,
+     0.4066,
+     0.4914,
+     0.5848,
+     0.7112,
+     0.8588,
+     0.936,
+     1.013,
+     1.092,
+     1.173,
+     1.257,
+     1.346,
+     1.452,
+     1.58,
+     1.677,
+     1.713]
+
+# Parameters for div_curl_tendency_outputs/LevelScale:
+# ==============================================================================
+div_curl_tendency_outputs/LevelScale.keys_to_scale = ['specific_humidity']
+div_curl_tendency_outputs/LevelScale.name = None
+div_curl_tendency_outputs/LevelScale.scales = \
+    [0.000132,
+     0.0002049,
+     0.000272,
+     0.0004865,
+     0.001441,
+     0.004704,
+     0.01322,
+     0.03009,
+     0.05781,
+     0.09648,
+     0.1486,
+     0.211,
+     0.2892,
+     0.3729,
+     0.471,
+     0.579,
+     0.6898,
+     0.8158,
+     0.9043,
+     0.9958,
+     1.102,
+     1.197,
+     1.308,
+     1.427,
+     1.519,
+     1.571,
+     1.599,
+     1.614,
+     1.629,
+     1.656,
+     1.745,
+     1.841]
+
+# Parameters for encode/LevelScale:
+# ==============================================================================
+encode/LevelScale.keys_to_scale = ['specific_humidity']
+encode/LevelScale.name = None
+encode/LevelScale.scales = \
+    [8.824e-05,
+     7.105e-05,
+     0.0001046,
+     0.0001851,
+     0.0007562,
+     0.002625,
+     0.007185,
+     0.01554,
+     0.02886,
+     0.04613,
+     0.07032,
+     0.09901,
+     0.1344,
+     0.175,
+     0.2226,
+     0.2802,
+     0.3438,
+     0.4144,
+     0.4684,
+     0.5257,
+     0.5895,
+     0.6645,
+     0.7505,
+     0.8483,
+     0.9499,
+     1.053,
+     1.159,
+     1.274,
+     1.399,
+     1.552,
+     1.709,
+     1.791]
+
+# Parameters for advance/MemoryVelocityAndValues:
+# ==============================================================================
+advance/MemoryVelocityAndValues.fields_to_include = None
+advance/MemoryVelocityAndValues.name = None
+
+# Parameters for decode/MlpUniform:
+# ==============================================================================
+decode/MlpUniform.activate_final = False
+decode/MlpUniform.activation = %ACTIVATION
+decode/MlpUniform.b_init = None
+decode/MlpUniform.b_init_final = None
+decode/MlpUniform.name = None
+decode/MlpUniform.num_hidden_layers = 0
+decode/MlpUniform.num_hidden_units = %LAYER_SIZE
+decode/MlpUniform.w_init = None
+decode/MlpUniform.w_init_final = None
+decode/MlpUniform.with_bias = False
+
+# Parameters for encode/MlpUniform:
+# ==============================================================================
+encode/MlpUniform.activate_final = False
+encode/MlpUniform.activation = %ACTIVATION
+encode/MlpUniform.b_init = None
+encode/MlpUniform.b_init_final = None
+encode/MlpUniform.name = None
+encode/MlpUniform.num_hidden_layers = 0
+encode/MlpUniform.num_hidden_units = 0
+encode/MlpUniform.w_init = None
+encode/MlpUniform.w_init_final = None
+encode/MlpUniform.with_bias = True
+
+# Parameters for process/MlpUniform:
+# ==============================================================================
+process/MlpUniform.activate_final = False
+process/MlpUniform.activation = %ACTIVATION
+process/MlpUniform.b_init = None
+process/MlpUniform.b_init_final = None
+process/MlpUniform.name = None
+process/MlpUniform.num_hidden_layers = 3
+process/MlpUniform.num_hidden_units = %LAYER_SIZE
+process/MlpUniform.w_init = None
+process/MlpUniform.w_init_final = None
+process/MlpUniform.with_bias = True
+
+# Parameters for surface_model_decode/MlpUniform:
+# ==============================================================================
+surface_model_decode/MlpUniform.activate_final = False
+surface_model_decode/MlpUniform.activation = %ACTIVATION
+surface_model_decode/MlpUniform.b_init = None
+surface_model_decode/MlpUniform.b_init_final = None
+surface_model_decode/MlpUniform.name = None
+surface_model_decode/MlpUniform.num_hidden_layers = 1
+surface_model_decode/MlpUniform.num_hidden_units = %SURFACE_MODEL_LAYER_SIZE
+surface_model_decode/MlpUniform.w_init = None
+surface_model_decode/MlpUniform.w_init_final = None
+surface_model_decode/MlpUniform.with_bias = False
+
+# Parameters for surface_model_encode/MlpUniform:
+# ==============================================================================
+surface_model_encode/MlpUniform.activate_final = False
+surface_model_encode/MlpUniform.activation = %ACTIVATION
+surface_model_encode/MlpUniform.b_init = None
+surface_model_encode/MlpUniform.b_init_final = None
+surface_model_encode/MlpUniform.name = None
+surface_model_encode/MlpUniform.num_hidden_layers = 0
+surface_model_encode/MlpUniform.num_hidden_units = 0
+surface_model_encode/MlpUniform.w_init = None
+surface_model_encode/MlpUniform.w_init_final = None
+surface_model_encode/MlpUniform.with_bias = True
+
+# Parameters for surface_model_process/MlpUniform:
+# ==============================================================================
+surface_model_process/MlpUniform.activate_final = False
+surface_model_process/MlpUniform.activation = %ACTIVATION
+surface_model_process/MlpUniform.b_init = None
+surface_model_process/MlpUniform.b_init_final = None
+surface_model_process/MlpUniform.name = None
+surface_model_process/MlpUniform.num_hidden_layers = 3
+surface_model_process/MlpUniform.num_hidden_units = %SURFACE_MODEL_LAYER_SIZE
+surface_model_process/MlpUniform.w_init = None
+surface_model_process/MlpUniform.w_init_final = None
+surface_model_process/MlpUniform.with_bias = True
+
+# Parameters for advance/ModalToNodalEmbedding:
+# ==============================================================================
+advance/ModalToNodalEmbedding.modal_to_nodal_features_module = \
+    @embedding_model/CombinedFeatures
+advance/ModalToNodalEmbedding.name = None
+advance/ModalToNodalEmbedding.nodal_mapping_module = @NodalVolumeMapping
+advance/ModalToNodalEmbedding.output_transform_module = @IdentityTransform
+
+# Parameters for sea_model/ModalToNodalEmbedding:
+# ==============================================================================
+sea_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \
+    @sea_model/CombinedFeatures
+sea_model/ModalToNodalEmbedding.name = None
+sea_model/ModalToNodalEmbedding.nodal_mapping_module = @sea_model/NodalMapping
+sea_model/ModalToNodalEmbedding.output_transform_module = @IdentityTransform
+
+# Parameters for MoistPrimitiveEquationsWithCloudMoisture:
+# ==============================================================================
+MoistPrimitiveEquationsWithCloudMoisture.include_vertical_advection = True
+MoistPrimitiveEquationsWithCloudMoisture.name = None
+MoistPrimitiveEquationsWithCloudMoisture.orography_module = @LearnedOrography
+
+# Parameters for advance/NodalLandSeaIceEmbedding:
+# ==============================================================================
+advance/NodalLandSeaIceEmbedding.land_embedding = None
+advance/NodalLandSeaIceEmbedding.name = None
+advance/NodalLandSeaIceEmbedding.sea_embedding = @sea_model/ModalToNodalEmbedding
+advance/NodalLandSeaIceEmbedding.sea_ice_embedding = None
+advance/NodalLandSeaIceEmbedding.static_vars_ds_path = None
+
+# Parameters for NodalMapping:
+# ==============================================================================
+NodalMapping.name = None
+NodalMapping.tower_factory = @EpdTower
+
+# Parameters for sea_model/NodalMapping:
+# ==============================================================================
+sea_model/NodalMapping.name = None
+sea_model/NodalMapping.tower_factory = @surface_model/EpdTower
+
+# Parameters for NodalVolumeMapping:
+# ==============================================================================
+NodalVolumeMapping.name = None
+NodalVolumeMapping.tower_factory = @VerticalConvTower
+
+# Parameters for NullFeatures:
+# ==============================================================================
+NullFeatures.name = None
+
+# Parameters for PerVariableDataFilter:
+# ==============================================================================
+PerVariableDataFilter.name = None
+PerVariableDataFilter.per_variable_filters = \
+    {'divergence': @divergence/DataExponentialFilter,
+     'log_surface_pressure': @lsp/DataExponentialFilter,
+     'sim_time': @DataNoFilter,
+     'temperature_variation': @temperature_variation/DataExponentialFilter,
+     'tracers': {'specific_cloud_ice_water_content': @q/DataExponentialFilter,
+                 'specific_cloud_liquid_water_content': @q/DataExponentialFilter,
+                 'specific_humidity': @q/DataExponentialFilter},
+     'vorticity': @vorticity/DataExponentialFilter}
+
+# Parameters for advance/PressureFeatures:
+# ==============================================================================
+advance/PressureFeatures.name = None
+
+# Parameters for embedding_model/PressureFeatures:
+# ==============================================================================
+embedding_model/PressureFeatures.name = None
+
+# Parameters for primitive_eq_specs_constructor:
+# ==============================================================================
+primitive_eq_specs_constructor.scale = @GET_ATMOSPHERIC_SCALE()
+
+# Parameters for primitive_eq_to_xarray:
+# ==============================================================================
+# None.
+
+# Parameters for PrimitiveToWeatherbenchDecoder:
+# ==============================================================================
+# None.
+
+# Parameters for advance/RadiationFeatures:
+# ==============================================================================
+advance/RadiationFeatures.name = None
+
+# Parameters for decoder_model/RadiationFeatures:
+# ==============================================================================
+decoder_model/RadiationFeatures.name = None
+
+# Parameters for encoder_data/RadiationFeatures:
+# ==============================================================================
+encoder_data/RadiationFeatures.name = None
+
+# Parameters for RealSphericalHarmonicsWithZeroImag:
+# ==============================================================================
+RealSphericalHarmonicsWithZeroImag.base_shape_multiple = None
+RealSphericalHarmonicsWithZeroImag.reverse_einsum_arg_order = None
+RealSphericalHarmonicsWithZeroImag.stacked_fourier_transforms = None
+RealSphericalHarmonicsWithZeroImag.transform_precision = 'tensorfloat32'
+
+# Parameters for orography/RealSphericalHarmonicsWithZeroImag:
+# ==============================================================================
+orography/RealSphericalHarmonicsWithZeroImag.base_shape_multiple = None
+orography/RealSphericalHarmonicsWithZeroImag.reverse_einsum_arg_order = None
+orography/RealSphericalHarmonicsWithZeroImag.stacked_fourier_transforms = None
+orography/RealSphericalHarmonicsWithZeroImag.transform_precision = 'tensorfloat32'
+
+# Parameters for dycore/SequentialStepFilter:
+# ==============================================================================
+dycore/SequentialStepFilter.filter_modules = \
+    (@dycore/ExponentialFilter, @stability/ExponentialFilter)
+dycore/SequentialStepFilter.name = None
+
+# Parameters for ml/SequentialStepFilter:
+# ==============================================================================
+ml/SequentialStepFilter.filter_modules = (@stability/ExponentialFilter,)
+ml/SequentialStepFilter.name = None
+
+# Parameters for advance/SequentialTransform:
+# ==============================================================================
+advance/SequentialTransform.name = None
+advance/SequentialTransform.transform_modules = \
+    (@advance/ShiftAndNormalize,
+     @advance/InverseLevelScale,
+     @advance/TruncateSigmaLevels,
+     @SoftClip)
+
+# Parameters for decoder/SequentialTransform:
+# ==============================================================================
+decoder/SequentialTransform.name = None
+decoder/SequentialTransform.transform_modules = \
+    (@decoder/InverseShiftAndNormalize, @decoder/LevelScale)
+
+# Parameters for decoder_model/SequentialTransform:
+# ==============================================================================
+decoder_model/SequentialTransform.name = None
+decoder_model/SequentialTransform.transform_modules = \
+    (@decoder_model/ShiftAndNormalize,
+     @decoder_model/InverseLevelScale,
+     @decoder_model/TruncateSigmaLevels)
+
+# Parameters for div_curl_tendency_outputs/SequentialTransform:
+# ==============================================================================
+div_curl_tendency_outputs/SequentialTransform.name = None
+div_curl_tendency_outputs/SequentialTransform.transform_modules = \
+    (@div_curl_tendency_outputs/InverseShiftAndNormalize,
+     @div_curl_tendency_outputs/LevelScale)
+
+# Parameters for encode/SequentialTransform:
+# ==============================================================================
+encode/SequentialTransform.name = None
+encode/SequentialTransform.transform_modules = \
+    (@encoder/InverseShiftAndNormalize, @encode/LevelScale)
+
+# Parameters for encoder_data/SequentialTransform:
+# ==============================================================================
+encoder_data/SequentialTransform.name = None
+encoder_data/SequentialTransform.transform_modules = \
+    (@encoder_data/ShiftAndNormalize, @encoder_data/InverseLevelScale)
+
+# Parameters for advance/ShiftAndNormalize:
+# ==============================================================================
+advance/ShiftAndNormalize.features_to_exclude = ()
+advance/ShiftAndNormalize.global_scale = None
+advance/ShiftAndNormalize.name = None
+advance/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3057,
+     'divergence': 0.08243,
+     'divergence_del2': 629.5,
+     'divergence_dlat': 6.037,
+     'divergence_dlon': 4.505,
+     'geopotential_at_surface': 0.009482,
+     'geopotential_at_surface_del2': 11.13,
+     'geopotential_at_surface_dlat': 0.1337,
+     'geopotential_at_surface_dlon': 0.1168,
+     'land_sea_mask': 0.4503,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1123,
+     'log_surface_pressure_del2': 120.6,
+     'log_surface_pressure_dlat': 1.41,
+     'log_surface_pressure_dlon': 1.285,
+     'memory_divergence': 0.08243,
+     'memory_log_surface_pressure': 0.1123,
+     'memory_specific_cloud_ice_water_content': 9.111e-06,
+     'memory_specific_cloud_liquid_water_content': 1.897e-05,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': 14.99,
+     'memory_u': 0.01485,
+     'memory_v': 0.01017,
+     'memory_vorticity': 0.2579,
+     'pressure': 1.644,
+     'radiation': 0.2864,
+     'sea_ice_cover': 0.387,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7057,
+     'specific_cloud_ice_water_content': 9.111e-06,
+     'specific_cloud_ice_water_content_del2': 0.03864,
+     'specific_cloud_ice_water_content_dlat': 0.0003813,
+     'specific_cloud_ice_water_content_dlon': 0.0003149,
+     'specific_cloud_liquid_water_content': 1.897e-05,
+     'specific_cloud_liquid_water_content_del2': 0.08082,
+     'specific_cloud_liquid_water_content_dlat': 0.0008049,
+     'specific_cloud_liquid_water_content_dlon': 0.0006388,
+     'specific_humidity': 0.003298,
+     'specific_humidity_del2': 2.389,
+     'specific_humidity_dlat': 0.02819,
+     'specific_humidity_dlon': 0.021,
+     'surface_embedding': 1.0,
+     'temperature_variation': 14.99,
+     'temperature_variation_del2': 5793.0,
+     'temperature_variation_dlat': 75.47,
+     'temperature_variation_dlon': 64.98,
+     'u': 0.01485,
+     'u_del2': 13.13,
+     'u_dlat': 1.223,
+     'u_dlon': 0.1378,
+     'v': 0.01017,
+     'v_del2': 11.26,
+     'v_dlat': 0.764,
+     'v_dlon': 0.1781,
+     'vorticity': 0.2579,
+     'vorticity_del2': 1189.0,
+     'vorticity_dlat': 11.67,
+     'vorticity_dlon': 9.514}
+advance/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.639,
+     'divergence': -0.0,
+     'divergence_del2': -0.009,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.069,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': 0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.716,
+     'log_surface_pressure_del2': 0.824,
+     'log_surface_pressure_dlat': 0.13,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.716,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': -4.946,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.214,
+     'sea_ice_cover': 0.24,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': 0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.946,
+     'temperature_variation_del2': 47.3,
+     'temperature_variation_dlat': 4.913,
+     'temperature_variation_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.187,
+     'u_dlat': 0.002,
+     'u_dlon': -0.0,
+     'v': -0.0,
+     'v_del2': -0.002,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.37,
+     'vorticity_dlat': 0.042,
+     'vorticity_dlon': -0.0}
+
+# Parameters for decoder_model/ShiftAndNormalize:
+# ==============================================================================
+decoder_model/ShiftAndNormalize.features_to_exclude = ()
+decoder_model/ShiftAndNormalize.global_scale = None
+decoder_model/ShiftAndNormalize.name = None
+decoder_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3057,
+     'divergence': 0.08243,
+     'divergence_del2': 629.5,
+     'divergence_dlat': 6.037,
+     'divergence_dlon': 4.505,
+     'geopotential_at_surface': 0.009482,
+     'geopotential_at_surface_del2': 11.13,
+     'geopotential_at_surface_dlat': 0.1337,
+     'geopotential_at_surface_dlon': 0.1168,
+     'land_sea_mask': 0.4503,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1123,
+     'log_surface_pressure_del2': 120.6,
+     'log_surface_pressure_dlat': 1.41,
+     'log_surface_pressure_dlon': 1.285,
+     'memory_divergence': 0.08243,
+     'memory_log_surface_pressure': 0.1123,
+     'memory_specific_cloud_ice_water_content': 9.111e-06,
+     'memory_specific_cloud_liquid_water_content': 1.897e-05,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': 14.99,
+     'memory_u': 0.01485,
+     'memory_v': 0.01017,
+     'memory_vorticity': 0.2579,
+     'pressure': 1.644,
+     'radiation': 0.2864,
+     'sea_ice_cover': 0.387,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7057,
+     'specific_cloud_ice_water_content': 9.111e-06,
+     'specific_cloud_ice_water_content_del2': 0.03864,
+     'specific_cloud_ice_water_content_dlat': 0.0003813,
+     'specific_cloud_ice_water_content_dlon': 0.0003149,
+     'specific_cloud_liquid_water_content': 1.897e-05,
+     'specific_cloud_liquid_water_content_del2': 0.08082,
+     'specific_cloud_liquid_water_content_dlat': 0.0008049,
+     'specific_cloud_liquid_water_content_dlon': 0.0006388,
+     'specific_humidity': 0.003298,
+     'specific_humidity_del2': 2.389,
+     'specific_humidity_dlat': 0.02819,
+     'specific_humidity_dlon': 0.021,
+     'surface_embedding': 1.0,
+     'temperature_variation': 14.99,
+     'temperature_variation_del2': 5793.0,
+     'temperature_variation_dlat': 75.47,
+     'temperature_variation_dlon': 64.98,
+     'u': 0.01485,
+     'u_del2': 13.13,
+     'u_dlat': 1.223,
+     'u_dlon': 0.1378,
+     'v': 0.01017,
+     'v_del2': 11.26,
+     'v_dlat': 0.764,
+     'v_dlon': 0.1781,
+     'vorticity': 0.2579,
+     'vorticity_del2': 1189.0,
+     'vorticity_dlat': 11.67,
+     'vorticity_dlon': 9.514}
+decoder_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.639,
+     'divergence': -0.0,
+     'divergence_del2': -0.009,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.069,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': 0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.716,
+     'log_surface_pressure_del2': 0.824,
+     'log_surface_pressure_dlat': 0.13,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.716,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': -4.946,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.214,
+     'sea_ice_cover': 0.24,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': 0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.946,
+     'temperature_variation_del2': 47.3,
+     'temperature_variation_dlat': 4.913,
+     'temperature_variation_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.187,
+     'u_dlat': 0.002,
+     'u_dlon': -0.0,
+     'v': -0.0,
+     'v_del2': -0.002,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.37,
+     'vorticity_dlat': 0.042,
+     'vorticity_dlon': -0.0}
+
+# Parameters for embedding_model/ShiftAndNormalize:
+# ==============================================================================
+embedding_model/ShiftAndNormalize.features_to_exclude = ()
+embedding_model/ShiftAndNormalize.global_scale = None
+embedding_model/ShiftAndNormalize.name = None
+embedding_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3057,
+     'divergence': 0.08243,
+     'divergence_del2': 629.5,
+     'divergence_dlat': 6.037,
+     'divergence_dlon': 4.505,
+     'geopotential_at_surface': 0.009482,
+     'geopotential_at_surface_del2': 11.13,
+     'geopotential_at_surface_dlat': 0.1337,
+     'geopotential_at_surface_dlon': 0.1168,
+     'land_sea_mask': 0.4503,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1123,
+     'log_surface_pressure_del2': 120.6,
+     'log_surface_pressure_dlat': 1.41,
+     'log_surface_pressure_dlon': 1.285,
+     'memory_divergence': 0.08243,
+     'memory_log_surface_pressure': 0.1123,
+     'memory_specific_cloud_ice_water_content': 9.111e-06,
+     'memory_specific_cloud_liquid_water_content': 1.897e-05,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': 14.99,
+     'memory_u': 0.01485,
+     'memory_v': 0.01017,
+     'memory_vorticity': 0.2579,
+     'pressure': 1.644,
+     'radiation': 0.2864,
+     'sea_ice_cover': 0.387,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7057,
+     'specific_cloud_ice_water_content': 9.111e-06,
+     'specific_cloud_ice_water_content_del2': 0.03864,
+     'specific_cloud_ice_water_content_dlat': 0.0003813,
+     'specific_cloud_ice_water_content_dlon': 0.0003149,
+     'specific_cloud_liquid_water_content': 1.897e-05,
+     'specific_cloud_liquid_water_content_del2': 0.08082,
+     'specific_cloud_liquid_water_content_dlat': 0.0008049,
+     'specific_cloud_liquid_water_content_dlon': 0.0006388,
+     'specific_humidity': 0.003298,
+     'specific_humidity_del2': 2.389,
+     'specific_humidity_dlat': 0.02819,
+     'specific_humidity_dlon': 0.021,
+     'surface_embedding': 1.0,
+     'temperature_variation': 14.99,
+     'temperature_variation_del2': 5793.0,
+     'temperature_variation_dlat': 75.47,
+     'temperature_variation_dlon': 64.98,
+     'u': 0.01485,
+     'u_del2': 13.13,
+     'u_dlat': 1.223,
+     'u_dlon': 0.1378,
+     'v': 0.01017,
+     'v_del2': 11.26,
+     'v_dlat': 0.764,
+     'v_dlon': 0.1781,
+     'vorticity': 0.2579,
+     'vorticity_del2': 1189.0,
+     'vorticity_dlat': 11.67,
+     'vorticity_dlon': 9.514}
+embedding_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.639,
+     'divergence': -0.0,
+     'divergence_del2': -0.009,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.069,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': 0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.716,
+     'log_surface_pressure_del2': 0.824,
+     'log_surface_pressure_dlat': 0.13,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.716,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': -4.946,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.214,
+     'sea_ice_cover': 0.24,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': 0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.946,
+     'temperature_variation_del2': 47.3,
+     'temperature_variation_dlat': 4.913,
+     'temperature_variation_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.187,
+     'u_dlat': 0.002,
+     'u_dlon': -0.0,
+     'v': -0.0,
+     'v_del2': -0.002,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.37,
+     'vorticity_dlat': 0.042,
+     'vorticity_dlon': -0.0}
+
+# Parameters for encoder_data/ShiftAndNormalize:
+# ==============================================================================
+encoder_data/ShiftAndNormalize.features_to_exclude = ()
+encoder_data/ShiftAndNormalize.global_scale = None
+encoder_data/ShiftAndNormalize.name = None
+encoder_data/ShiftAndNormalize.scales = \
+    {'cos_latitude': 0.3057,
+     'geopotential_at_surface': 0.009482,
+     'geopotential_at_surface_del2': 11.13,
+     'geopotential_at_surface_dlat': 0.1337,
+     'geopotential_at_surface_dlon': 0.1168,
+     'land_sea_mask': 0.4503,
+     'learned_positional_features': 1.0,
+     'radiation': 0.2864,
+     'sin_latitude': 0.7057,
+     'specific_cloud_ice_water_content': 7.538e-06,
+     'specific_cloud_ice_water_content_del2': 0.03283,
+     'specific_cloud_ice_water_content_dlat': 0.0003762,
+     'specific_cloud_ice_water_content_dlon': 0.0002613,
+     'specific_cloud_liquid_water_content': 1.979e-05,
+     'specific_cloud_liquid_water_content_del2': 0.08809,
+     'specific_cloud_liquid_water_content_dlat': 0.001038,
+     'specific_cloud_liquid_water_content_dlon': 0.000684,
+     'specific_humidity': 0.003487,
+     'specific_humidity_del2': 2.381,
+     'specific_humidity_dlat': 0.03179,
+     'specific_humidity_dlon': 0.02059,
+     't': 28.0,
+     't_del2': 4397.0,
+     't_dlat': 64.14,
+     't_dlon': 50.07,
+     'u': 0.01933,
+     'u_del2': 12.27,
+     'u_dlat': 0.1743,
+     'u_dlon': 0.1296,
+     'v': 0.01033,
+     'v_del2': 12.08,
+     'v_dlat': 0.1152,
+     'v_dlon': 0.1617,
+     'z': 0.1496,
+     'z_del2': 0.332,
+     'z_dlat': 0.01598,
+     'z_dlon': 0.008823}
+encoder_data/ShiftAndNormalize.shifts = \
+    {'cos_latitude': 0.639,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.069,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': 0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'radiation': 0.214,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': -0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': -0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     't': 247.07,
+     't_del2': 35.964,
+     't_dlat': 1.763,
+     't_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.045,
+     'u_dlat': -0.0,
+     'u_dlon': 0.0,
+     'v': 0.0,
+     'v_del2': 0.005,
+     'v_dlat': 0.0,
+     'v_dlon': 0.0,
+     'z': 0.145,
+     'z_del2': 0.011,
+     'z_dlat': 0.001,
+     'z_dlon': 0.0}
+
+# Parameters for sea_model/ShiftAndNormalize:
+# ==============================================================================
+sea_model/ShiftAndNormalize.features_to_exclude = ()
+sea_model/ShiftAndNormalize.global_scale = None
+sea_model/ShiftAndNormalize.name = None
+sea_model/ShiftAndNormalize.scales = \
+    {'CNN1D_0': 0.1,
+     'CNN1D_1': 0.1,
+     'CNN1D_10': 0.1,
+     'CNN1D_11': 0.1,
+     'CNN1D_12': 0.1,
+     'CNN1D_13': 0.1,
+     'CNN1D_14': 0.1,
+     'CNN1D_15': 0.1,
+     'CNN1D_16': 0.1,
+     'CNN1D_17': 0.1,
+     'CNN1D_18': 0.1,
+     'CNN1D_19': 0.1,
+     'CNN1D_2': 0.1,
+     'CNN1D_20': 0.1,
+     'CNN1D_21': 0.1,
+     'CNN1D_22': 0.1,
+     'CNN1D_23': 0.1,
+     'CNN1D_24': 0.1,
+     'CNN1D_25': 0.1,
+     'CNN1D_26': 0.1,
+     'CNN1D_27': 0.1,
+     'CNN1D_28': 0.1,
+     'CNN1D_29': 0.1,
+     'CNN1D_3': 0.1,
+     'CNN1D_30': 0.1,
+     'CNN1D_31': 0.1,
+     'CNN1D_32': 0.1,
+     'CNN1D_33': 0.1,
+     'CNN1D_34': 0.1,
+     'CNN1D_35': 0.1,
+     'CNN1D_36': 0.1,
+     'CNN1D_37': 0.1,
+     'CNN1D_38': 0.1,
+     'CNN1D_39': 0.1,
+     'CNN1D_4': 0.1,
+     'CNN1D_40': 0.1,
+     'CNN1D_41': 0.1,
+     'CNN1D_42': 0.1,
+     'CNN1D_43': 0.1,
+     'CNN1D_44': 0.1,
+     'CNN1D_45': 0.1,
+     'CNN1D_46': 0.1,
+     'CNN1D_47': 0.1,
+     'CNN1D_48': 0.1,
+     'CNN1D_49': 0.1,
+     'CNN1D_5': 0.1,
+     'CNN1D_50': 0.1,
+     'CNN1D_51': 0.1,
+     'CNN1D_52': 0.1,
+     'CNN1D_53': 0.1,
+     'CNN1D_54': 0.1,
+     'CNN1D_55': 0.1,
+     'CNN1D_56': 0.1,
+     'CNN1D_57': 0.1,
+     'CNN1D_58': 0.1,
+     'CNN1D_59': 0.1,
+     'CNN1D_6': 0.1,
+     'CNN1D_60': 0.1,
+     'CNN1D_61': 0.1,
+     'CNN1D_62': 0.1,
+     'CNN1D_63': 0.1,
+     'CNN1D_7': 0.1,
+     'CNN1D_8': 0.1,
+     'CNN1D_9': 0.1,
+     'cos_latitude': 0.3057,
+     'divergence': 0.08243,
+     'divergence_del2': 629.5,
+     'divergence_dlat': 6.037,
+     'divergence_dlon': 4.505,
+     'geopotential_at_surface': 0.009482,
+     'geopotential_at_surface_del2': 11.13,
+     'geopotential_at_surface_dlat': 0.1337,
+     'geopotential_at_surface_dlon': 0.1168,
+     'land_sea_mask': 0.4503,
+     'learned_positional_features': 1.0,
+     'log_surface_pressure': 0.1123,
+     'log_surface_pressure_del2': 120.6,
+     'log_surface_pressure_dlat': 1.41,
+     'log_surface_pressure_dlon': 1.285,
+     'memory_divergence': 0.08243,
+     'memory_log_surface_pressure': 0.1123,
+     'memory_specific_cloud_ice_water_content': 9.111e-06,
+     'memory_specific_cloud_liquid_water_content': 1.897e-05,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': 14.99,
+     'memory_u': 0.01485,
+     'memory_v': 0.01017,
+     'memory_vorticity': 0.2579,
+     'pressure': 1.644,
+     'radiation': 0.2864,
+     'sea_ice_cover': 0.387,
+     'sea_surface_temperature': 11.93,
+     'sin_latitude': 0.7057,
+     'specific_cloud_ice_water_content': 9.111e-06,
+     'specific_cloud_ice_water_content_del2': 0.03864,
+     'specific_cloud_ice_water_content_dlat': 0.0003813,
+     'specific_cloud_ice_water_content_dlon': 0.0003149,
+     'specific_cloud_liquid_water_content': 1.897e-05,
+     'specific_cloud_liquid_water_content_del2': 0.08082,
+     'specific_cloud_liquid_water_content_dlat': 0.0008049,
+     'specific_cloud_liquid_water_content_dlon': 0.0006388,
+     'specific_humidity': 0.003298,
+     'specific_humidity_del2': 2.389,
+     'specific_humidity_dlat': 0.02819,
+     'specific_humidity_dlon': 0.021,
+     'surface_embedding': 1.0,
+     'temperature_variation': 14.99,
+     'temperature_variation_del2': 5793.0,
+     'temperature_variation_dlat': 75.47,
+     'temperature_variation_dlon': 64.98,
+     'u': 0.01485,
+     'u_del2': 13.13,
+     'u_dlat': 1.223,
+     'u_dlon': 0.1378,
+     'v': 0.01017,
+     'v_del2': 11.26,
+     'v_dlat': 0.764,
+     'v_dlon': 0.1781,
+     'vorticity': 0.2579,
+     'vorticity_del2': 1189.0,
+     'vorticity_dlat': 11.67,
+     'vorticity_dlon': 9.514}
+sea_model/ShiftAndNormalize.shifts = \
+    {'CNN1D_0': 0.0,
+     'CNN1D_1': 0.0,
+     'CNN1D_10': 0.0,
+     'CNN1D_11': 0.0,
+     'CNN1D_12': 0.0,
+     'CNN1D_13': 0.0,
+     'CNN1D_14': 0.0,
+     'CNN1D_15': 0.0,
+     'CNN1D_16': 0.0,
+     'CNN1D_17': 0.0,
+     'CNN1D_18': 0.0,
+     'CNN1D_19': 0.0,
+     'CNN1D_2': 0.0,
+     'CNN1D_20': 0.0,
+     'CNN1D_21': 0.0,
+     'CNN1D_22': 0.0,
+     'CNN1D_23': 0.0,
+     'CNN1D_24': 0.0,
+     'CNN1D_25': 0.0,
+     'CNN1D_26': 0.0,
+     'CNN1D_27': 0.0,
+     'CNN1D_28': 0.0,
+     'CNN1D_29': 0.0,
+     'CNN1D_3': 0.0,
+     'CNN1D_30': 0.0,
+     'CNN1D_31': 0.0,
+     'CNN1D_32': 0.0,
+     'CNN1D_33': 0.0,
+     'CNN1D_34': 0.0,
+     'CNN1D_35': 0.0,
+     'CNN1D_36': 0.0,
+     'CNN1D_37': 0.0,
+     'CNN1D_38': 0.0,
+     'CNN1D_39': 0.0,
+     'CNN1D_4': 0.0,
+     'CNN1D_40': 0.0,
+     'CNN1D_41': 0.0,
+     'CNN1D_42': 0.0,
+     'CNN1D_43': 0.0,
+     'CNN1D_44': 0.0,
+     'CNN1D_45': 0.0,
+     'CNN1D_46': 0.0,
+     'CNN1D_47': 0.0,
+     'CNN1D_48': 0.0,
+     'CNN1D_49': 0.0,
+     'CNN1D_5': 0.0,
+     'CNN1D_50': 0.0,
+     'CNN1D_51': 0.0,
+     'CNN1D_52': 0.0,
+     'CNN1D_53': 0.0,
+     'CNN1D_54': 0.0,
+     'CNN1D_55': 0.0,
+     'CNN1D_56': 0.0,
+     'CNN1D_57': 0.0,
+     'CNN1D_58': 0.0,
+     'CNN1D_59': 0.0,
+     'CNN1D_6': 0.0,
+     'CNN1D_60': 0.0,
+     'CNN1D_61': 0.0,
+     'CNN1D_62': 0.0,
+     'CNN1D_63': 0.0,
+     'CNN1D_7': 0.0,
+     'CNN1D_8': 0.0,
+     'CNN1D_9': 0.0,
+     'cos_latitude': 0.639,
+     'divergence': -0.0,
+     'divergence_del2': -0.009,
+     'divergence_dlat': -0.001,
+     'divergence_dlon': 0.0,
+     'geopotential_at_surface': 0.004,
+     'geopotential_at_surface_del2': -0.069,
+     'geopotential_at_surface_dlat': -0.01,
+     'geopotential_at_surface_dlon': 0.0,
+     'land_sea_mask': 0.335,
+     'learned_positional_features': 0.0,
+     'log_surface_pressure': 1.716,
+     'log_surface_pressure_del2': 0.824,
+     'log_surface_pressure_dlat': 0.13,
+     'log_surface_pressure_dlon': 0.0,
+     'memory_divergence': -0.0,
+     'memory_log_surface_pressure': 1.716,
+     'memory_specific_cloud_ice_water_content': 0.0,
+     'memory_specific_cloud_liquid_water_content': 0.0,
+     'memory_specific_humidity': 0.003298,
+     'memory_temperature_variation': -4.946,
+     'memory_u': 0.007,
+     'memory_v': -0.0,
+     'memory_vorticity': -0.002,
+     'pressure': 2.796,
+     'radiation': 0.214,
+     'sea_ice_cover': 0.24,
+     'sea_surface_temperature': 285.14,
+     'sin_latitude': -0.0,
+     'specific_cloud_ice_water_content': 0.0,
+     'specific_cloud_ice_water_content_del2': 0.0,
+     'specific_cloud_ice_water_content_dlat': 0.0,
+     'specific_cloud_ice_water_content_dlon': -0.0,
+     'specific_cloud_liquid_water_content': 0.0,
+     'specific_cloud_liquid_water_content_del2': -0.0,
+     'specific_cloud_liquid_water_content_dlat': 0.0,
+     'specific_cloud_liquid_water_content_dlon': 0.0,
+     'specific_humidity': 0.0,
+     'specific_humidity_del2': 0.0,
+     'specific_humidity_dlat': 0.0,
+     'specific_humidity_dlon': 0.0,
+     'surface_embedding': 0.0,
+     'temperature_variation': -4.946,
+     'temperature_variation_del2': 47.3,
+     'temperature_variation_dlat': 4.913,
+     'temperature_variation_dlon': -0.0,
+     'u': 0.007,
+     'u_del2': 0.187,
+     'u_dlat': 0.002,
+     'u_dlon': -0.0,
+     'v': -0.0,
+     'v_del2': -0.002,
+     'v_dlat': -0.0,
+     'v_dlon': 0.0,
+     'vorticity': -0.002,
+     'vorticity_del2': -0.37,
+     'vorticity_dlat': 0.042,
+     'vorticity_dlon': -0.0}
+
+# Parameters for SigmaCoordinatesEquidistant:
+# ==============================================================================
+SigmaCoordinatesEquidistant.layers = %N_SIGMA_LAYERS
+
+# Parameters for custom_corrds/SigmaCoordinatesEquidistant:
+# ==============================================================================
+custom_corrds/SigmaCoordinatesEquidistant.layers = %N_SIGMA_LAYERS
+
+# Parameters for advance/SoftClip:
+# ==============================================================================
+advance/SoftClip.hinge_softness = 1.0
+advance/SoftClip.max_value = 16
+advance/SoftClip.name = None
+
+# Parameters for StochasticModularStepModel:
+# ==============================================================================
+StochasticModularStepModel.advance_module = @StochasticPhysicsParameterizationStep
+StochasticModularStepModel.decoder_module = \
+    @DimensionalLearnedPrimitiveToWeatherbenchDecoder
+StochasticModularStepModel.encoder_module = \
+    @DimensionalLearnedWeatherbenchToPrimitiveWithMemoryEncoder
+StochasticModularStepModel.forcing_module = @DynamicDataForcing
+StochasticModularStepModel.name = None
+
+# Parameters for StochasticPhysicsParameterizationStep:
+# ==============================================================================
+StochasticPhysicsParameterizationStep.checkpoint_substep = False
+StochasticPhysicsParameterizationStep.corrector_module = %CORRECTOR_MODULE
+StochasticPhysicsParameterizationStep.name = None
+StochasticPhysicsParameterizationStep.num_substeps = %NUM_SUBSTEPS
+StochasticPhysicsParameterizationStep.physics_parameterization_module = \
+    @DivCurlNeuralParameterization
+StochasticPhysicsParameterizationStep.randomness_module = @ZerosRandomField
+
+# Parameters for advance/ToModalDiffOperators:
+# ==============================================================================
+advance/ToModalDiffOperators.name = None
+
+# Parameters for encoder_data/ToModalDiffOperators:
+# ==============================================================================
+encoder_data/ToModalDiffOperators.name = None
+
+# Parameters for with_grads/ToModalDiffOperators:
+# ==============================================================================
+with_grads/ToModalDiffOperators.name = None
+
+# Parameters for trajectory_from_step:
+# ==============================================================================
+trajectory_from_step.checkpoint_multistep = False
+trajectory_from_step.checkpoint_post_process = True
+trajectory_from_step.checkpoint_step = True
+
+# Parameters for advance/TruncateSigmaLevels:
+# ==============================================================================
+advance/TruncateSigmaLevels.name = None
+advance/TruncateSigmaLevels.sigma_ranges = \
+    {'divergence': (0.0, 1),
+     'geopotential': (0.0, 1),
+     'specific_cloud_ice_water_content': (0.0, 1),
+     'specific_cloud_liquid_water_content': (0.0, 1),
+     'specific_humidity': (0.0, 1),
+     't': (0.0, 1),
+     'temperature_variation': (0.0, 1),
+     'u': (0.0, 1),
+     'u_component_of_wind': (0.0, 1),
+     'v': (0.0, 1),
+     'v_component_of_wind': (0.0, 1),
+     'vorticity': (0.0, 1),
+     'z': (0.0, 1)}
+
+# Parameters for decoder_model/TruncateSigmaLevels:
+# ==============================================================================
+decoder_model/TruncateSigmaLevels.name = None
+decoder_model/TruncateSigmaLevels.sigma_ranges = \
+    {'divergence': (0.0, 1),
+     'geopotential': (0.0, 1),
+     'specific_cloud_ice_water_content': (0.0, 1),
+     'specific_cloud_liquid_water_content': (0.0, 1),
+     'specific_humidity': (0.0, 1),
+     't': (0.0, 1),
+     'temperature_variation': (0.0, 1),
+     'u': (0.0, 1),
+     'u_component_of_wind': (0.0, 1),
+     'v': (0.0, 1),
+     'v_component_of_wind': (0.0, 1),
+     'vorticity': (0.0, 1),
+     'z': (0.0, 1)}
+
+# Parameters for advance/VelocityAndPrognostics:
+# ==============================================================================
+advance/VelocityAndPrognostics.compute_gradients_module = @ToModalDiffOperators
+advance/VelocityAndPrognostics.fields_to_include = \
+    ['divergence',
+     'vorticity',
+     'u',
+     'v',
+     'temperature_variation',
+     'log_surface_pressure',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+advance/VelocityAndPrognostics.name = None
+
+# Parameters for embedding_model/VelocityAndPrognostics:
+# ==============================================================================
+embedding_model/VelocityAndPrognostics.fields_to_include = \
+    ['divergence',
+     'vorticity',
+     'u',
+     'v',
+     'temperature_variation',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+embedding_model/VelocityAndPrognostics.name = None
+
+# Parameters for encoder_data/VelocityAndPrognostics:
+# ==============================================================================
+encoder_data/VelocityAndPrognostics.compute_gradients_module = \
+    @ToModalDiffOperators
+encoder_data/VelocityAndPrognostics.fields_to_include = \
+    ['u',
+     'v',
+     't',
+     'z',
+     'specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content']
+encoder_data/VelocityAndPrognostics.name = None
+
+# Parameters for model/VelocityAndPrognostics:
+# ==============================================================================
+model/VelocityAndPrognostics.fields_to_include = None
+model/VelocityAndPrognostics.name = None
+
+# Parameters for VerticalConvTower:
+# ==============================================================================
+VerticalConvTower.activate_final = False
+VerticalConvTower.activation = %ACTIVATION
+VerticalConvTower.channels = [64, 64, 64, 64]
+VerticalConvTower.checkpoint_tower = True
+VerticalConvTower.kernel_shape = 5
+VerticalConvTower.name = None
+VerticalConvTower.with_bias = True
+
+# Parameters for WhirlModel:
+# ==============================================================================
+WhirlModel.from_xarray_fn = @xarray_to_state_and_dynamic_covariate_data
+WhirlModel.model_cls = @StochasticModularStepModel
+WhirlModel.to_xarray_fn = @data_to_xarray_with_renaming
+
+# Parameters for xarray_to_data_with_renaming:
+# ==============================================================================
+xarray_to_data_with_renaming.renaming_dict = \
+    {'geopotential': 'z',
+     'latitude': 'lat',
+     'longitude': 'lon',
+     'temperature': 't',
+     'u_component_of_wind': 'u',
+     'v_component_of_wind': 'v'}
+xarray_to_data_with_renaming.xarray_to_data_fn = @xarray_to_weatherbench_data
+
+# Parameters for xarray_to_dynamic_covariate_data:
+# ==============================================================================
+xarray_to_dynamic_covariate_data.covariates_to_include = \
+    ('sea_ice_cover', 'sea_surface_temperature')
+
+# Parameters for xarray_to_state_and_dynamic_covariate_data:
+# ==============================================================================
+xarray_to_state_and_dynamic_covariate_data.values = 'values'
+xarray_to_state_and_dynamic_covariate_data.xarray_to_dynamic_covariate_data_fn = \
+    @xarray_to_dynamic_covariate_data
+xarray_to_state_and_dynamic_covariate_data.xarray_to_state_data_fn = \
+    @xarray_to_data_with_renaming
+
+# Parameters for xarray_to_weatherbench_data:
+# ==============================================================================
+xarray_to_weatherbench_data.diagnostics_to_include = ()
+xarray_to_weatherbench_data.tracers_to_include = \
+    ('specific_humidity',
+     'specific_cloud_liquid_water_content',
+     'specific_cloud_ice_water_content')
+
+# Parameters for ZerosRandomField:
+# ==============================================================================
+ZerosRandomField.prefer_nodal = True

model/reference_code/reader.py

@@ -0,0 +1,624 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# pylint: disable=line-too-long
+# pyformat: disable
+"""Xarray based readers for feeding time-series into tf.data."""
+# pyformat: enable
+from __future__ import annotations
+from collections import abc
+import concurrent.futures
+import dataclasses
+import logging
+import math
+import random
+from typing import Callable, Optional, TypeVar
+
+import numpy as np
+import tensorflow as tf
+import xarray
+
+
+# pylint: disable=logging-fstring-interpolation
+
+
+def _xarray_bytes_per_element(
+    source: xarray.Dataset, exclude_dims: set[str]
+) -> int:
+  bytes_per_element = 0
+  for variable in source.values():
+    items_per_element = math.prod(
+        size for dim, size in variable.sizes.items() if dim not in exclude_dims
+    )
+    bytes_per_element += variable.dtype.itemsize * items_per_element
+  return bytes_per_element
+
+
+def _calculate_block_size(
+    source: xarray.Dataset,
+    block_dims: list[str],
+    bytes_per_request: float,
+    min_elements_per_request: int = 1,
+) -> int:
+  """Calculate the size of blocks to read simultaneously from disk."""
+  bytes_per_element = _xarray_bytes_per_element(source, set(block_dims))
+  elements_per_request = round(bytes_per_request / bytes_per_element)
+  max_elements = math.prod(source.sizes[dim] for dim in block_dims)
+  elements_per_request = min(
+      max(elements_per_request, min_elements_per_request), max_elements
+  )
+  return elements_per_request
+
+
+def _iterate_windowed_block_slices(
+    sample_size: int,
+    total_size: int,
+    block_size: int,
+    stride_between_samples: int = 1,
+    output_window_stride: int = 1,
+    first_sample_offset: int = 0,
+) -> abc.Iterator[slice]:
+  """Yields slices for every block needed to generate windowed samples.
+
+  Args:
+    sample_size: size of each sample.
+    total_size: total size of the dimension being sampled along.
+    block_size: desired size of blocks to read from disk.
+    stride_between_samples: shift between starts of sampled windows.
+    output_window_stride: shift between samples within a window.
+    first_sample_offset: offset of the first sample.
+
+  Yields:
+    Slice objects with integer bounds for each block.
+  """
+  assert stride_between_samples >= 1
+  assert output_window_stride >= 1
+  assert first_sample_offset >= 0
+
+  sample_input_size = (
+      range(0, sample_size * output_window_stride, output_window_stride)[-1] + 1
+  )
+  assert 0 < sample_input_size <= block_size <= total_size
+
+  sample_stop = 0  # unused
+
+  # first block
+  block_start = first_sample_offset
+  block_stop = first_sample_offset + block_size
+
+  # iterate through all slices, in order
+  for start in range(
+      first_sample_offset,
+      total_size - sample_input_size + 1,
+      stride_between_samples,
+  ):
+    prev_sample_stop = sample_stop
+    sample_stop = start + sample_input_size
+
+    if sample_stop > block_stop:
+      # yield previous block
+      assert prev_sample_stop > 0
+      yield slice(block_start, prev_sample_stop)
+
+      # begin new block
+      block_start = start
+      block_stop = start + block_size
+
+  if sample_stop > block_start:
+    # yield the final block
+    yield slice(block_start, sample_stop)
+
+
+def _drop_static_vars(dataset: xarray.Dataset) -> xarray.Dataset:
+  """Drop fields that are static and do not vary with time."""
+  vars_to_drop = [k for k, v in dataset.items() if 'time' not in v.dims[0]]  # pytype: disable=unsupported-operands
+  return dataset.drop_vars(vars_to_drop)
+
+
+NestedTensors = TypeVar('NestedTensors', tf.Tensor, dict[str, tf.Tensor])
+
+
+@tf.function(jit_compile=True, autograph=False)
+def rolling_window_tensors(
+    inputs: NestedTensors, /, size: int, shift: int = 1, stride: int = 1
+) -> NestedTensors:
+  """Calculate a tensor of rolling windows.
+
+  Example usage:
+
+    >>> rolling_window_tensors(tf.range(10), size=6, shift=2)
+    <tf.Tensor: shape=(3, 6), dtype=int32, numpy=
+    array([[0, 1, 2, 3, 4, 5],
+          [2, 3, 4, 5, 6, 7],
+          [4, 5, 6, 7, 8, 9]], dtype=int32)>
+
+    >>> rolling_window_tensors(tf.range(10), size=4, stride=2)
+    <tf.Tensor: shape=(4, 4), dtype=int32, numpy=
+    array([[0, 2, 4, 6],
+          [1, 3, 5, 7],
+          [2, 4, 6, 8],
+          [3, 5, 7, 9]], dtype=int32)>
+
+  Args:
+    inputs: nested data structure with tf.Tensor values of shape [T, ...].
+    size: size of the time dimension in rolling window samples.
+    shift: shift between subsequent window samples along time.
+    stride: shift within a window along time.
+
+  Returns:
+    Nested tensors of shape [S, W, ...] sampled from inputs, where S is the
+    number of samples and W is the window size.
+  """
+
+  def calculate_windows(tensor):
+    shifts = tf.range(0, tf.shape(tensor)[0] - stride * (size - 1), shift)
+    indices = tf.range(0, size * stride, stride)
+    samples = tf.vectorized_map(
+        lambda shift: tf.gather(tensor, shift + indices), shifts
+    )
+    samples = tf.ensure_shape(samples, [None, size] + tensor.shape[1:])
+    return samples
+
+  return tf.nest.map_structure(calculate_windows, inputs)
+
+
+class Sampler:
+  """Base class for sampling from blocks."""
+
+  def list_block_slices(self, block_size: int, total_size: int) -> list[slice]:
+    """Returns a list of slices bounding blocks to sample from."""
+    raise NotImplementedError
+
+  def sample_block(self, data: NestedTensors) -> NestedTensors:
+    """Returns sample tensors from block tensors."""
+    raise NotImplementedError
+
+  @property
+  def example_size(self) -> int:
+    """Size of each example."""
+    raise NotImplementedError
+
+  def examples_per_block(self, block_size: int) -> int:
+    """Number of examples per block."""
+    raise NotImplementedError
+
+
+@dataclasses.dataclass
+class Splitter(Sampler):
+  """Split samples along the first axis."""
+
+
+  def list_block_slices(self, block_size: int, total_size: int) -> list[slice]:
+    return [
+        slice(start, min(start + block_size, total_size))
+        for start in range(0, total_size, block_size)
+    ]
+
+  def sample_block(self, data: NestedTensors) -> NestedTensors:
+    # Insert a dummy dimension for time-series length, which is always one.
+    return tf.nest.map_structure(lambda x: x[:, tf.newaxis, ...], data)
+
+  @property
+  def example_size(self) -> int:
+    return 1
+
+  def examples_per_block(self, block_size: int) -> int:
+    return block_size
+
+
+@dataclasses.dataclass
+class Windower(Sampler):
+  """Sample rolling windows along the first axis.
+
+  Attributes:
+    window_size: size of output windows.
+    stride_between_windows: offset between starting sequential windows.
+    output_window_stride: separation between between observations within a
+      window.
+    first_window_offset: offset of starting the first window.
+  """
+
+  window_size: int
+  stride_between_windows: int
+  output_window_stride: int = 1
+  first_window_offset: int = 0
+
+  def list_block_slices(self, block_size: int, total_size: int) -> list[slice]:
+    return list(
+        _iterate_windowed_block_slices(
+            sample_size=self.window_size,
+            block_size=block_size,
+            total_size=total_size,
+            stride_between_samples=self.stride_between_windows,
+            output_window_stride=self.output_window_stride,
+            first_sample_offset=self.first_window_offset,
+        )
+    )
+
+  def sample_block(self, data: NestedTensors) -> NestedTensors:
+    # NOTE(shoyer): It is tempting to try to use tf.data.Dataset.window instead
+    # for sampling windows, but that method does something different: it
+    # calculates windows over Dataset elements, rather than calculating windows
+    # within each Dataset element.
+    return rolling_window_tensors(
+        data,
+        size=self.window_size,
+        shift=self.stride_between_windows,
+        stride=self.output_window_stride,
+    )
+
+  @property
+  def example_size(self) -> int:
+    return self.window_size
+
+  def examples_per_block(self, block_size: int) -> int:
+    stop = block_size - self.output_window_stride * (self.window_size - 1)
+    return len(range(0, stop, self.stride_between_windows))
+
+
+@dataclasses.dataclass
+class WindowerAtOffsets(Sampler):
+  """Sample rolling windows along the first axis at specified offsets."""
+
+  window_size: int
+  window_offsets: list[int]
+  output_window_stride: int = 1
+
+  def list_block_slices(self, block_size: int, total_size: int) -> list[slice]:
+    stride = self.output_window_stride
+    # This suffices for now because generally we cache evaluation data.
+    sample_input_size = range(0, self.window_size * stride, stride)[-1] + 1
+    assert 0 < sample_input_size <= block_size <= total_size
+    slices = []
+    for start in self.window_offsets:
+      stop = start + sample_input_size
+      if stop > total_size:
+        raise ValueError(
+            f'offset at {start} needs data through {stop=}, which is beyond'
+            f' {total_size=}'
+        )
+      slices.append(slice(start, stop))
+    return slices
+
+  def sample_block(self, data: NestedTensors) -> NestedTensors:
+    def strided_sample(tensor):
+      # Insert a dummy batch/sample dimension.
+      return tf.ensure_shape(
+          tensor[tf.newaxis, :: self.output_window_stride],
+          [1, self.window_size] + tensor.shape[1:],
+      )
+
+    return tf.nest.map_structure(strided_sample, data)
+
+  @property
+  def example_size(self) -> int:
+    return self.window_size
+
+  def examples_per_block(self, block_size: int) -> int:
+    del block_size  # unused
+    return 1
+
+
+class Selector:
+  """Base class for block selection."""
+
+  def select(self, blocks: list[slice]) -> list[slice]:
+    """Select a subset of blocks for sampling."""
+    raise NotImplementedError
+
+
+class CompleteSelector(Selector):
+
+  def select(self, blocks: list[slice]) -> list[slice]:
+    return blocks
+
+
+@dataclasses.dataclass
+class ShardSelector(Selector):
+  shard_index: int
+  shard_count: int
+
+  def select(self, blocks: list[slice]) -> list[slice]:
+    return [
+        block
+        for i, block in enumerate(blocks)
+        if i % self.shard_count == self.shard_index
+    ]
+
+
+@dataclasses.dataclass
+class ShuffleSelector(Selector):
+  seed: int = 0
+  reshuffle_each_iteration: bool = True
+
+  def select(self, blocks: list[slice]) -> list[slice]:
+    rng = random.Random(self.seed)
+    if self.reshuffle_each_iteration:
+      self.seed = rng.randrange(2**63)
+    return rng.sample(blocks, k=len(blocks))
+
+
+@dataclasses.dataclass
+class ComposedSelector(Selector):
+  components: list[Selector]
+
+  def select(self, blocks: list[slice]) -> list[slice]:
+    for component in self.components:
+      blocks = component.select(blocks)
+    return blocks
+
+
+@dataclasses.dataclass
+class CustomSelector(Selector):
+  select: Callable[[list[slice]], list[slice]]
+
+
+def _thread_pool_loader(max_workers: int = 100):
+  """Dataset loader using a large thread pool for concurrency."""
+  # We use a separate thread for reading each data variable in each block.
+  executor = concurrent.futures.ThreadPoolExecutor(max_workers)
+
+  def load(dataset: xarray.Dataset) -> xarray.Dataset:
+    arrays = executor.map(lambda var: var.values, dataset.values())
+    return dataset.copy(data={k: v for k, v in zip(dataset, arrays)})
+
+  return load
+
+
+class _Reader:
+  """Class for reading an xarray.Dataset."""
+
+  def __init__(
+      self,
+      source: xarray.Dataset,
+      sampler: Sampler,
+      block_selector: Selector = CompleteSelector(),
+      *,
+      sample_dim: str = 'time',
+      block_size_in_bytes: float = 1e8,
+      parallel_block_reads: int = tf.data.AUTOTUNE,
+      parallel_samples: int = tf.data.AUTOTUNE,
+      dataset_loader: Optional[
+          Callable[[xarray.Dataset], xarray.Dataset]
+      ] = None,
+  ):
+    if dataset_loader is None:
+      # In principle, it could make sense to support passing alternative
+      # loaders, such as xarray_tensorstore.read() or a dask loader that calls
+      # .compute(). We don't yet have any use cases where this seems to make a
+      # difference, though. (The thread pool loader works as well as
+      # xarray_tensorstore.read.)
+      dataset_loader = _thread_pool_loader()
+
+    if sample_dim not in source.dims:
+      raise ValueError(
+          'source does not include variables with a'
+          f' {sample_dim!r} dimension:\n{source}'
+      )
+    source = _drop_static_vars(source)
+    source = source.transpose(sample_dim, ...)
+
+    block_size = _calculate_block_size(
+        source,
+        block_dims=[sample_dim],
+        bytes_per_request=block_size_in_bytes,
+        min_elements_per_request=sampler.example_size,
+    )
+
+    block_slices = sampler.list_block_slices(
+        block_size, source.sizes[sample_dim]
+    )
+
+    bytes_per_element = _xarray_bytes_per_element(source, {sample_dim})
+    bytes_per_example = sampler.example_size * bytes_per_element
+    examples_per_block = sampler.examples_per_block(block_size)
+    sample_bytes_per_block = bytes_per_example * examples_per_block
+    expansion = sample_bytes_per_block / block_size_in_bytes
+    logging.info(
+        f'picked block_size={block_size}, corresponding to {len(block_slices)} '
+        f'blocks with examples_per_block={examples_per_block}, based on '
+        f'sampler={sampler} and {block_size_in_bytes=:g}. '
+        f'{sample_bytes_per_block=:g} is a {expansion:1.2f}x expansion.'
+    )
+
+    self.source = source
+    self.sampler = sampler
+    self.block_selector = block_selector
+    self.parallel_block_reads = parallel_block_reads
+    self.parallel_samples = parallel_samples
+    self.dataset_loader = dataset_loader
+
+    self.block_size = block_size
+    self.block_slices = block_slices
+    self.bytes_per_example = bytes_per_example
+    self.examples_per_block = examples_per_block
+
+  def read(self) -> tf.data.Dataset:
+    """Read this dataset into a tf.data.Dataset."""
+
+    def generate_blocks():
+      for block in self.block_selector.select(self.block_slices):
+        yield (block.start, block.stop)
+
+    def np_read_block(start: np.ndarray, stop: np.ndarray) -> list[np.ndarray]:
+      selection = self.source.isel(time=slice(start, stop))
+      loaded = self.dataset_loader(selection)
+      arrays = [x.values for x in loaded.values()]
+      return arrays
+
+    def tf_read_block(start: tf.Tensor, stop: tf.Tensor):
+      dtypes = [v.dtype for v in self.source.values()]
+      shapes = [(None,) + v.shape[1:] for v in self.source.values()]
+      tensors = tf.numpy_function(np_read_block, [start, stop], dtypes)
+      for tensor, shape in zip(tensors, shapes):
+        tensor.set_shape(shape)
+      return dict(zip(self.source.keys(), tensors))
+
+    data = tf.data.Dataset.from_generator(
+        generate_blocks, output_signature=2 * (tf.TensorSpec((), tf.int64),)
+    )
+    data = data.map(tf_read_block, num_parallel_calls=self.parallel_block_reads)
+    data = data.map(
+        self.sampler.sample_block, num_parallel_calls=self.parallel_samples
+    )
+
+    data = data.unbatch()
+
+    return data
+
+
+def read_timeseries(
+    source: xarray.Dataset,
+    sampler: Sampler,
+    block_selector: Selector = CompleteSelector(),
+    *,
+    sample_dim: str = 'time',
+    block_size_in_bytes: float = 1e8,
+    parallel_block_reads: int = tf.data.AUTOTUNE,
+    parallel_samples: int = tf.data.AUTOTUNE,
+) -> tf.data.Dataset:
+  """Read a time-series xarray.Dataset into a tf.data.Dataset of windows.
+
+  See go/whirl-zarr-reader for a detailed description of the design.
+
+  Args:
+    source: lazy xarray.Dataset, e.g., opened from a Zarr file with
+      `open_zarr(..., chunks=None)`. All data variables with a 'time' dimension
+      will be sampled. Note: setting `chunks=None` to avoid using Dask is
+      preferred for optimal performance.
+    sampler: specification of what time-series samples of this dataset should
+      look like. Currently the only supported sampler is Windower.
+    block_selector: selector called at each pass through the source dataset,
+      indicating the blocks to read in order. The returned blocks should be a
+      subset of passed in blocks.
+    sample_dim: name of the dimension to sample along.
+    block_size_in_bytes: number of bytes to use for each reading a "block" of
+      data from the source data. Larger block sizes are more efficient.
+    parallel_block_reads: number of blocks to read in parallel.
+    parallel_samples: number of threads to use for generating samples from
+      blocks.
+
+  Returns:
+    tf.data.Dataset where each element is a dict of arrays.
+  """
+  return _Reader(
+      source=source,
+      sampler=sampler,
+      block_selector=block_selector,
+      sample_dim=sample_dim,
+      block_size_in_bytes=block_size_in_bytes,
+      parallel_block_reads=parallel_block_reads,
+      parallel_samples=parallel_samples,
+  ).read()
+
+
+def read_shuffled_shard(
+    source: xarray.Dataset,
+    sampler: Sampler,
+    *,
+    sample_dim: str = 'time',
+    block_size_in_bytes: float = 1e8,
+    buffer_size_in_bytes: float = 1e10,
+    min_buffer_blocks: float = 10,
+    parallel_block_reads: int = tf.data.AUTOTUNE,
+    parallel_samples: int = tf.data.AUTOTUNE,
+    shard_index: Optional[int] = None,
+    shard_count: Optional[int] = None,
+    seed: int = 0,
+    reshuffle_each_iteration: bool = True,
+) -> tf.data.Dataset:
+  """Read a time-series with samples in randomly shuffled order.
+
+  Args:
+    source: lazy xarray.Dataset, e.g., opened from a Zarr file with
+      `open_zarr(..., chunks=None)`. All data variables with a 'time' dimension
+      will be sampled. Note: setting `chunks=None` to avoid using Dask is
+      preferred for optimal performance.
+    sampler: specification of what time-series samples of this dataset should
+      look like.
+    sample_dim: name of the dimension to sample along.
+    block_size_in_bytes: number of bytes to use for each reading a "block" of
+      data from the source data. Larger block sizes are more efficient.
+    buffer_size_in_bytes: number of bytes to use in the shuffle buffer.
+    min_buffer_blocks: minimum number of blocks that must be represented in the
+      shuffle buffer, if more than one sample is taken from each block.
+      Typically this should be at least as large as the batch size.
+    parallel_block_reads: number of blocks to read in parallel.
+    parallel_samples: number of threads to use for generating samples from
+      blocks.
+    shard_index: integer index for this shard of the data, in the range `[0,
+      shard_count)`. In a multi-host JAX training setup, this should equal
+      `jax.process_index()`.
+    shard_count: total number of data shards. In a multi-host JAX training
+      setup, this should equal `jax.process_count()`.
+    seed: seed to use for random number generation.
+    reshuffle_each_iteration: whether to use a new shuffle order for elements
+      after each iteration through `source` or not.
+
+  Returns:
+    tf.data.Dataset where each element is a dict of arrays.
+  """
+  if shard_index is None and shard_count is None:
+    shard_index = 0
+    shard_count = 1
+
+  if shard_index is None or shard_count is None:
+    raise ValueError('must set both or neither of shard_index and shard_count')
+
+  selector = ComposedSelector([
+      ShardSelector(shard_index, shard_count),
+      ShuffleSelector(seed, reshuffle_each_iteration),
+  ])
+
+  def _make_reader(block_size_in_bytes):
+    reader = _Reader(
+        source=source,
+        sampler=sampler,
+        sample_dim=sample_dim,
+        block_selector=selector,
+        block_size_in_bytes=block_size_in_bytes,
+        parallel_block_reads=parallel_block_reads,
+        parallel_samples=parallel_samples,
+    )
+    buffer_size = int(buffer_size_in_bytes / reader.bytes_per_example)
+    logging.info(
+        f'picked shuffle buffer size of {buffer_size} based on '
+        f'{buffer_size_in_bytes=:g}'
+    )
+    return reader, buffer_size
+
+  reader, buffer_size = _make_reader(block_size_in_bytes)
+
+  if buffer_size:
+    examples_per_block = reader.examples_per_block
+    buffer_blocks = buffer_size / examples_per_block
+    if examples_per_block > 1 and buffer_blocks < min_buffer_blocks:
+      block_size_in_bytes = reader.bytes_per_example
+      logging.warning(
+          'insufficient diversity in proposed shuffle buffer: '
+          f'{examples_per_block=} and {buffer_size=} means that on average '
+          f'only {buffer_blocks:g} blocks will be represented in the shuffle '
+          f'buffer, which is less than {min_buffer_blocks=}. Falling back to '
+          f'one example per block ({block_size_in_bytes=:g}).'
+      )
+      reader, buffer_size = _make_reader(block_size_in_bytes)
+      assert reader.examples_per_block == 1
+
+  data = reader.read()
+
+  if buffer_size:
+    # for testing, disable the shuffle buffer if it has size zero (the shuffle
+    # method does not support size zero buffers)
+    data = data.shuffle(buffer_size, seed, reshuffle_each_iteration)
+
+  return data

model/reference_code/stochastic_losses.py

@@ -0,0 +1,458 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Stochastic losses for NeuralGCM."""
+import abc
+from typing import Callable, Optional, Sequence
+from dinosaur import typing
+import gin
+import jax
+import jax.numpy as jnp
+import model.reference_code.linear_transforms as linear_transforms
+import model.reference_code.metrics_base as metrics_base
+import model.reference_code.metrics_util as metrics_util
+from model.legacy import model_utils
+import numpy as np
+
+
+Pytree = typing.Pytree
+TrajectoryRepresentations = typing.TrajectoryRepresentations
+
+AggregationTransformConstructor = metrics_util.AggregationTransformConstructor
+
+tree_leaves = jax.tree_util.tree_leaves
+tree_map = jax.tree_util.tree_map
+
+
+def replicate(
+    x: Pytree,
+    axis_name: str = 'batch',
+    times: Optional[int] = None,
+) -> Pytree:
+  """Replicated a pytree across devices."""
+  if times is None:
+    times = jax.local_device_count()
+
+  def _replicate(_):
+    return x
+
+  return jax.pmap(_replicate, axis_name)(np.ones(times))
+
+
+class EnergyLikeLoss(metrics_base.Loss, abc.ABC):
+  """Energy-score like loss function.
+
+  Both CRPS and EnergyScore take the form (with E expectation)
+    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
+  where for CRPS ‖⋅‖ is the L1 norm, and for EnergyScore it is the L2 norm.
+
+  To create a general implementation, we decompose the norm as
+    ‖Z‖ := _norm_reduction_fn(_norm_inner_fn(Z))
+
+  For more see (21) and (22) in [1]; http://shortn/_Lyu0etEy1F
+
+  References:
+    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
+         prediction, and estimation. Journal of the American statistical
+         Association, 102(477), 359-378.
+  """
+
+  def __init__(
+      self,
+      trajectory_spec: metrics_util.TrajectorySpec,
+      components: Sequence[linear_transforms.LinearTransformConstructor],
+      time_step: Optional[int | slice] = None,
+      level: Optional[int] = None,
+      getter: Callable[[Pytree], Pytree] = (
+          metrics_util.filter_sim_time_and_diagnostics
+      ),
+      beta: float = 1.0,
+      ensemble_term_weight: float = 0.5,
+      is_nodal: bool = True,
+      is_encoded: bool = False,
+      coarsen_aggregation: AggregationTransformConstructor = (
+          metrics_util.AggregateIdentity
+      ),
+      vector_norm_squared_aggregation: AggregationTransformConstructor = (
+          metrics_util.AggregateIdentity
+      ),
+  ):
+    """Constructs an instance of EnergyLikeLoss.
+
+    Args:
+      trajectory_spec: Specification of spatial and temporal trajectory sizes.
+      components: Sequence of linear transformations to be applied to errors.
+      time_step: Step or slice at which to compute loss, or None for all steps.
+      level: Level to compute loss at, or None to use mean over all levels.
+      getter: Function for extracting a sub-pytree on which errors are computed.
+      beta: Power parameter of the loss. For energy score to be strictly proper
+        beta must be belong to `(0, 2)`.
+      ensemble_term_weight: Coefficient that specifcies how much weight is put
+        on the terms that captures the spread of the 2-ensemble. For standard
+        energy score this value should be set to `0.5`. It can be used to
+        interpolate to other scoring rules that are not strictly proper. For
+        example setting this value to `0.0` and setting `beta = 2.0` will result
+        in a squared error loss.
+      is_nodal: Indicator whether loss is computed in nodal space.
+      is_encoded: Indicator whether loss is computed in encoded(model) space.
+      coarsen_aggregation: Transform class that is used to aggregate errors
+        before computing the loss elements. This enables defining losses on
+        coarser representations that accentuate larger scale structure.
+        Currently this argument should be used only by PatchEnergyLoss. Example
+        coarsening operators include `RegriddingAggregation`, `TimeWindowSum`.
+      vector_norm_squared_aggregation: Transform class that is used to aggregate
+        components of the squared errors to form the distance for computing the
+        energy score. Currently this argument should be used only by
+        PatchEnergyLoss. Suitable aggregation methods include
+        `RegriddingAggregation`, `TimeWindowSum`, `SumVariables`, which would
+        correspond to vectors of (1) single level, time, variable, horizontal
+        neighbors; (2) single level, variable, lon-lat, sequence of time values;
+        (3) all variables at a single level, time, lon-lat.
+    """
+    self.coarsen_fn = coarsen_aggregation(
+        trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded
+    )
+    self.vector_norm_squared_fn = vector_norm_squared_aggregation(
+        self.coarsen_fn.out_trajectory_spec,
+        is_nodal=is_nodal,
+        is_encoded=is_encoded,
+    )
+    # parent class reductions are done on the final out_trajectory_spec.
+    super().__init__(
+        self.vector_norm_squared_fn.out_trajectory_spec,
+        is_nodal=is_nodal,
+        is_encoded=is_encoded,
+    )
+    self.components = components
+    self.time_step = time_step
+    self.level = level
+    self.getter = getter
+    # transform is applied to raw inputs which are aligned with trajectory_spec.
+    self.transform = linear_transforms.ComposedTransformForLoss(
+        trajectory_spec, self.components
+    )
+    self._beta = beta
+    self._ensemble_term_weight = ensemble_term_weight
+
+  def a_minus_cb(self, a: Pytree, c: float, b: Pytree) -> Pytree:
+    """A - c * B."""
+    return tree_map(lambda a_i, b_i: a_i - c * b_i, a, b)
+
+  def ca_minus_b(self, c: float, a: Pytree, b: Pytree) -> Pytree:
+    """c * A - B."""
+    return tree_map(lambda a_i, b_i: c * a_i - b_i, a, b)
+
+  def component_mean(self, tree: Pytree) -> jax.Array:
+    """Mean over variable, time, pressure, lat, lon."""
+    leaf_means = tree_leaves(self.mean_per_variable(tree))
+    return sum(leaf_means) / len(leaf_means)
+
+  def ensemble_mean(self, tree: Pytree) -> Pytree:
+    return jax.lax.pmean(tree, 'ensemble')
+
+  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
+    """Prepares target or predictions."""
+    # Cannot consolidate with RMSE.prepare since this one
+    # * does not take ensemble mean of trajectory.
+    trajectory = metrics_util.extract_variable(
+        trajectory,
+        self.trajectory_spec,
+        self.time_step,
+        self.level,
+        self.getter,
+        self.is_nodal,
+        self.is_encoded,
+    )
+    return trajectory
+
+  def evaluate(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    """Evaluates giving values of interest."""
+    pv2ss = self._per_variable_spread_skill_errors(prediction, target)
+    return self._spread_skill_and_loss(
+        x_minus_y=pv2ss['x_minus_y'],
+        x_minus_xprime=pv2ss['x_minus_xprime'],
+    )['loss']
+
+  def debug_loss_terms_instance(self) -> metrics_base.EvaluateFunctionWrapper:
+    """Returns class that evaluates rel loss per variable and spread/skill."""
+
+    def evaluate_fn(
+        prediction: TrajectoryRepresentations,
+        target: TrajectoryRepresentations,
+    ) -> Pytree:
+      # self.loss.evaluate takes ensemble mean (to evaluate on ensemble mean) if
+      # needed.
+      pv2ss = self._per_variable_spread_skill_errors(prediction, target)
+      overall_spread_skill_loss = self._spread_skill_and_loss(
+          x_minus_y=pv2ss['x_minus_y'],
+          x_minus_xprime=pv2ss['x_minus_xprime'],
+      )
+
+      all_vars = pv2ss['x_minus_y'].keys()
+
+      per_variable_terms = {
+          var: self._spread_skill_and_loss(
+              x_minus_y=pv2ss['x_minus_y'][var],
+              x_minus_xprime=pv2ss['x_minus_xprime'][var],
+          )
+          for var in all_vars
+      }
+      # here we reduce terms by summation to expose relative contributions,
+      # even though the actual total_loss might be different.
+      per_variable_losses = {
+          var: per_variable_terms[var]['loss'] for var in all_vars
+      }
+      sum_of_losses = sum(per_variable_losses.values())
+      per_variable_relative_losses = tree_map(
+          lambda x: x / sum_of_losses, per_variable_losses
+      )
+      return {
+          'relative_loss': per_variable_relative_losses,
+          'overall': overall_spread_skill_loss,
+          'per_variable_spread': {
+              var: per_variable_terms[var]['spread'] for var in all_vars
+          },
+          'per_variable_skill': {
+              var: per_variable_terms[var]['skill'] for var in all_vars
+          },
+      }
+
+    return metrics_base.EvaluateFunctionWrapper(evaluate_fn)
+
+  def _per_variable_spread_skill_errors(
+      self,
+      prediction: TrajectoryRepresentations,
+      target: TrajectoryRepresentations,
+  ) -> Pytree:
+    """Computes non-reduced loss terms (skill and spread) for each variable.
+
+    Args:
+      prediction: predicted 2-ensemble of trajectories with each component
+        having shape [2, time_steps, vertical, lat_axis, lon_axis], with leading
+        axis corresponding to different ensemble members and last two axes being
+        either spherical harmonics numbers or lat, lon values.
+      target: target trajectory replicated along the ensemble axis. The shape is
+        expected to be exactly the same as `trajectory`.
+
+    Returns:
+      A dictionary with keys containing transformed variables.
+        `x_minus_y` = prediction - target
+        `x_minus_xprime` = difference of ensemble predictions
+        `prediction` = prediction
+    """
+    ensemble_size = jax.lax.psum(1, 'ensemble')
+    if ensemble_size != 2:
+      raise ValueError(f'{ensemble_size=} is not 2')
+
+    prediction = self.transform(self._prepare(prediction), target)
+    target = self.transform(self._prepare(target), target)
+
+    x_minus_y = tree_map(jnp.subtract, prediction, target)  # X_i - Y
+
+    xprime = jax.lax.pshuffle(prediction, 'ensemble', (1, 0))
+    x_minus_xprime = tree_map(jnp.subtract, prediction, xprime)  # X_i - X_j≠i
+
+    return {
+        'x_minus_y': x_minus_y,
+        'x_minus_xprime': x_minus_xprime,
+        'prediction': prediction,
+    }
+
+  @abc.abstractmethod
+  def _spread_skill_and_loss(
+      self,
+      x_minus_y: Pytree,
+      x_minus_xprime: Pytree,
+  ) -> dict[str, jax.Array]:
+    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""
+
+
+@gin.register(
+    denylist=['coarsen_aggregation', 'vector_norm_squared_aggregation']
+)
+class CRPSLoss(EnergyLikeLoss):
+  """CRPS loss on linearly transformed errors.
+
+  CRPS takes the form (with E expectation)
+    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
+  where ‖⋅‖ is the L1 norm. It can be thought of as the sum of component-wise
+  energy score losses.
+
+  Based on formula 21 in [1]; http://shortn/_Lyu0etEy1F
+
+  References:
+    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
+         prediction, and estimation. Journal of the American statistical
+         Association, 102(477), 359-378.
+  """
+
+  def _spread_skill_and_loss(
+      self,
+      x_minus_y: Pytree,
+      x_minus_xprime: Pytree,
+  ) -> dict[str, jax.Array]:
+    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""
+    a_minus_cb = self.a_minus_cb
+    ensemble_mean = self.ensemble_mean
+    component_mean = self.component_mean
+
+    def abs_beta(tree: Pytree) -> Pytree:
+      return tree_map(lambda x: jnp.abs(x) ** self._beta, tree)
+
+    # With X, X' two i.i.d. predictions,
+    #   Skill  = (1/2)[ (1/N)Σₙ|Xₙ-Yₙ| + (1/N)Σₙ|Xₙ'-Yₙ| ]
+    #   Spread = (1/N) Σₙ|Xₙ-Xₙ'|
+
+    # Recall x_minus_y = X-Y on one device and X'-Y on another. So the ensemble
+    # mean of this (which is all-reduced) is exactly Skill above.
+    skill = component_mean(ensemble_mean(abs_beta(x_minus_y)))
+
+    # One device has X-X' and the other has X'-X, so the ensemble mean is the
+    # same on both devices.
+    spread = component_mean(ensemble_mean(abs_beta(x_minus_xprime)))
+
+    # Then CRPS = Skill - (1/2) Spread
+    # However, this is unstable if Spread = 2Skill + ε, where |ε| << |Spread|.
+    # In particular, up to numerical precision, CRPS will equal 0!
+    # This can happen if Prob[Xₙ = 1] = p << 1, and Prob[Xₙ = 0] = 1 - p.
+    # a stable estimate of CRPS is
+    #   CRPS = C + C'  (an ensemble mean)
+    # where
+    #      C = (1/N) Σₙ[ |Xₙ-Yₙ| - (1/2) |Xₙ-Xₙ'| ]
+    #      C'= (1/N) Σₙ[ |Xₙ'-Yₙ| - (1/2) |Xₙ'-Xₙ| ]
+    # which should be re-written as
+    #   CRPS = (1/(2N)) Σₙ[ |Xₙ-Yₙ| + |Xₙ'-Yₙ| - |Xₙ-Xₙ'| ]
+    # The triangle inequality ensures the summands are non-negative.
+    crps = component_mean(
+        ensemble_mean(
+            a_minus_cb(  # |Xₙ-Yₙ| - (1/2) |Xₙ-Xₙ'|
+                abs_beta(x_minus_y),
+                self._ensemble_term_weight,
+                abs_beta(x_minus_xprime),
+            )
+        )
+    )
+    return {'spread': spread, 'skill': skill, 'loss': crps}
+
+
+@gin.register(
+    denylist=['coarsen_aggregation', 'vector_norm_squared_aggregation']
+)
+class EnergyScoreLoss(EnergyLikeLoss):
+  """Energy score loss on linearly transformed errors.
+
+  EnergyScoreLoss takes the form (with E expectation)
+    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
+  where ‖⋅‖ is the L2 norm. It is a generalization of CRPS to
+  multiple-dimensions.
+
+  Based on formula 22 in [1]; http://shortn/_Lyu0etEy1F
+
+  References:
+    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
+         prediction, and estimation. Journal of the American Statistical
+         Association, 102(477), 359-378.
+  """
+
+  def _spread_skill_and_loss(
+      self,
+      x_minus_y: Pytree,
+      x_minus_xprime: Pytree,
+  ) -> dict[str, jax.Array]:
+    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""
+    a_minus_cb = self.a_minus_cb
+    ensemble_mean = self.ensemble_mean
+    component_mean = self.component_mean
+
+    def sqrt_beta(x: jax.Array) -> jax.Array:
+      return model_utils.safe_sqrt(x) ** self._beta
+
+    def square(tree: Pytree) -> Pytree:
+      return tree_map(jnp.square, tree)
+
+    # With X, X' two i.i.d. predictions,
+    #   Skill  = (1/2)[ ‖X-Y‖ + ‖X'-Y‖ ]
+    #   Spread = ‖Xₙ-Xₙ'‖
+
+    # Recall x_minus_y = X-Y on one device and X'-Y on another. So the ensemble
+    # mean of this (which is all-reduced) is exactly Skill above.
+    skill = ensemble_mean(sqrt_beta(component_mean(square(x_minus_y))))
+
+    # One device has X-X' and the other has X'-X, so the ensemble mean is the
+    # same on both devices. The call to ensemble_mean simply removes the
+    # ensemble dim.
+    spread = ensemble_mean(sqrt_beta(component_mean(square(x_minus_xprime))))
+
+    # The straightforward implementation will lose resolution when the relative
+    # difference between
+    #   ‖X - X'‖  AND ‖X - Y‖ + ‖X' - Y‖,
+    # is less than 1e-6. This is so unlikely that we do will not handle it.
+    es_straightforward = a_minus_cb(skill, self._ensemble_term_weight, spread)
+    es = es_straightforward
+
+    # Unused demonstration of how to handle this co-linear case with lots of
+    # extra complex operations.
+    # if float(self._beta) != 1:
+    #   es = es_straightforward
+    # else:
+    #   # If beta == 1, there is a high resolution fix.
+    #   # See http://screen/BkvX57d9B9eqMrB
+    #   #
+    #   # alpha = ‖X - Y‖²
+    #   alpha = component_mean(square(x_minus_y))
+    #   # And if ensemble_term_weight == 1/2,
+    #   # gamma_minus_alpha = ‖X - X'‖²/4 - ‖X - Y‖²
+    #   #                  = (1/N) Σₙ[ (Xₙ-Xₙ')²/4 - (Xₙ-Yₙ)² ]
+    #   gamma_minus_alpha = component_mean(
+    #       self.ca_minus_b(
+    #           self._ensemble_term_weight**2,
+    #           square(x_minus_xprime),
+    #           square(x_minus_y),
+    #       )
+    #   )
+    #   # If gamma = ‖X - X'‖²/4 = 0 (e.g. at step=0), then (γ-α)/α = -1,
+    #   # and then grad(sqrt1pm1) is NaN. However, in this case we can use the
+    #   # straightforward version with no issues.
+    #   gamma_minus_alpha_div_alpha = gamma_minus_alpha / alpha
+
+    #   # Construct a "safe" input to use in the go/tf-where-nan trick.
+    #   cutoff = -0.1
+    #   safe_gamma_minus_alpha_div_alpha = jnp.maximum(
+    #       gamma_minus_alpha / alpha, cutoff
+    #   )
+
+    #   # For γ ≈ α, safe_gamma_minus_alpha_div_alpha =
+    #   # gamma_minus_alpha_div_alpha, and this code block will be used.
+    #   es_for_small_diffs = ensemble_mean(
+    #       # sqrt1pm1(z) = sqrt(z + 1) - 1, so
+    #       # sqrt(α) * -1 * sqrt1pm1((γ-α)/α)
+    #       # = sqrt(α) * (1 - sqrt((γ-α)/α) + 1)
+    #       # = sqrt(α) * (1 - sqrt(γ/α))
+    #       # = sqrt(α) - sqrt(γ)
+    #       # = sqrt(‖X - Y‖²) - sqrt(‖X - X'‖²/4)
+    #       jnp.sqrt(alpha)
+    #       * -1
+    #       * tfp.math.sqrt1pm1(safe_gamma_minus_alpha_div_alpha)
+    #   )
+    #   es = jnp.where(
+    #       # Reminder that the triangle-inequality shows γ <= α always. So
+    #       # (γ - α) / α < 0.1 is a "large diff" (despite being negative).
+    #       gamma_minus_alpha_div_alpha < cutoff,
+    #       es_straightforward,
+    #       es_for_small_diffs,
+    #   )
+    return {'spread': spread, 'skill': skill, 'loss': es}

model/reference_code/train_utils.py

@@ -0,0 +1,662 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Training utility functions for NeuralGCM."""
+
+import collections
+from collections import abc
+import functools
+import logging
+import math
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+from dinosaur import pytree_utils
+from dinosaur import typing
+import einops
+import gin
+import haiku as hk
+import jax
+from jax.experimental import mesh_utils
+import jax.numpy as jnp
+from model.legacy import optimization
+import numpy as np
+import optax
+
+
+# pylint: disable=logging-fstring-interpolation
+
+
+PRNGKeyArray = typing.PRNGKeyArray
+Array = Union[np.ndarray, jnp.ndarray]
+PyTree = Any
+Forcing = typing.Forcing
+
+IntOrArray = Union[int, Array]
+OptState = optimization.OptState
+ModelParams = Any
+ModelGradients = ModelParams
+EMAParams = ModelParams
+StepAndOptState = Tuple[IntOrArray, OptState]
+StepOptAndEMAState = Tuple[IntOrArray, OptState, ModelParams]
+LossValue = Array
+LossFunction = Callable[[PyTree, PyTree], LossValue]
+LossAndGradFunction = Callable[
+    [ModelParams, PRNGKeyArray, PyTree, Forcing],
+    Tuple[LossValue, ModelGradients],
+]
+MetricFunction = Callable[[PyTree, PyTree], Union[Array, Mapping[str, Array]]]
+TrainStepFunction = Callable[
+    [PRNGKeyArray, StepAndOptState, PyTree, Forcing],
+    Tuple[StepAndOptState, LossValue],
+]
+EvalStepFunction = Callable[
+    [ModelParams, PRNGKeyArray, PyTree, Forcing], Mapping[str, Array]
+]
+TrajectoryFunction = Callable[
+    [ModelParams, PRNGKeyArray, PyTree, Forcing], Tuple[PyTree, PyTree]
+]
+
+
+def flatten_dict(
+    inputs: Mapping[str, Any],
+    parent_key: str = '',
+    sep: str = ' ',
+) -> Mapping[str, Array]:
+  """Returns a flattened version of `inputs` dictionary."""
+  items = []
+  for k, v in inputs.items():
+    new_key = parent_key + sep + k if parent_key else k
+    if isinstance(v, Mapping):
+      items.extend(flatten_dict(v, new_key, sep=sep).items())
+    else:
+      items.append((new_key, v))
+  keys, counts = np.unique(np.array([x[0] for x in items]), return_counts=True)
+  if (counts > 1).any():
+    raise ValueError(f'got duplicate keys {keys[counts > 1]}')
+  return dict(items)
+
+
+#
+#  Note that all functions below deal with *batched* inputs.
+#
+
+
+def loss_and_gradient(
+    trajectory_fn: TrajectoryFunction,
+    loss_fn: LossFunction,
+) -> LossAndGradFunction:
+  """Returns a function that computes loss and the gradient of the loss.
+
+  Args:
+    trajectory_fn: a function that accepts `params` and `initial_velocity` and
+      returns a trajectory of velocities.
+    loss_fn: a function that accepts a predicted trajectory and a ground truth
+      trajectory, returning a scalar loss value.
+
+  Returns:
+    A function that accepts `params, initial_velocity, target_trajectory` and
+    returns the loss and the gradient of the loss.
+  """
+
+  def _loss(
+      params: ModelParams,
+      rng: PRNGKeyArray,
+      target_trajectory: PyTree,
+      forcing_data: typing.ForcingData,
+  ) -> LossValue:
+    """Returns loss value and gradient with respect to model parameters."""
+    _, predicted_trajectory = trajectory_fn(
+        params, rng, target_trajectory, forcing_data
+    )
+    loss = loss_fn(predicted_trajectory, target_trajectory)  # type: ignore
+    return loss
+
+  return jax.value_and_grad(_loss)
+
+
+def train_step(
+    loss_and_grad_fn: LossAndGradFunction,
+    optimizer: optax.GradientTransformation,
+) -> TrainStepFunction:
+  """Returns a function that performs a single training step.
+
+  Args:
+    loss_and_grad_fn: a function that accepts `params, initial_velocity,
+      target_trajectory` and returns the loss and the gradient of the loss.
+    optimizer: Optax optimizer to update params and internal state.
+
+  Returns:
+    A function that performs a single training step.
+  """
+
+  def _train_step(
+      rng: PRNGKeyArray,
+      step_and_state: StepAndOptState,
+      target_trajectory: PyTree,
+      forcing_data: typing.ForcingData,
+  ) -> Tuple[StepAndOptState, LossValue]:
+    """A function that performs a single training step."""
+    step, opt_state = step_and_state
+    loss, grad = loss_and_grad_fn(
+        opt_state.params, rng, target_trajectory, forcing_data
+    )
+
+    updates, new_state = optimizer.update(
+        grad, opt_state.state, opt_state.params
+    )
+    new_params = optax.apply_updates(opt_state.params, updates)
+    new_opt_state = OptState(state=new_state, params=new_params)
+
+    return (step + 1, new_opt_state), loss
+
+  return _train_step
+
+
+def eval_batch(
+    trajectory_fn: TrajectoryFunction,
+    metric_funcs: Mapping[str, MetricFunction],
+) -> EvalStepFunction:
+  """Returns a function that performs a single evaluation step.
+
+  Args:
+    trajectory_fn: a function that accepts `params` and `initial_velocity` and
+      returns a trajectory of velocities.
+    metric_funcs: a dictionary mapping strings to metric funcutils, each
+      returning either a metric scalar or a dictionary of such.
+
+  Returns:
+    A function that performs a single evaluation step.
+  """
+
+  def _eval_batch(
+      params: ModelParams,
+      rng: PRNGKeyArray,
+      target_trajectory: PyTree,
+      forcing_data: typing.ForcingData,
+  ) -> Mapping[str, Array]:
+    """A function that performs a single evaluation step."""
+    _, predicted_trajectory = trajectory_fn(
+        params, rng, target_trajectory, forcing_data
+    )
+    metric_values = {
+        k: metric(predicted_trajectory, target_trajectory)
+        for k, metric in metric_funcs.items()
+    }
+    results = flatten_dict(metric_values)
+    return results
+
+  return _eval_batch
+
+
+def streaming_mean(
+    rngs: Iterable[PRNGKeyArray],
+    batch_and_forcing: Iterable[Tuple[PyTree, Forcing]],
+    eval_fn: Callable[[PRNGKeyArray, PyTree, Forcing], Mapping[str, Array]],
+    data_preprocess_fn: Callable[..., PyTree] = lambda x: x,
+) -> Mapping[str, Array]:
+  """Runs evaluation on `eval_data`.
+
+  Args:
+    rngs: an iterable of random number keys to be used for evaluation.
+    batch_and_forcing: an iterable of batched velocity trajectories and forcing.
+    eval_fn: a function that performs a single evaluation step.
+    data_preprocess_fn: a preprocessing function be applied to each batch.
+
+  Returns:
+    A dict mapping strings to metric values.
+
+  Raises:
+    RuntimeError: if there are no batches to iterate over.
+  """
+  eval_metrics = collections.defaultdict(float)
+  count = 0
+  for rng, (batch, forcing) in zip(rngs, batch_and_forcing):
+    batch = data_preprocess_fn(batch)
+    batch_metrics = eval_fn(rng, batch, forcing)
+    for k, v in batch_metrics.items():
+      eval_metrics[k] += v
+    count += 1
+  if not count:
+    raise RuntimeError('no batches to iterate over')
+  return {k: v / count for k, v in eval_metrics.items()}
+
+
+@gin.register
+def identity(batch: Tuple[Array, ...], rng: Array = None) -> Tuple[Array, ...]:  # pytype: disable=annotation-type-mismatch  # jax-ndarray
+  """Identity preprocessing function that does not modify the `batch`."""
+  del rng  # unused.
+  return batch
+
+
+@gin.configurable
+def add_noise_to_input_frame(
+    batch: Tuple[Array, ...], rng: Array, scale: float = 1e-2, **kwargs
+) -> Tuple[Array, ...]:
+  """Adds noise to the 0th time frame in the `batch`.
+
+  Args:
+    batch: original batch to which the noise will be added.
+    rng: random number key to be used to generate noise.
+    scale: scale of the normal noise to be added.
+    **kwargs: other keyword arguments. Not used.
+
+  Returns:
+    batch with noise added along the 0th time slice.
+  """
+  del kwargs  # unused.
+  time_zero_slice = pytree_utils.slice_along_axis(batch, 1, 0)
+  shapes = jax.tree.map(np.shape, time_zero_slice)
+  rngs = jax.random.split(rng, len(jax.tree.leaves(time_zero_slice)))
+  rngs = jax.tree.unflatten(jax.tree.structure(time_zero_slice), rngs)
+
+  def noise_fn(key, s):
+    return scale * jax.random.truncated_normal(key, -2.0, 2.0, s)
+
+  noise = jax.tree.map(noise_fn, rngs, shapes)
+  add_noise_fn = lambda x, n: x.at[:, 0, ...].add(n)
+  return jax.tree.map(add_noise_fn, batch, noise)
+
+
+def preprocess(
+    data_iterator: Iterator[Tuple[Array, ...]],
+    rng_stream: Iterator[Array],
+    preprocess_fn: Callable[..., Tuple[Array, ...]],
+):
+  """Generator that applies `preprocess_fn` to entries of the `data_iterator`.
+
+  Args:
+    data_iterator: numpy iterator holding the data.
+    rng_stream: stream of random numbers to be used by `preprocess_fn`.
+    preprocess_fn: preprocessing function to be applied to each batch of data.
+
+  Yields:
+    Batch of data from `data_iterator` preprocessed with `preprocess_fn`.
+  """
+  preprocess_fn = jax.jit(preprocess_fn)
+  while True:
+    rng = next(rng_stream)
+    yield preprocess_fn(next(data_iterator), rng)
+
+
+def split_rngs(rngs: PRNGKeyArray, num: int) -> PRNGKeyArray:
+  """Splits `rngs` into `num` along the last batch axis."""
+  ndim = rngs.ndim
+  split_fn = jax.random.split
+  for _ in range(ndim - 1):
+    split_fn = jax.vmap(split_fn, (0, None), 1)
+  return split_fn(rngs, num)
+
+
+@functools.partial(jax.jit, static_argnames=['batch_shape'])
+def _split_rmgs_by_batch_shape(
+    rngs: PRNGKeyArray,
+    batch_shape: tuple[int, ...],
+) -> PRNGKeyArray:
+  for batch_size in batch_shape[::-1]:
+    rngs = split_rngs(rngs, batch_size)
+  return rngs
+
+
+class BatchedPRNGSequence(Iterator):
+  """Iterator of JAX different random keys split by `batch_shape`."""
+
+  def __init__(
+      self,
+      key_or_seed: Union[int, PRNGKeyArray],
+      batch_shape: Optional[Tuple[int, ...]] = None,
+  ):
+    """Creates an instance a class.
+
+    Args:
+      key_or_seed: Key or seed to initialize the random sequence.
+      batch_shape: Batch shape of the sequence.
+    """
+    self._key = hk.PRNGSequence(key_or_seed)
+    self.batch_shape = batch_shape
+
+  def reserve(self, num: int):
+    """Splits an additional ``num`` keys for later use."""
+    self._key = self._key.reserve(num)
+
+  def __next__(self):
+    rngs = next(self._key)
+    return _split_rmgs_by_batch_shape(rngs, self.batch_shape)
+
+
+@jax.jit
+def _combine_rng_seeds(seeds: jax.Array) -> jax.Array:
+  key = jax.random.PRNGKey(seeds[0])
+  for seed in seeds[1:]:
+    key = jax.random.fold_in(key, seed)
+  return jax.random.bits(key, shape=(), dtype=jnp.uint32)
+
+
+def combine_rng_seeds(*seeds: int) -> int:
+  """Combine uint32 seeds into a single Python integer RNG seed."""
+  # Put the seeds on the first CPU device so that JAX runs the entire
+  # computation on the CPU.
+  seeds = jax.device_put(
+      np.array(seeds), device=jax.local_devices(backend='cpu')[0]
+  )
+  return int(_combine_rng_seeds(seeds))
+
+
+def ensure_sharded_rng_key(
+    rng_key: jax.Array, *, mesh: jax.sharding.Mesh
+) -> jax.Array:
+  """Ensure that a batched PRNG key is sharded across all devices."""
+  spec = P('batch', 'ensemble', None)
+  sharding = jax.sharding.NamedSharding(mesh, spec)
+  return jax.lax.with_sharding_constraint(rng_key, sharding)
+
+
+def get_tpu_physical_mesh_shape() -> tuple[int, int, int] | None:
+  """Get the shape of the TPU connectivity torus for v4 or v5 chips."""
+  jax_devices = jax.devices()
+  try:
+    device_coords = [d.coords for d in jax_devices]
+  except AttributeError:
+    return None  # no "coords" attribute (e.g., using CPU devices)
+  dims = tuple(d + 1 for d in max(device_coords))
+  if len(dims) != 3 or math.prod(dims) != len(jax_devices):
+    return None
+  return dims
+
+
+# dict of dicts of indicating how to rearrange from physical TPU mesh layouts
+# (X, Y, Z) into logical mesh layouts (batch, ensemble, z, x, y) with
+# einops.rearrange for model training.
+# {tpu_topology: {(ensemble_shards, z_shard, x_shards, y_shards): ...}}
+_TPU_LAYOUT_REARRANGEMENTS = {
+    '2x2x2': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 2, 1, 1): 'z b0 b1 -> (b0 b1) () z () ()',
+        (2, 1, 1, 1): 'e b0 b1 -> (b0 b1) e () () ()',
+    },
+    '2x2x4': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 2, 1, 1): 'z b0 b1 -> (b0 b1) () z () ()',
+        (1, 4, 1, 1): 'b0 b1 z -> (b0 b1) () z () ()',
+        (2, 1, 1, 1): 'e b0 b1 -> (b0 b1) e () () ()',
+    },
+    '2x4x4': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 2, 1, 1): 'z b0 b1 -> (b0 b1) () z () ()',
+        (1, 4, 1, 1): 'b0 b1 z -> (b0 b1) () z () ()',
+        (2, 1, 1, 1): 'e b0 b1 -> (b0 b1) e () () ()',
+        (2, 2, 1, 1): 'z (b0 e) b1 -> (b0 b1) e z () ()',
+    },
+    '4x4x4': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 2, 1, 1): '(b0 z) b1 b2 -> (b0 b1 b2) () z () ()',
+        (1, 4, 1, 1): 'b0 b1 z -> (b0 b1) () z () ()',
+        (1, 2, 2, 1): '(b0 z) (b1 x) b2 -> (b0 b1 b2) () z x ()',
+        (2, 1, 1, 1): '(b0 e) b1 b2 -> (b0 b1 b2) e () () ()',
+        (2, 2, 1, 1): '(b0 e) (b1 z) b2 -> (b0 b1 b2) e z () ()',
+        (2, 2, 2, 1): '(b0 e) (b1 z) (b2 x) -> (b0 b1 b2) e z x ()',
+    },
+    '4x4x8': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 4, 2, 1): 'z (b0 x) b1 -> (b0 b1) () z x ()',
+        (1, 4, 2, 2): 'z (b0 x) (b1 y) -> (b0 b1) () z x y',
+        (2, 4, 2, 1): 'z (b0 x) (b1 e) -> (b0 b1) e z x ()',
+    },
+    '4x8x8': {
+        (1, 1, 1, 1): 'b0 b1 b2 -> (b0 b1 b2) () () () ()',
+        (1, 4, 2, 1): 'z (b0 x) b1 -> (b0 b1) () z x ()',
+        (1, 4, 2, 2): 'z (b0 x) (b1 y) -> (b0 b1) () z x y',
+        (2, 4, 2, 1): 'z (b0 e) (b1 x) -> (b0 b1) e z x ()',
+        (2, 4, 2, 2): 'z (b0 e x) (b1 y) -> (b0 b1) e z x y',
+    },
+    '2x2x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): 'z b0 () -> b0 () z () ()',
+        (2, 1, 1, 1): 'e b0 () -> b0 e () () ()',
+        (2, 2, 1, 1): 'e z () -> () e z () ()',
+    },
+    '2x4x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): 'z b0 () -> b0 () z () ()',
+        (2, 1, 1, 1): 'e b0 () -> b0 e () () ()',
+        (2, 2, 1, 1): 'z (b0 e) -> b0 e z () ()',
+    },
+    '4x4x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): '(b0 z) b1 () -> (b0 b1) () z () ()',
+        (2, 1, 1, 1): '(b0 e) b1 () -> (b0 b1) e () () ()',
+        (2, 2, 1, 1): '(b0 e) (b1 z) () -> (b0 b1) e z () ()',
+    },
+    '4x8x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): '(b0 z) b1 () -> (b0 b1) () z () ()',
+        (1, 2, 2, 1): '(b0 z) (b1 x) () -> (b0 b1) () z x ()',
+        (2, 1, 1, 1): '(b0 e) b1 () -> (b0 b1) e () () ()',
+        (2, 2, 1, 1): '(b0 e) (b1 z) () -> (b0 b1) e z () ()',
+    },
+    '8x8x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): '(b0 z) b1 () -> (b0 b1) () z () ()',
+        (1, 2, 2, 1): '(b0 z) (b1 x) () -> (b0 b1) () z x ()',
+        (2, 1, 1, 1): '(b0 e) b1 () -> (b0 b1) e () () ()',
+        (2, 2, 1, 1): '(b0 e) (b1 z) () -> (b0 b1) e z () ()',
+        (2, 2, 2, 1): '(b0 e z) (b1 x) () -> (b0 b1) e z x ()',
+    },
+    '8x16x1': {
+        (1, 1, 1, 1): 'b0 b1 () -> (b0 b1) () () () ()',
+        (1, 2, 1, 1): '(b0 z) b1 () -> (b0 b1) () z () ()',
+        (1, 2, 2, 1): '(b0 z) (b1 x) () -> (b0 b1) () z x ()',
+        (2, 1, 1, 1): '(b0 e) b1 () -> (b0 b1) e () () ()',
+        (2, 2, 1, 1): '(b0 e) (b1 z) () -> (b0 b1) e z () ()',
+        (2, 2, 2, 1): '(b0 e z) (b1 x) () -> (b0 b1) e z x ()',
+    },
+}
+
+
+def create_spmd_mesh(sizes: dict[str, int]) -> jax.sharding.Mesh:
+  """Create an SPMD mesh suitable for data & model parallelism.
+
+  Args:
+    sizes: dictionary mapping from dimension names (batch, z, x, and y) to the
+      number of devices desired along that axis in the parallel mesh.
+
+  Returns:
+    Mesh with axis names ['batch', 'ensemble', 'x', 'y', 'z'] and the desired
+    axis sizes.
+  """
+  axis_names = ['batch', 'ensemble', 'z', 'x', 'y']
+  for name in sizes:
+    if name not in axis_names:
+      raise ValueError(f'unrecognized {name!r} not in {axis_names}')
+
+  logical_mesh_shape = tuple(
+      sizes.get(axis_name, 1) for axis_name in axis_names
+  )
+  if math.prod(logical_mesh_shape) != jax.device_count():
+    raise ValueError(
+        f'{logical_mesh_shape=} is incompatible with {jax.device_count()=}'
+    )
+
+  physical_mesh_shape = get_tpu_physical_mesh_shape()
+  if physical_mesh_shape is None:
+    try:
+      # only succeeds if the logical mesh shape perfectly matches the physical
+      # mesh, e.g., in the case of pure data parallelism
+      mesh_devices = mesh_utils.create_device_mesh(logical_mesh_shape)
+    except (AssertionError, NotImplementedError):
+      mesh_devices = np.reshape(jax.devices(), logical_mesh_shape)
+  else:
+    devices = np.empty(physical_mesh_shape, dtype=object)
+    for device in jax.devices():
+      devices[tuple(device.coords)] = device
+
+    topology = 'x'.join(map(str, physical_mesh_shape))
+    logical_mesh_shape = tuple(
+        sizes[dim] for dim in ['ensemble', 'z', 'x', 'y']
+    )
+    rearrangement = _TPU_LAYOUT_REARRANGEMENTS[topology][logical_mesh_shape]
+
+    abbreviated_sizes = {
+        'e': sizes['ensemble'],
+        'z': sizes['z'],
+        'x': sizes['x'],
+        'y': sizes['y'],
+    }
+    abbreviated_sizes = {k: v for k, v in abbreviated_sizes.items() if v != 1}
+    mesh_devices = einops.rearrange(devices, rearrangement, **abbreviated_sizes)
+
+  return jax.sharding.Mesh(mesh_devices, axis_names)
+
+
+P = jax.sharding.PartitionSpec
+
+
+def make_distributed_array_from_local_arrays(
+    pytree: PyTree,
+    mesh: jax.sharding.Mesh,
+    spatial_partitions: jax.sharding.PartitionSpec,
+    global_batch_size: int,
+) -> PyTree:
+  """Creates a pytree of global jax arrays for data/model parallelsm.
+
+  This function exists for loading spatially partitioned data, which is assumed
+  to be replicated across the ensemble dimension.
+
+  Args:
+    pytree: PyTree of NumPy arrays to convert into distributed JAX arrays. The
+      leading "batch" dimension is divided between different local devices.
+    mesh: SPDM sharding mesh.
+    spatial_partitions: JAX partition spec (of length 3) to use for partitioning
+      spatial dimensions (z, x, y).
+    global_batch_size: number distinct examples in a single batch across all
+      devices. Does not include the ensemble.
+
+  Returns:
+    Pytree with the same structure as the inputs, but with arrays replaced by
+    distributed JAX arrays.
+  """
+  if len(spatial_partitions) != 3:
+    raise ValueError(f'invalid {spatial_partitions=}')
+
+  def get_shard_count(spec_part: None | str | tuple[str, ...]) -> int:
+    # calculate the number of shards corresponding to an element in a
+    # PartitionSpec
+    if spec_part is None:
+      return 1
+    elif isinstance(spec_part, str):
+      return mesh.shape[spec_part]
+    else:
+      return math.prod(mesh.shape[x] for x in spec_part)
+
+  def shard_array(x: np.ndarray) -> jax.Array:
+    if x.ndim <= 3:
+      # handle sim_time [batch]
+      global_shape = (global_batch_size,) + x.shape[1:]
+      partition_spec = P('batch', *([None] * (x.ndim - 1)))
+    elif x.ndim == 4:
+      # This is currently needed to handle surface data that has shape:
+      # [batch, time, x, y].
+      _, x_shards, y_shards = map(get_shard_count, spatial_partitions)
+      global_shape = (
+          global_batch_size,
+          x.shape[1],
+          x.shape[2] * x_shards,
+          x.shape[3] * y_shards,
+      )
+      partition_spec = P('batch', None, *spatial_partitions[1:])
+    else:
+      # everything else has dimensions [batch, time, z, x, y]
+      assert x.ndim == 5, x.shape
+      z_shards, x_shards, y_shards = map(get_shard_count, spatial_partitions)
+      if x.shape[2] == 1:
+        z_shards = 1
+      global_shape = (
+          global_batch_size,
+          x.shape[1],
+          x.shape[2] * z_shards,
+          x.shape[3] * x_shards,
+          x.shape[4] * y_shards,
+      )
+      partition_spec = P('batch', None, *spatial_partitions)
+
+    sharding = jax.sharding.NamedSharding(mesh, partition_spec)
+    single_device_arrays = put_to_devices(x, jax.local_devices(), axis=0)
+    return jax.make_array_from_single_device_arrays(
+        global_shape, sharding, single_device_arrays
+    )
+
+  try:
+    return jax.tree_util.tree_map(shard_array, pytree)
+  except Exception as e:
+    shape_tree = jax.tree_util.tree_map(jnp.shape, pytree)
+    raise RuntimeError(
+        f'failed to shard arrays with shapes {shape_tree!r}'
+    ) from e
+
+
+def put_to_devices(
+    host_array: np.ndarray, local_devices: abc.Sequence[Any], axis: int
+) -> list[Any]:
+  """Transfers a host array to local devices, split on the first dimension."""
+  local_device_count = len(local_devices)
+  try:
+    per_device_arrays = np.split(host_array, local_device_count, axis=axis)
+  except ValueError as array_split_error:
+    raise ValueError(
+        f'Unable to put to devices shape {host_array.shape} with '
+        f'local device count {local_device_count}'
+    ) from array_split_error
+  device_buffers = [
+      jax.device_put(arr, d) for arr, d in zip(per_device_arrays, local_devices)
+  ]
+  return device_buffers
+
+
+def ensure_replicated(pytree: PyTree, *, mesh: jax.sharding.Mesh) -> PyTree:
+  """Ensure that a pytree is replicated across all devices."""
+
+  def replicate(x):
+    x = jnp.asarray(x)
+    spec = jax.sharding.PartitionSpec(*([None] * x.ndim))
+    sharding = jax.sharding.NamedSharding(mesh, spec)
+    return jax.lax.with_sharding_constraint(x, sharding)
+
+  return jax.tree_util.tree_map(replicate, pytree)
+
+
+T = TypeVar('T')
+
+
+def jit_once(f: T, **jit_kwargs) -> T:
+  """Like jax.jit, but raises an error instead of compiling multiple times."""
+  compiled = None
+
+  def g(*args, **kwargs):
+    nonlocal compiled
+    if compiled is None:
+      logging.info(f'lowering {f}')
+      lowered = jax.jit(f, **jit_kwargs).lower(*args, **kwargs)
+      logging.info(f'compiling {f}')
+      compiled = lowered.compile()
+      logging.info(f'finishing compiling {f}')
+    return compiled(*args, **kwargs)
+
+  return g

scripts/checkpoint_info.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Print parameter and serialization sizes for NeuralGCM checkpoints."""
+from __future__ import annotations
+
+import argparse
+import pickle
+import sys
+
+try:
+    from common import PROJECT_ROOT, resolve_path
+except ModuleNotFoundError:  # supports ``python -m scripts.checkpoint_info``
+    from scripts.common import PROJECT_ROOT, resolve_path
+
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.NeuralGCM import checkpoint_mode, format_parameter_summary
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoints", nargs="+")
+    args = parser.parse_args()
+    for value in args.checkpoints:
+        path = resolve_path(value)
+        with path.open("rb") as handle:
+            payload = pickle.load(handle)
+        if not isinstance(payload, dict) or "params" not in payload:
+            raise ValueError(f"{path} does not contain an official params tree")
+        mode = payload.get("mode") or checkpoint_mode(payload) or "unknown"
+        training_state = payload.get("training_state")
+        resume_text = (
+            f"resumable=true step={training_state.get('step')}"
+            if isinstance(training_state, dict)
+            else "resumable=false"
+        )
+        print(
+            f"checkpoint={path.name} mode={mode} "
+            f"file.bytes={path.stat().st_size:,} "
+            f"{resume_text} {format_parameter_summary(payload['params'])}"
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/common.py

@@ -0,0 +1,405 @@
+"""Shared config, channel and OneScience ERA5Dataset helpers."""
+
+from __future__ import annotations
+
+import json
+import sys
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+SYNTHETIC_GENERATOR_VERSION = "neuralgcm-hydrostatic-v2"
+
+
+def load_config(path: str | Path | None = None) -> dict[str, Any]:
+    path = Path(path or PROJECT_ROOT / "conf/config.yaml")
+    with path.open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def resolve_path(value: str | Path, config_path: str | Path | None = None) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    base = Path(config_path or PROJECT_ROOT / "conf/config.yaml").resolve().parent.parent
+    return base / path
+
+
+def channel_order(config: dict[str, Any]) -> list[str]:
+    return list(config["data"]["channel_order"])
+
+
+def pressure_levels(config: dict[str, Any]) -> list[int]:
+    return list(config["model"]["pressure_levels_hpa"])
+
+
+def as_time_major_frames(value: Any, *, name: str = "frames"):
+    """Normalize OneScience ERA5Dataset output to ``(T, C, H, W)``.
+
+    ERA5Dataset squeezes the leading time dimension when ``output_steps=1``;
+    callers must restore it before indexing forecast frames. Input frames are
+    allowed to remain ``(C, H, W)`` and should not use this helper.
+    """
+    import numpy as np
+
+    if hasattr(value, "detach"):
+        value = value.detach().cpu().numpy()
+    value = np.asarray(value)
+    if value.ndim == 3:
+        value = value[None, ...]
+    if value.ndim != 4:
+        raise ValueError(
+            f"{name} must have shape (T,C,H,W) or (C,H,W), got {value.shape}"
+        )
+    return value
+
+
+def load_era5_dataset(config: dict[str, Any], years: list[int], *, input_steps: int | None = None, output_steps: int | None = None):
+    """Construct the required OneScience ERA5Dataset, without replacing it."""
+    try:
+        from onescience.datapipes.climate import ERA5Dataset
+    except Exception as exc:
+        # Source-tree fallback mirrors the earth examples and keeps this
+        # project usable before OneScience is installed as a wheel.
+        local_src = Path("/public/home/yangzt01/onescience/src")
+        if local_src.exists() and str(local_src) not in sys.path:
+            sys.path.insert(0, str(local_src))
+        try:
+            from onescience.datapipes.climate import ERA5Dataset
+        except Exception as fallback_exc:
+            raise RuntimeError(
+            "OneScience ERA5Dataset import failed; load OneScience and its "
+            f"runtime modules first: {type(fallback_exc).__name__}: {fallback_exc}"
+            ) from fallback_exc
+    data_dir = resolve_path(config["data"]["data_dir"])
+    return ERA5Dataset(
+        dataset_dir=str(data_dir),
+        used_years=years,
+        used_variables=channel_order(config),
+        input_steps=input_steps or int(config["data"]["input_steps"]),
+        output_steps=output_steps or int(config["data"]["output_steps"]),
+        normalize=bool(config["data"].get("normalize", False)),
+    )
+
+
+def era5_data_is_synthetic(config: dict[str, Any], years: list[int]) -> bool:
+    """Return true only when every requested HDF5 file declares synthetic data."""
+    import h5py
+
+    data_dir = resolve_path(config["data"]["data_dir"]) / "data"
+    paths = [data_dir / f"{year}.h5" for year in years]
+    if not paths or any(not path.exists() for path in paths):
+        return False
+    try:
+        for path in paths:
+            with h5py.File(path, "r") as handle:
+                fields = handle[config["data"].get("field_key", "fields")]
+                if not bool(fields.attrs.get("synthetic", False)):
+                    return False
+    except (KeyError, OSError):
+        return False
+    return True
+
+
+def validate_synthetic_era5_version(
+    config: dict[str, Any], years: list[int]
+) -> None:
+    """Reject obsolete virtual fields that are known to destabilize the model."""
+    import h5py
+
+    data_dir = resolve_path(config["data"]["data_dir"]) / "data"
+    for year in years:
+        path = data_dir / f"{year}.h5"
+        with h5py.File(path, "r") as handle:
+            fields = handle[config["data"].get("field_key", "fields")]
+            if not bool(fields.attrs.get("synthetic", False)):
+                continue
+            version = fields.attrs.get("generator_version")
+            if isinstance(version, bytes):
+                version = version.decode()
+            if version != SYNTHETIC_GENERATOR_VERSION:
+                raise RuntimeError(
+                    f"Synthetic ERA5 file {path} uses obsolete generator_version="
+                    f"{version!r}; expected {SYNTHETIC_GENERATOR_VERSION!r}. "
+                    "Regenerate it with scripts/fake_data.py before running a "
+                    "NeuralGCM rollout."
+                )
+
+
+def write_json(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8")
+
+
+def era5_sample_to_xarray(sample: Any, config: dict[str, Any], *, timestamp: Any):
+    """Convert one ERA5Dataset frame to the official NeuralGCM xarray contract.
+
+    The HDF5 loader returns flattened channels in ``[C, latitude, longitude]``;
+    official NeuralGCM expects named variables with pressure ``level`` and
+    explicit latitude/longitude coordinates. Spatial interpolation to the
+    configured native grid is performed before the model API sees the data.
+    """
+    import numpy as np
+    import xarray as xr
+
+    invar = sample[0]
+    if hasattr(invar, "detach"):
+        invar = invar.detach().cpu().numpy()
+    channels = channel_order(config)
+    levels = pressure_levels(config)
+    height, width = invar.shape[-2:]
+    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
+    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
+    # ERA5Dataset stores (latitude, longitude), while NeuralGCM's xarray API
+    # expects (longitude, latitude) for horizontal fields.
+    dataset = xr.Dataset(coords={"latitude": lat, "longitude": lon, "time": [np.datetime64(timestamp)]})
+    grouped: dict[str, list[tuple[int, Any]]] = {}
+    for index, name in enumerate(channels):
+        if name in {"sea_ice_cover", "sea_surface_temperature"}:
+            values = xr.DataArray(invar[index].T, dims=("longitude", "latitude"), coords={"latitude": lat, "longitude": lon})
+        else:
+            base, _, suffix = name.rpartition("_")
+            if not suffix.isdigit() or base not in config["model"]["input_variables"] + config["model"].get("optional_input_variables", []):
+                continue
+            values = xr.DataArray(invar[index].T, dims=("longitude", "latitude"), coords={"latitude": lat, "longitude": lon}).expand_dims(level=[int(suffix)])
+            values = values.expand_dims(time=[np.datetime64(timestamp)])
+            grouped.setdefault(base, []).append((int(suffix), values))
+            continue
+        values = values.expand_dims(time=[np.datetime64(timestamp)])
+        dataset[name] = values
+    for base, entries in grouped.items():
+        entries.sort(key=lambda item: levels.index(item[0]) if item[0] in levels else item[0])
+        merged = xr.concat([value for _, value in entries], dim="level")
+        dataset[base] = merged.transpose("time", "level", "longitude", "latitude") if "time" in merged.dims else merged.transpose("level", "longitude", "latitude")
+    return dataset
+
+
+def era5_frames_to_xarray(
+    frames: Any,
+    config: dict[str, Any],
+    *,
+    start_time: Any,
+):
+    """Vectorized ERA5 ``(T,C,H,W)`` to NeuralGCM xarray conversion.
+
+    This is equivalent to concatenating ``era5_sample_to_xarray`` outputs, but
+    constructs every multi-level variable in one operation. It avoids hundreds
+    of small DataArray allocations per training window.
+    """
+    import numpy as np
+    import xarray as xr
+
+    frames = as_time_major_frames(frames, name="ERA5 trajectory")
+    channels = channel_order(config)
+    if frames.shape[1] != len(channels):
+        raise ValueError(
+            f"ERA5 trajectory has {frames.shape[1]} channels, expected "
+            f"{len(channels)}"
+        )
+    n_time, _, height, width = frames.shape
+    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
+    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
+    step_hours = int(config["data"].get("time_step_hours", 6))
+    times = np.datetime64(start_time) + np.arange(n_time) * np.timedelta64(step_hours, "h")
+    coords = {"time": times, "latitude": lat, "longitude": lon}
+    dataset = xr.Dataset(coords=coords)
+
+    level_indices: dict[str, list[tuple[int, int]]] = {}
+    allowed = set(config["model"]["input_variables"])
+    allowed.update(config["model"].get("optional_input_variables", []))
+    for channel_index, name in enumerate(channels):
+        if name in {"sea_ice_cover", "sea_surface_temperature"}:
+            dataset[name] = (
+                ("time", "longitude", "latitude"),
+                np.asarray(frames[:, channel_index]).transpose(0, 2, 1),
+            )
+            continue
+        base, _, suffix = name.rpartition("_")
+        if suffix.isdigit() and base in allowed:
+            level_indices.setdefault(base, []).append((int(suffix), channel_index))
+
+    configured_levels = pressure_levels(config)
+    for base, entries in level_indices.items():
+        entries.sort(
+            key=lambda item: configured_levels.index(item[0])
+            if item[0] in configured_levels
+            else item[0]
+        )
+        indices = [index for _, index in entries]
+        levels = [level for level, _ in entries]
+        values = np.asarray(frames[:, indices]).transpose(0, 1, 3, 2)
+        dataset[base] = (
+            ("time", "level", "longitude", "latitude"),
+            values,
+        )
+        dataset = dataset.assign_coords(level=np.asarray(levels))
+    return dataset
+
+
+def _target_grid(mode: str):
+    from dinosaur import spherical_harmonic
+
+    targets = {
+        "weather_forecast": spherical_harmonic.Grid.TL255,
+        "climate_scale": spherical_harmonic.Grid.TL127,
+        "forecast_2_8_deg": spherical_harmonic.Grid.TL63,
+        "stochastic_1_4_deg": spherical_harmonic.Grid.TL127,
+    }
+    try:
+        return targets[mode]()
+    except KeyError as exc:
+        raise ValueError(f"Unknown model mode {mode!r}") from exc
+
+
+@lru_cache(maxsize=16)
+def _profile_regridder(
+    height: int,
+    width: int,
+    mode: str,
+    latitude_spacing: str,
+    longitude_offset: float,
+):
+    """Construct and cache the profile's conservative regridder."""
+    from dinosaur import horizontal_interpolation, spherical_harmonic
+
+    source_grid = spherical_harmonic.Grid(
+        latitude_nodes=height,
+        longitude_nodes=width,
+        latitude_spacing=latitude_spacing,
+        longitude_offset=longitude_offset,
+    )
+    return horizontal_interpolation.ConservativeRegridder(
+        source_grid, _target_grid(mode), skipna=True
+    )
+
+
+def regrid_for_neuralgcm(dataset: Any, official_model: Any):
+    """Conservatively regrid ERA5 fields to the checkpoint's Gaussian grid."""
+    from dinosaur import horizontal_interpolation
+    from dinosaur import spherical_harmonic
+    from dinosaur import xarray_utils
+
+    source_grid = spherical_harmonic.Grid(
+        latitude_nodes=dataset.sizes["latitude"],
+        longitude_nodes=dataset.sizes["longitude"],
+        latitude_spacing=xarray_utils.infer_latitude_spacing(dataset.latitude),
+        longitude_offset=xarray_utils.infer_longitude_offset(dataset.longitude),
+    )
+    regridder = horizontal_interpolation.ConservativeRegridder(
+        source_grid, official_model.data_coords.horizontal, skipna=True
+    )
+    regridded = xarray_utils.regrid(dataset, regridder)
+    return xarray_utils.fill_nan_with_nearest(regridded)
+
+
+def regrid_for_profile(dataset: Any, mode: str):
+    """Regrid to the Gaussian data grid selected by an official Gin profile."""
+    from dinosaur import xarray_utils
+
+    regridder = _profile_regridder(
+        dataset.sizes["latitude"],
+        dataset.sizes["longitude"],
+        mode,
+        xarray_utils.infer_latitude_spacing(dataset.latitude),
+        float(xarray_utils.infer_longitude_offset(dataset.longitude)),
+    )
+    return xarray_utils.fill_nan_with_nearest(xarray_utils.regrid(dataset, regridder))
+
+
+@lru_cache(maxsize=16)
+def _load_static_features(
+    path_text: str,
+    mode: str | None,
+    target_height: int,
+    target_width: int,
+):
+    """Load and, only when necessary, regrid a reusable static dataset."""
+    import xarray as xr
+
+    with xr.open_dataset(path_text) as source:
+        static = source[["geopotential_at_surface", "land_sea_mask"]].load()
+    source_shape = (
+        static.sizes.get("latitude"),
+        static.sizes.get("longitude"),
+    )
+    if source_shape != (target_height, target_width):
+        if mode is None:
+            return None
+        static = regrid_for_profile(static, mode)
+    if (
+        static.sizes.get("latitude"),
+        static.sizes.get("longitude"),
+    ) != (target_height, target_width):
+        return None
+    return static
+
+
+def add_static_features(
+    dataset: Any,
+    config: dict[str, Any] | None = None,
+    *,
+    mode: str | None = None,
+    prefer_profile: bool = True,
+):
+    """Attach official profile static fields, with a synthetic fallback.
+
+    Callers attach fields after regridding the dynamic ERA5 trajectory. This
+    preserves the exact Gaussian-grid topography and land/sea mask bundled in
+    the official checkpoints. ``data.static_file`` remains a source-grid
+    fallback for installations that do not carry the released checkpoints.
+    """
+    import numpy as np
+
+    required = ("geopotential_at_surface", "land_sea_mask")
+    if config is not None and not set(required).issubset(dataset):
+        data_cfg = config.get("data", {})
+        profile_path = (
+            data_cfg.get("static_files", {}).get(mode) if mode else None
+        )
+        fallback_path = data_cfg.get("static_file")
+        candidates = []
+        if not prefer_profile and fallback_path:
+            candidates.append(fallback_path)
+        if mode:
+            if profile_path:
+                candidates.append(profile_path)
+        if prefer_profile and fallback_path:
+            candidates.append(fallback_path)
+        for value in candidates:
+            static_path = resolve_path(value)
+            if not static_path.exists():
+                continue
+            static = _load_static_features(
+                str(static_path.resolve()),
+                mode,
+                int(dataset.sizes["latitude"]),
+                int(dataset.sizes["longitude"]),
+            )
+            if static is None:
+                continue
+            for name in required:
+                if name not in dataset:
+                    # Both arrays are on the same profile Gaussian grid. Assign
+                    # by position rather than xarray label alignment: checkpoint
+                    # coordinates are float64 while regridded ERA5 coordinates
+                    # can be float32, and exact-label alignment would inject NaN.
+                    values = static[name].transpose("longitude", "latitude")
+                    dataset[name] = (
+                        ("longitude", "latitude"),
+                        np.asarray(values.values),
+                    )
+                    dataset[name].attrs.update(values.attrs)
+            dataset.attrs["static_features_source"] = str(static_path)
+            break
+    if "geopotential_at_surface" not in dataset:
+        dataset["geopotential_at_surface"] = (("longitude", "latitude"), np.zeros((dataset.sizes["longitude"], dataset.sizes["latitude"]), np.float32))
+    if "land_sea_mask" not in dataset:
+        dataset["land_sea_mask"] = (("longitude", "latitude"), np.zeros((dataset.sizes["longitude"], dataset.sizes["latitude"]), np.float32))
+    # Gin FloatDataFeatures parses units from these static fields exactly as in
+    # the official ERA5 pipeline.
+    dataset["geopotential_at_surface"].attrs.setdefault("units", "m**2 s**-2")
+    dataset["land_sea_mask"].attrs.setdefault("units", "dimensionless")
+    return dataset