+ + NeuralGCM + +
+ +# 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. diff --git a/conf/config.yaml b/conf/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bef69d25aee8d3f35f2ea62eaafbe250a6b1956d --- /dev/null +++ b/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 diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..357eaa5801a3516ff560c55b2b698f593c8e1e3e --- /dev/null +++ b/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" + ] +} diff --git a/configuration.json b/configuration.json new file mode 100644 index 0000000000000000000000000000000000000000..710668a83dc7fea1c14d593db1bdf54640cc1cd1 --- /dev/null +++ b/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" +} diff --git a/model/NeuralGCM.py b/model/NeuralGCM.py new file mode 100644 index 0000000000000000000000000000000000000000..8cc73ca8928058fec75a5c252ad9db2a01aea118 --- /dev/null +++ b/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 diff --git a/model/NeuralGCM_demo.py b/model/NeuralGCM_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..be5d8b92a62df6fc53d7c3978c791195f5352e65 --- /dev/null +++ b/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) diff --git a/model/data/era5_tl31_19590102T00.nc b/model/data/era5_tl31_19590102T00.nc new file mode 100644 index 0000000000000000000000000000000000000000..4c6376af0e4385de69711ec3caf329165dc09a52 --- /dev/null +++ b/model/data/era5_tl31_19590102T00.nc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:18f66e795af9f564a2b6e0d861b9a51e74ce831a82675b47ff957be709554a5e +size 2141788 diff --git a/model/data/tl63_stochastic_mini.pkl b/model/data/tl63_stochastic_mini.pkl new file mode 100644 index 0000000000000000000000000000000000000000..86595e4b996e5c22fcc60db3fc4ab679920c6043 --- /dev/null +++ b/model/data/tl63_stochastic_mini.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:de0f8e8cb31ecca59469104440808a3ea8e31418d5a8a48220e3c1fa64303baa +size 1030438 diff --git a/model/legacy/api.py b/model/legacy/api.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd01133c36c7c8253da762a2c3ee2a296911eea --- /dev/null +++ b/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) diff --git a/model/legacy/correctors.py b/model/legacy/correctors.py new file mode 100644 index 0000000000000000000000000000000000000000..333268bb342035388bef0b1f0668e4d596a39d6f --- /dev/null +++ b/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) diff --git a/model/legacy/decoders.py b/model/legacy/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..16983d53438fd9b99d1131e79c996dd3b772966b --- /dev/null +++ b/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)) diff --git a/model/legacy/diagnostics.py b/model/legacy/diagnostics.py new file mode 100644 index 0000000000000000000000000000000000000000..d818d5c2b497adacd3fc8ad547f4ab87655636a6 --- /dev/null +++ b/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} diff --git a/model/legacy/embeddings.py b/model/legacy/embeddings.py new file mode 100644 index 0000000000000000000000000000000000000000..a5ea8bcba1de92dcecef7a91cf3cf758567ef17e --- /dev/null +++ b/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) diff --git a/model/legacy/encoders.py b/model/legacy/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..9b0a499239ada7020082d5c456fd3279c2f34e1d --- /dev/null +++ b/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] + ) diff --git a/model/legacy/equations.py b/model/legacy/equations.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce904bea667f4ede5b9c751e84ce15200c354fb --- /dev/null +++ b/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 diff --git a/model/legacy/features.py b/model/legacy/features.py new file mode 100644 index 0000000000000000000000000000000000000000..f2e428b6731bcc6ba46ea3c6f7a22d02e5ddfaf5 --- /dev/null +++ b/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) diff --git a/model/legacy/filters.py b/model/legacy/filters.py new file mode 100644 index 0000000000000000000000000000000000000000..4e74f4afb9197f1376c4aa49c80380111a223e54 --- /dev/null +++ b/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)) diff --git a/model/legacy/forcings.py b/model/legacy/forcings.py new file mode 100644 index 0000000000000000000000000000000000000000..82b96828e6c2d0121b02da462324a0cce388fc3a --- /dev/null +++ b/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])}' + ) diff --git a/model/legacy/gin_utils.py b/model/legacy/gin_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fee0f18d6c6258ec94f85f0ddb994dccd7406676 --- /dev/null +++ b/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) diff --git a/model/legacy/initializers.py b/model/legacy/initializers.py new file mode 100644 index 0000000000000000000000000000000000000000..8a5ed26f727f52fbbb81efb937b4ccd60af3dcb4 --- /dev/null +++ b/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) diff --git a/model/legacy/integrators.py b/model/legacy/integrators.py new file mode 100644 index 0000000000000000000000000000000000000000..1ae94b2e1eb62626e533f8dab7720eda4d1d9ebe --- /dev/null +++ b/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) diff --git a/model/legacy/layers.py b/model/legacy/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..e32594028c365f82034eaedba44e2893677d3ce8 --- /dev/null +++ b/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 diff --git a/model/legacy/mappings.py b/model/legacy/mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..a3e4a5825fd661ad590f1482c1255312df512e7b --- /dev/null +++ b/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) diff --git a/model/legacy/model_builder.py b/model/legacy/model_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..1f27b38889bdd507d9219bf25f80aca93d5690e8 --- /dev/null +++ b/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] diff --git a/model/legacy/model_utils.py b/model/legacy/model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..45f04e6bc16f64186c2db80d239e12ef9e19da9a --- /dev/null +++ b/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 diff --git a/model/legacy/optimization.py b/model/legacy/optimization.py new file mode 100644 index 0000000000000000000000000000000000000000..bd1deb1121b9226297640f94d37ae72d1c67a490 --- /dev/null +++ b/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, + ) diff --git a/model/legacy/orographies.py b/model/legacy/orographies.py new file mode 100644 index 0000000000000000000000000000000000000000..9152a49126e5e28078e92a6df544f30f4fe5e3e6 --- /dev/null +++ b/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 diff --git a/model/legacy/parameterizations.py b/model/legacy/parameterizations.py new file mode 100644 index 0000000000000000000000000000000000000000..f7dfeb9041205e99f0f0ad859e0bcf2314273c23 --- /dev/null +++ b/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 diff --git a/model/legacy/perturbations.py b/model/legacy/perturbations.py new file mode 100644 index 0000000000000000000000000000000000000000..c0d1316577499dae3e5d0f18e77cc730d85ee5e8 --- /dev/null +++ b/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) diff --git a/model/legacy/physics_specifications.py b/model/legacy/physics_specifications.py new file mode 100644 index 0000000000000000000000000000000000000000..4e606e661d92132605b16a05b192307dbab9ae7e --- /dev/null +++ b/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) diff --git a/model/legacy/steps.py b/model/legacy/steps.py new file mode 100644 index 0000000000000000000000000000000000000000..84843b25d337e40d8b80732331dd744eefa9d266 --- /dev/null +++ b/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 + ) diff --git a/model/legacy/stochastic.py b/model/legacy/stochastic.py new file mode 100644 index 0000000000000000000000000000000000000000..b537b67b595f6e2cca6fb51aad37ca49013727bb --- /dev/null +++ b/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') diff --git a/model/legacy/towers.py b/model/legacy/towers.py new file mode 100644 index 0000000000000000000000000000000000000000..0a0001f1748d082af48c0797461c722bc0365954 --- /dev/null +++ b/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 diff --git a/model/legacy/transforms.py b/model/legacy/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..09bf161806e9692660274aa32cbdd9bd5f5026d0 --- /dev/null +++ b/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 diff --git a/model/reference_code/datasets.py b/model/reference_code/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..4cf2fc05e041e5557967407347b9b5ef1a1f5ba8 --- /dev/null +++ b/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 diff --git a/model/reference_code/experiment.py b/model/reference_code/experiment.py new file mode 100644 index 0000000000000000000000000000000000000000..215babbc0c9f466082cbc8b72f9c1c1044d0baa1 --- /dev/null +++ b/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)) diff --git a/model/reference_code/linear_transforms.py b/model/reference_code/linear_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..cff00a0f493d8b37e41fcc447221156c6973afb1 --- /dev/null +++ b/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) diff --git a/model/reference_code/metrics.py b/model/reference_code/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..605d9e5dcbb24f8244449b1b6da2fa3e5456d6be --- /dev/null +++ b/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 diff --git a/model/reference_code/metrics_base.py b/model/reference_code/metrics_base.py new file mode 100644 index 0000000000000000000000000000000000000000..16dbd43fa9ddc308cb6cbf920679e00af415e544 --- /dev/null +++ b/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) diff --git a/model/reference_code/metrics_util.py b/model/reference_code/metrics_util.py new file mode 100644 index 0000000000000000000000000000000000000000..613355e27f130050b72c6dca1970b2270180a9a9 --- /dev/null +++ b/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) diff --git a/model/reference_code/paper_configs/deterministic_0_7_deg.gin b/model/reference_code/paper_configs/deterministic_0_7_deg.gin new file mode 100644 index 0000000000000000000000000000000000000000..a4a39ebcb0cf7b0439abcb08bbfbde858615ad32 --- /dev/null +++ b/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 diff --git a/model/reference_code/paper_configs/deterministic_1_4_deg.gin b/model/reference_code/paper_configs/deterministic_1_4_deg.gin new file mode 100644 index 0000000000000000000000000000000000000000..18a9b75dbbd4db2d355ff4d153214984ecc5348c --- /dev/null +++ b/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 diff --git a/model/reference_code/paper_configs/deterministic_2_8_deg.gin b/model/reference_code/paper_configs/deterministic_2_8_deg.gin new file mode 100644 index 0000000000000000000000000000000000000000..36c297500836f1d8f5b962c62bbb0a0d1273cc9e --- /dev/null +++ b/model/reference_code/paper_configs/deterministic_2_8_deg.gin @@ -0,0 +1,2911 @@ +# Macros: +# ============================================================================== +ACTIVATION = @gelu +CORRECTOR_MODULE = @CustomCoordsCorrector +CORRECTOR_SCALE = 0.01 +DATA_FILTER_ATTENUATION = 0.0 +DYCORE_FILTER_ORDER = 3 +DYCORE_GRID = @smaller/GridWithWavenumbers() +DYCORE_INTEGRATOR = @imex_rk_sil3 +DYCORE_TAU = '120 minutes' +GLOBAL_OUT_SCALE = 0.02 +INPUT_FILTER_ATTEN = 5 +INPUT_FILTER_ORDER = 15 +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 = 1 +PARAMETERIZATION_FILTER = @ml/SequentialStepFilter +POSITIONAL_LATENT_SIZE = 8 +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, + @decoder_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 land_model/CombinedFeatures: +# ============================================================================== +land_model/CombinedFeatures.feature_module_names_to_exclude = () +land_model/CombinedFeatures.feature_modules = (@land_model/VelocityAndPrognostics,) +land_model/CombinedFeatures.features_to_exclude = () +land_model/CombinedFeatures.features_transform_module = \ + @land_model/ShiftAndNormalize +land_model/CombinedFeatures.name = None + +# Parameters for sea_ice_model/CombinedFeatures: +# ============================================================================== +sea_ice_model/CombinedFeatures.feature_module_names_to_exclude = () +sea_ice_model/CombinedFeatures.feature_modules = \ + (@sea_ice_model/VelocityAndPrognostics, @sea_ice_model/ForcingFeatures) +sea_ice_model/CombinedFeatures.features_to_exclude = () +sea_ice_model/CombinedFeatures.features_transform_module = \ + @sea_ice_model/ShiftAndNormalize +sea_ice_model/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 = None +coordinate_system_from_dataset.truncation = 'LINEAR' + +# Parameters for CoordinateSystem: +# ============================================================================== +CoordinateSystem.horizontal = @GridTL63() +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 encoder/DataExponentialFilter: +# ============================================================================== +encoder/DataExponentialFilter.attenuation = %INPUT_FILTER_ATTEN +encoder/DataExponentialFilter.cutoff = 0 +encoder/DataExponentialFilter.name = None +encoder/DataExponentialFilter.order = %INPUT_FILTER_ORDER + +# Parameters for orography/DataExponentialFilter: +# ============================================================================== +orography/DataExponentialFilter.attenuation = %DATA_FILTER_ATTENUATION +orography/DataExponentialFilter.cutoff = 0 +orography/DataExponentialFilter.name = None +orography/DataExponentialFilter.order = 1 + +# 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 = \ + @FilteredCustomOrography +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 = (@encoder/DataExponentialFilter,) +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 = 10 +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_ice_model/ForcingFeatures: +# ============================================================================== +sea_ice_model/ForcingFeatures.forcing_to_include = ('sea_ice_cover',) +sea_ice_model/ForcingFeatures.name = None + +# 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 = \ + [223.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 GridTL63: +# ============================================================================== +# None. + +# Parameters for smaller/GridWithWavenumbers: +# ============================================================================== +smaller/GridWithWavenumbers.dealiasing = 'quadratic' +smaller/GridWithWavenumbers.latitude_spacing = 'gauss' +smaller/GridWithWavenumbers.longitude_offset = 0.0 +smaller/GridWithWavenumbers.longitude_wavenumbers = 63 +smaller/GridWithWavenumbers.radius = None + +# Parameters for advance/IdentityTransform: +# ============================================================================== +advance/IdentityTransform.name = None + +# Parameters for land_model/IdentityTransform: +# ============================================================================== +land_model/IdentityTransform.name = None + +# Parameters for sea_ice_model/IdentityTransform: +# ============================================================================== +sea_ice_model/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.825e-05, + 7.041e-05, + 0.0001044, + 0.0001832, + 0.0007474, + 0.002588, + 0.007067, + 0.01524, + 0.02838, + 0.04518, + 0.06902, + 0.09718, + 0.1316, + 0.1716, + 0.2182, + 0.2748, + 0.3379, + 0.4073, + 0.4611, + 0.5177, + 0.5798, + 0.6547, + 0.7397, + 0.8362, + 0.9373, + 1.041, + 1.149, + 1.265, + 1.394, + 1.551, + 1.71, + 1.792] + +# 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.825e-05, + 7.041e-05, + 0.0001044, + 0.0001832, + 0.0007474, + 0.002588, + 0.007067, + 0.01524, + 0.02838, + 0.04518, + 0.06902, + 0.09718, + 0.1316, + 0.1716, + 0.2182, + 0.2748, + 0.3379, + 0.4073, + 0.4611, + 0.5177, + 0.5798, + 0.6547, + 0.7397, + 0.8362, + 0.9373, + 1.041, + 1.149, + 1.265, + 1.394, + 1.551, + 1.71, + 1.792] + +# 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.792e-05, + 5.989e-05, + 7.112e-05, + 8.21e-05, + 8.842e-05, + 9.109e-05, + 7.706e-05, + 7.156e-05, + 8.221e-05, + 9.83e-05, + 0.0001427, + 0.0002927, + 0.001, + 0.002714, + 0.006034, + 0.01158, + 0.01969, + 0.04488, + 0.08384, + 0.1362, + 0.2039, + 0.2933, + 0.3992, + 0.4827, + 0.5746, + 0.6994, + 0.8461, + 0.9233, + 1.001, + 1.08, + 1.162, + 1.248, + 1.339, + 1.448, + 1.579, + 1.675, + 1.712] + +# 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': 0.00012966, + 'specific_cloud_liquid_water_content': 0.0003412, + 'specific_humidity': 0.003474}, + 'u': 0.01924, + 'v': 0.0102, + '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.0497, + 'sim_time': 0.0, + 'temperature_variation': 33.62, + 'tracers': {'specific_cloud_ice_water_content': 0.0007559999999999999, + 'specific_cloud_liquid_water_content': 0.0014856000000000001, + 'specific_humidity': 0.005841}, + 'u': 0.0576, + 'v': 0.0506} +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.05232, + 'log_surface_pressure': 0.1096, + 'sim_time': 0.0, + 'temperature_variation': 14.73, + 'tracers': {'specific_cloud_ice_water_content': 0.00015690000000000002, + 'specific_cloud_liquid_water_content': 0.00032960000000000004, + 'specific_humidity': 0.003281}, + 'vorticity': 0.2168} +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 = 1e-05 +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.792e-05, + 5.989e-05, + 7.112e-05, + 8.21e-05, + 8.842e-05, + 9.109e-05, + 7.706e-05, + 7.156e-05, + 8.221e-05, + 9.83e-05, + 0.0001427, + 0.0002927, + 0.001, + 0.002714, + 0.006034, + 0.01158, + 0.01969, + 0.04488, + 0.08384, + 0.1362, + 0.2039, + 0.2933, + 0.3992, + 0.4827, + 0.5746, + 0.6994, + 0.8461, + 0.9233, + 1.001, + 1.08, + 1.162, + 1.248, + 1.339, + 1.448, + 1.579, + 1.675, + 1.712] + +# 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.0001361, + 0.0002096, + 0.000281, + 0.0004981, + 0.001449, + 0.004681, + 0.01312, + 0.02983, + 0.05774, + 0.09596, + 0.1486, + 0.211, + 0.2882, + 0.3727, + 0.47, + 0.5779, + 0.6897, + 0.8138, + 0.9033, + 0.9937, + 1.095, + 1.191, + 1.299, + 1.411, + 1.498, + 1.549, + 1.578, + 1.599, + 1.628, + 1.679, + 1.783, + 1.882] + +# Parameters for encode/LevelScale: +# ============================================================================== +encode/LevelScale.keys_to_scale = ['specific_humidity'] +encode/LevelScale.name = None +encode/LevelScale.scales = \ + [8.825e-05, + 7.041e-05, + 0.0001044, + 0.0001832, + 0.0007474, + 0.002588, + 0.007067, + 0.01524, + 0.02838, + 0.04518, + 0.06902, + 0.09718, + 0.1316, + 0.1716, + 0.2182, + 0.2748, + 0.3379, + 0.4073, + 0.4611, + 0.5177, + 0.5798, + 0.6547, + 0.7397, + 0.8362, + 0.9373, + 1.041, + 1.149, + 1.265, + 1.394, + 1.551, + 1.71, + 1.792] + +# 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 land_model/ModalToNodalEmbedding: +# ============================================================================== +land_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \ + @land_model/CombinedFeatures +land_model/ModalToNodalEmbedding.name = None +land_model/ModalToNodalEmbedding.nodal_mapping_module = @land_model/NodalMapping +land_model/ModalToNodalEmbedding.output_transform_module = @IdentityTransform + +# Parameters for sea_ice_model/ModalToNodalEmbedding: +# ============================================================================== +sea_ice_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \ + @sea_ice_model/CombinedFeatures +sea_ice_model/ModalToNodalEmbedding.name = None +sea_ice_model/ModalToNodalEmbedding.nodal_mapping_module = \ + @sea_ice_model/NodalMapping +sea_ice_model/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 = @land_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.name = None +advance/NodalLandSeaIceEmbedding.sea_embedding = @sea_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.sea_ice_embedding = \ + @sea_ice_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.static_vars_ds_path = None + +# Parameters for NodalMapping: +# ============================================================================== +NodalMapping.name = None +NodalMapping.tower_factory = @EpdTower + +# Parameters for land_model/NodalMapping: +# ============================================================================== +land_model/NodalMapping.name = None +land_model/NodalMapping.tower_factory = @surface_model/EpdTower + +# Parameters for sea_ice_model/NodalMapping: +# ============================================================================== +sea_ice_model/NodalMapping.name = None +sea_ice_model/NodalMapping.tower_factory = @surface_model/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 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 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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + '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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + '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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + '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.3036, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'radiation': 0.2867, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 6.483e-06, + 'specific_cloud_ice_water_content_del2': 0.008699, + 'specific_cloud_ice_water_content_dlat': 0.0001996, + 'specific_cloud_ice_water_content_dlon': 0.0001299, + 'specific_cloud_liquid_water_content': 1.706e-05, + 'specific_cloud_liquid_water_content_del2': 0.02237, + 'specific_cloud_liquid_water_content_dlat': 0.0005234, + 'specific_cloud_liquid_water_content_dlon': 0.0003232, + 'specific_humidity': 0.003474, + 'specific_humidity_del2': 0.8351, + 'specific_humidity_dlat': 0.0224, + 'specific_humidity_dlon': 0.01323, + 't': 28.0, + 't_del2': 1795.0, + 't_dlat': 52.11, + 't_dlon': 38.9, + 'u': 0.01924, + 'u_del2': 5.206, + 'u_dlat': 0.1515, + 'u_dlon': 0.09435, + 'v': 0.0102, + 'v_del2': 5.058, + 'v_dlat': 0.08895, + 'v_dlon': 0.1268, + 'z': 0.1496, + 'z_del2': 0.1961, + 'z_dlat': 0.01578, + 'z_dlon': 0.008513} +encoder_data/ShiftAndNormalize.shifts = \ + {'cos_latitude': 0.642, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + '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.123, + 't_del2': 31.568, + 't_dlat': 1.729, + 't_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.04, + 'u_dlat': 0.0, + 'u_dlon': -0.0, + 'v': 0.0, + 'v_del2': -0.003, + '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 land_model/ShiftAndNormalize: +# ============================================================================== +land_model/ShiftAndNormalize.features_to_exclude = () +land_model/ShiftAndNormalize.global_scale = None +land_model/ShiftAndNormalize.name = None +land_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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +land_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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + 'vorticity_dlon': -0.0} + +# Parameters for sea_ice_model/ShiftAndNormalize: +# ============================================================================== +sea_ice_model/ShiftAndNormalize.features_to_exclude = () +sea_ice_model/ShiftAndNormalize.global_scale = None +sea_ice_model/ShiftAndNormalize.name = None +sea_ice_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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +sea_ice_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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + 'vorticity_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.3036, + 'divergence': 0.05232, + 'divergence_del2': 108.7, + 'divergence_dlat': 2.176, + 'divergence_dlon': 1.488, + 'geopotential_at_surface': 0.00926, + 'geopotential_at_surface_del2': 4.511, + 'geopotential_at_surface_dlat': 0.1074, + 'geopotential_at_surface_dlon': 0.08614, + 'land_sea_mask': 0.4403, + 'learned_positional_features': 1.0, + 'log_surface_pressure': 0.1096, + 'log_surface_pressure_del2': 48.61, + 'log_surface_pressure_dlat': 1.095, + 'log_surface_pressure_dlon': 0.9575, + 'memory_divergence': 0.05232, + 'memory_log_surface_pressure': 0.1123, + 'memory_specific_cloud_ice_water_content': 7.845e-06, + 'memory_specific_cloud_liquid_water_content': 1.648e-05, + 'memory_specific_humidity': 0.003298, + 'memory_temperature_variation': 14.73, + 'memory_u': 0.01472, + 'memory_v': 0.01001, + 'memory_vorticity': 0.2168, + 'pressure': 1.643, + 'radiation': 0.2867, + 'sea_ice_cover': 0.387, + 'sea_surface_temperature': 11.93, + 'sin_latitude': 0.7043, + 'specific_cloud_ice_water_content': 7.845e-06, + 'specific_cloud_ice_water_content_del2': 0.01029, + 'specific_cloud_ice_water_content_dlat': 0.000203, + 'specific_cloud_ice_water_content_dlon': 0.0001601, + 'specific_cloud_liquid_water_content': 1.648e-05, + 'specific_cloud_liquid_water_content_del2': 0.02065, + 'specific_cloud_liquid_water_content_dlat': 0.0004055, + 'specific_cloud_liquid_water_content_dlon': 0.000309, + 'specific_humidity': 0.003281, + 'specific_humidity_del2': 0.8403, + 'specific_humidity_dlat': 0.01955, + 'specific_humidity_dlon': 0.01361, + 'surface_embedding': 1.0, + 'temperature_variation': 14.73, + 'temperature_variation_del2': 2386.0, + 'temperature_variation_dlat': 62.26, + 'temperature_variation_dlon': 50.23, + 'u': 0.01472, + 'u_del2': 6.478, + 'u_dlat': 0.682, + 'u_dlon': 0.1073, + 'v': 0.01001, + 'v_del2': 5.028, + 'v_dlat': 0.3098, + 'v_dlon': 0.1431, + 'vorticity': 0.2168, + 'vorticity_del2': 290.5, + 'vorticity_dlat': 5.741, + 'vorticity_dlon': 4.558} +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.642, + 'divergence': -0.0, + 'divergence_del2': 0.006, + 'divergence_dlat': 0.001, + 'divergence_dlon': -0.0, + 'geopotential_at_surface': 0.004, + 'geopotential_at_surface_del2': -0.068, + 'geopotential_at_surface_dlat': -0.011, + 'geopotential_at_surface_dlon': -0.0, + 'land_sea_mask': 0.334, + 'learned_positional_features': 0.0, + 'log_surface_pressure': 1.716, + 'log_surface_pressure_del2': 0.819, + 'log_surface_pressure_dlat': 0.135, + '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.798, + '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': -5.126, + 'temperature_variation_del2': 46.171, + 'temperature_variation_dlat': 5.054, + 'temperature_variation_dlon': 0.0, + 'u': 0.007, + 'u_del2': 0.147, + 'u_dlat': 0.002, + 'u_dlon': 0.0, + 'v': -0.0, + 'v_del2': -0.001, + 'v_dlat': -0.0, + 'v_dlon': 0.0, + 'vorticity': -0.002, + 'vorticity_del2': -0.399, + 'vorticity_dlat': 0.036, + '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 land_model/TakeSurfaceAdjacentSigmaLevel: +# ============================================================================== +land_model/TakeSurfaceAdjacentSigmaLevel.name = None + +# Parameters for sea_ice_model/TakeSurfaceAdjacentSigmaLevel: +# ============================================================================== +sea_ice_model/TakeSurfaceAdjacentSigmaLevel.name = None + +# Parameters for advance/ToModalDiffOperators: +# ============================================================================== +advance/ToModalDiffOperators.name = None + +# Parameters for decoder_model/ToModalDiffOperators: +# ============================================================================== +decoder_model/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 decoder_model/VelocityAndPrognostics: +# ============================================================================== +decoder_model/VelocityAndPrognostics.compute_gradients_module = \ + @ToModalDiffOperators +decoder_model/VelocityAndPrognostics.fields_to_include = \ + ['divergence', + 'vorticity', + 'u', + 'v', + 'temperature_variation', + 'specific_humidity'] +decoder_model/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 land_model/VelocityAndPrognostics: +# ============================================================================== +land_model/VelocityAndPrognostics.features_transform_module = \ + @TakeSurfaceAdjacentSigmaLevel +land_model/VelocityAndPrognostics.fields_to_include = \ + ['temperature_variation', 'specific_humidity'] +land_model/VelocityAndPrognostics.name = None + +# Parameters for sea_ice_model/VelocityAndPrognostics: +# ============================================================================== +sea_ice_model/VelocityAndPrognostics.features_transform_module = \ + @TakeSurfaceAdjacentSigmaLevel +sea_ice_model/VelocityAndPrognostics.fields_to_include = \ + ['temperature_variation', 'specific_humidity'] +sea_ice_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 diff --git a/model/reference_code/paper_configs/stochastic_1_4_deg.gin b/model/reference_code/paper_configs/stochastic_1_4_deg.gin new file mode 100644 index 0000000000000000000000000000000000000000..9bd7fc26b24230a16187902df5c2b696b75bc6b8 --- /dev/null +++ b/model/reference_code/paper_configs/stochastic_1_4_deg.gin @@ -0,0 +1,3442 @@ +# Macros: +# ============================================================================== +ACTIVATION = @gelu +ADVANCE_RANDOMNESS_MODULE = @DictOfGaussianRandomFieldModules +CHECKPOINT_MULTISTEP = False +CORRECTOR_MODULE = @CustomCoordsCorrector +CORRECTOR_SCALE = 0.01 +DATA_FILTER_ATTENUATION = 0.0 +DECODE_RANDOMNESS_MODULE = @NoRandomField +DYCORE_FILTER_ORDER = 3 +DYCORE_GRID = @smaller/GridWithWavenumbers() +DYCORE_INTEGRATOR = @imex_rk_sil3 +DYCORE_TAU = '120 minutes' +ENCODE_RANDOMNESS_MODULE = @DictOfGaussianRandomFieldModules +FIELD_SUBSET = (0, 2, 4, 6, 8, 10, 12, 14, 16, 18) +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 = 8 +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, + @with_grads/FloatDataFeatures, + @without_grads/FloatDataFeatures, + @LearnedPositionalFeatures, + @RandomnessFeatures) +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, + @decoder_model/VelocityAndPrognostics, + @with_grads/FloatDataFeatures, + @without_grads/FloatDataFeatures, + @LearnedPositionalFeatures, + @RandomnessFeatures) +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 encoder_model/CombinedFeatures: +# ============================================================================== +encoder_model/CombinedFeatures.feature_module_names_to_exclude = () +encoder_model/CombinedFeatures.feature_modules = (@RandomnessFeatures,) +encoder_model/CombinedFeatures.features_to_exclude = () +encoder_model/CombinedFeatures.features_transform_module = \ + @encoder_model/SequentialTransform +encoder_model/CombinedFeatures.name = None + +# Parameters for land_model/CombinedFeatures: +# ============================================================================== +land_model/CombinedFeatures.feature_module_names_to_exclude = () +land_model/CombinedFeatures.feature_modules = (@land_model/VelocityAndPrognostics,) +land_model/CombinedFeatures.features_to_exclude = () +land_model/CombinedFeatures.features_transform_module = \ + @land_model/ShiftAndNormalize +land_model/CombinedFeatures.name = None + +# Parameters for sea_ice_model/CombinedFeatures: +# ============================================================================== +sea_ice_model/CombinedFeatures.feature_module_names_to_exclude = () +sea_ice_model/CombinedFeatures.feature_modules = \ + (@sea_ice_model/VelocityAndPrognostics, @sea_ice_model/ForcingFeatures) +sea_ice_model/CombinedFeatures.features_to_exclude = () +sea_ice_model/CombinedFeatures.features_transform_module = \ + @sea_ice_model/ShiftAndNormalize +sea_ice_model/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 orography/DataExponentialFilter: +# ============================================================================== +orography/DataExponentialFilter.attenuation = %DATA_FILTER_ATTENUATION +orography/DataExponentialFilter.cutoff = 0 +orography/DataExponentialFilter.name = None +orography/DataExponentialFilter.order = 1 + +# Parameters for DictOfGaussianRandomFieldModules: +# ============================================================================== +DictOfGaussianRandomFieldModules.clip = 6.0 +DictOfGaussianRandomFieldModules.field_names = \ + ('GRF_001_hours_00084_km', + 'GRF_001_hours_00134_km', + 'GRF_002_hours_00201_km', + 'GRF_003_hours_00292_km', + 'GRF_004_hours_00410_km', + 'GRF_006_hours_00561_km', + 'GRF_008_hours_00751_km', + 'GRF_011_hours_00984_km', + 'GRF_014_hours_01268_km', + 'GRF_018_hours_01610_km', + 'GRF_023_hours_02016_km', + 'GRF_029_hours_02495_km', + 'GRF_035_hours_03054_km', + 'GRF_043_hours_03703_km', + 'GRF_052_hours_04449_km', + 'GRF_062_hours_05304_km', + 'GRF_073_hours_06275_km', + 'GRF_086_hours_07374_km', + 'GRF_100_hours_08612_km', + 'GRF_117_hours_10000_km') +DictOfGaussianRandomFieldModules.field_subset = %FIELD_SUBSET +DictOfGaussianRandomFieldModules.initial_correlation_lengths = \ + ('84 km', + '134 km', + '201 km', + '292 km', + '410 km', + '561 km', + '751 km', + '984 km', + '1268 km', + '1610 km', + '2016 km', + '2495 km', + '3054 km', + '3703 km', + '4449 km', + '5304 km', + '6275 km', + '7374 km', + '8612 km', + '10000 km') +DictOfGaussianRandomFieldModules.initial_correlation_times = \ + ('1 hours', + '1 hours', + '2 hours', + '3 hours', + '4 hours', + '6 hours', + '8 hours', + '11 hours', + '14 hours', + '18 hours', + '23 hours', + '29 hours', + '35 hours', + '43 hours', + '52 hours', + '62 hours', + '73 hours', + '86 hours', + '100 hours', + '117 hours') +DictOfGaussianRandomFieldModules.name = None +DictOfGaussianRandomFieldModules.variances = \ + (1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0) + +# 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 = \ + @FilteredCustomOrography +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.randomness_module = \ + %DECODE_RANDOMNESS_MODULE +DimensionalLearnedPrimitiveToWeatherbenchDecoder.time_axis = 0 + +# Parameters for DimensionalLearnedWeatherbenchToPrimitiveEncoder: +# ============================================================================== +DimensionalLearnedWeatherbenchToPrimitiveEncoder.correction_transform_module = \ + @encode/SequentialTransform +DimensionalLearnedWeatherbenchToPrimitiveEncoder.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'} +DimensionalLearnedWeatherbenchToPrimitiveEncoder.modal_to_nodal_data_features_module = \ + @encoder_data/CombinedFeatures +DimensionalLearnedWeatherbenchToPrimitiveEncoder.modal_to_nodal_model_features_module = \ + @encoder_model/CombinedFeatures +DimensionalLearnedWeatherbenchToPrimitiveEncoder.name = None +DimensionalLearnedWeatherbenchToPrimitiveEncoder.nodal_mapping_module = \ + @NodalMapping +DimensionalLearnedWeatherbenchToPrimitiveEncoder.orography_module = \ + @FilteredCustomOrography +DimensionalLearnedWeatherbenchToPrimitiveEncoder.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} +DimensionalLearnedWeatherbenchToPrimitiveEncoder.randomness_module = \ + %ENCODE_RANDOMNESS_MODULE +DimensionalLearnedWeatherbenchToPrimitiveEncoder.time_axis = 0 +DimensionalLearnedWeatherbenchToPrimitiveEncoder.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 = None +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,) + +# 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 = 10 +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_ice_model/ForcingFeatures: +# ============================================================================== +sea_ice_model/ForcingFeatures.forcing_to_include = ('sea_ice_cover',) +sea_ice_model/ForcingFeatures.name = None + +# 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 smaller/GridWithWavenumbers: +# ============================================================================== +smaller/GridWithWavenumbers.dealiasing = 'quadratic' +smaller/GridWithWavenumbers.latitude_spacing = 'gauss' +smaller/GridWithWavenumbers.longitude_offset = 0.0 +smaller/GridWithWavenumbers.longitude_wavenumbers = 126 +smaller/GridWithWavenumbers.radius = None +smaller/GridWithWavenumbers.spherical_harmonics_impl = \ + @RealSphericalHarmonicsWithZeroImag + +# Parameters for advance/IdentityTransform: +# ============================================================================== +advance/IdentityTransform.name = None + +# Parameters for land_model/IdentityTransform: +# ============================================================================== +land_model/IdentityTransform.name = None + +# Parameters for sea_ice_model/IdentityTransform: +# ============================================================================== +sea_ice_model/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 = 1e-05 +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 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 land_model/ModalToNodalEmbedding: +# ============================================================================== +land_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \ + @land_model/CombinedFeatures +land_model/ModalToNodalEmbedding.name = None +land_model/ModalToNodalEmbedding.nodal_mapping_module = @land_model/NodalMapping +land_model/ModalToNodalEmbedding.output_transform_module = @IdentityTransform + +# Parameters for sea_ice_model/ModalToNodalEmbedding: +# ============================================================================== +sea_ice_model/ModalToNodalEmbedding.modal_to_nodal_features_module = \ + @sea_ice_model/CombinedFeatures +sea_ice_model/ModalToNodalEmbedding.name = None +sea_ice_model/ModalToNodalEmbedding.nodal_mapping_module = \ + @sea_ice_model/NodalMapping +sea_ice_model/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 = @land_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.name = None +advance/NodalLandSeaIceEmbedding.sea_embedding = @sea_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.sea_ice_embedding = \ + @sea_ice_model/ModalToNodalEmbedding +advance/NodalLandSeaIceEmbedding.static_vars_ds_path = None + +# Parameters for NodalMapping: +# ============================================================================== +NodalMapping.name = None +NodalMapping.tower_factory = @EpdTower + +# Parameters for land_model/NodalMapping: +# ============================================================================== +land_model/NodalMapping.name = None +land_model/NodalMapping.tower_factory = @surface_model/EpdTower + +# Parameters for sea_ice_model/NodalMapping: +# ============================================================================== +sea_ice_model/NodalMapping.name = None +sea_ice_model/NodalMapping.tower_factory = @surface_model/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 NoRandomField: +# ============================================================================== +NoRandomField.prefer_nodal = True + +# Parameters for NullFeatures: +# ============================================================================== +NullFeatures.name = None + +# 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 advance/RandomnessFeatures: +# ============================================================================== +advance/RandomnessFeatures.name = None + +# Parameters for decoder_model/RandomnessFeatures: +# ============================================================================== +decoder_model/RandomnessFeatures.name = None + +# Parameters for encoder_model/RandomnessFeatures: +# ============================================================================== +encoder_model/RandomnessFeatures.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 encoder_model/SequentialTransform: +# ============================================================================== +encoder_model/SequentialTransform.name = None +encoder_model/SequentialTransform.transform_modules = \ + (@encoder_model/ShiftAndNormalize,) + +# 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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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 encoder_model/ShiftAndNormalize: +# ============================================================================== +encoder_model/ShiftAndNormalize.features_to_exclude = () +encoder_model/ShiftAndNormalize.global_scale = None +encoder_model/ShiftAndNormalize.name = None +encoder_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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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} +encoder_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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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 land_model/ShiftAndNormalize: +# ============================================================================== +land_model/ShiftAndNormalize.features_to_exclude = () +land_model/ShiftAndNormalize.global_scale = None +land_model/ShiftAndNormalize.name = None +land_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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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} +land_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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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 sea_ice_model/ShiftAndNormalize: +# ============================================================================== +sea_ice_model/ShiftAndNormalize.features_to_exclude = () +sea_ice_model/ShiftAndNormalize.global_scale = None +sea_ice_model/ShiftAndNormalize.name = None +sea_ice_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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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_ice_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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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 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, + 'GRF_001_hours_00084_km': 1.0, + 'GRF_001_hours_00134_km': 1.0, + 'GRF_002_hours_00201_km': 1.0, + 'GRF_003_hours_00292_km': 1.0, + 'GRF_004_hours_00410_km': 1.0, + 'GRF_006_hours_00561_km': 1.0, + 'GRF_008_hours_00751_km': 1.0, + 'GRF_011_hours_00984_km': 1.0, + 'GRF_014_hours_01268_km': 1.0, + 'GRF_018_hours_01610_km': 1.0, + 'GRF_023_hours_02016_km': 1.0, + 'GRF_029_hours_02495_km': 1.0, + 'GRF_035_hours_03054_km': 1.0, + 'GRF_043_hours_03703_km': 1.0, + 'GRF_052_hours_04449_km': 1.0, + 'GRF_062_hours_05304_km': 1.0, + 'GRF_073_hours_06275_km': 1.0, + 'GRF_086_hours_07374_km': 1.0, + 'GRF_100_hours_08612_km': 1.0, + 'GRF_117_hours_10000_km': 1.0, + '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, + '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, + 'GRF_001_hours_00084_km': 0.0, + 'GRF_001_hours_00134_km': 0.0, + 'GRF_002_hours_00201_km': 0.0, + 'GRF_003_hours_00292_km': 0.0, + 'GRF_004_hours_00410_km': 0.0, + 'GRF_006_hours_00561_km': 0.0, + 'GRF_008_hours_00751_km': 0.0, + 'GRF_011_hours_00984_km': 0.0, + 'GRF_014_hours_01268_km': 0.0, + 'GRF_018_hours_01610_km': 0.0, + 'GRF_023_hours_02016_km': 0.0, + 'GRF_029_hours_02495_km': 0.0, + 'GRF_035_hours_03054_km': 0.0, + 'GRF_043_hours_03703_km': 0.0, + 'GRF_052_hours_04449_km': 0.0, + 'GRF_062_hours_05304_km': 0.0, + 'GRF_073_hours_06275_km': 0.0, + 'GRF_086_hours_07374_km': 0.0, + 'GRF_100_hours_08612_km': 0.0, + 'GRF_117_hours_10000_km': 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, + '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 = \ + @DimensionalLearnedWeatherbenchToPrimitiveEncoder +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 = \ + %ADVANCE_RANDOMNESS_MODULE + +# Parameters for land_model/TakeSurfaceAdjacentSigmaLevel: +# ============================================================================== +land_model/TakeSurfaceAdjacentSigmaLevel.name = None + +# Parameters for sea_ice_model/TakeSurfaceAdjacentSigmaLevel: +# ============================================================================== +sea_ice_model/TakeSurfaceAdjacentSigmaLevel.name = None + +# Parameters for advance/ToModalDiffOperators: +# ============================================================================== +advance/ToModalDiffOperators.name = None + +# Parameters for decoder_model/ToModalDiffOperators: +# ============================================================================== +decoder_model/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 = %CHECKPOINT_MULTISTEP +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 decoder_model/VelocityAndPrognostics: +# ============================================================================== +decoder_model/VelocityAndPrognostics.compute_gradients_module = \ + @ToModalDiffOperators +decoder_model/VelocityAndPrognostics.fields_to_include = \ + ['divergence', + 'vorticity', + 'u', + 'v', + 'temperature_variation', + 'specific_humidity'] +decoder_model/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 land_model/VelocityAndPrognostics: +# ============================================================================== +land_model/VelocityAndPrognostics.features_transform_module = \ + @TakeSurfaceAdjacentSigmaLevel +land_model/VelocityAndPrognostics.fields_to_include = \ + ['temperature_variation', 'specific_humidity'] +land_model/VelocityAndPrognostics.name = None + +# Parameters for sea_ice_model/VelocityAndPrognostics: +# ============================================================================== +sea_ice_model/VelocityAndPrognostics.features_transform_module = \ + @TakeSurfaceAdjacentSigmaLevel +sea_ice_model/VelocityAndPrognostics.fields_to_include = \ + ['temperature_variation', 'specific_humidity'] +sea_ice_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') diff --git a/model/reference_code/reader.py b/model/reference_code/reader.py new file mode 100644 index 0000000000000000000000000000000000000000..a62e5e16fb5cb83d7b7041ca189f641b5f48bd8d --- /dev/null +++ b/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) +