diff --git a/configs/baselines/era5-aimip/README.md b/configs/baselines/era5-aimip/README.md new file mode 100644 index 000000000..7aa6ef965 --- /dev/null +++ b/configs/baselines/era5-aimip/README.md @@ -0,0 +1,70 @@ +# ERA5 AIMIP Baseline + +This directory contains scripts and configurations for training and running an ACE2 model on ERA5 +data for the AIMIP evaluation protocol. This configuration is referred to as **ACE2.1-ERA5**. The +model predicts pressure-level diagnostic variables over 1978–2024 with multiple initial conditions +and SST perturbation scenarios. + +## Workflow + +The full experiment proceeds in five stages. Checkpoint IDs embedded in the scripts must be updated +manually after evaluating each stage. + +### 1. Train + +```bash +bash run-ace-train.sh +``` + +Trains 4 random-seed ensemble members (RS0–RS3) on ERA5 1979–2008, with validation on 2009–2014. +Config: `ace-train-config.yaml`. + +### 2. Evaluate training seeds + +```bash +bash run-ace-evaluator-seed-selection.sh +bash run-ace-evaluator-seed-selection-single.sh +``` + +Evaluate all 4 trained checkpoints to select the best seed for fine-tuning. + +- `run-ace-evaluator-seed-selection.sh` — 7x 5-year evaluations (starting in 1980, 1985, 1990, + 1995, 2000, 2005, 2010). Config: `ace-evaluator-seed-selection-config.yaml`. +- `run-ace-evaluator-seed-selection-single.sh` — single continuous 36-year run (1978-10-01 to + 2014-12-31). Config: `ace-evaluator-seed-selection-single-config.yaml`. + +The best seed is chosen based on comparing the time-mean climate and trend skill, both in +the 7x 5-year and 36-year evaluations; this is somewhat subjective. The chosen seed is used +in `run-ace-fine-tune-decoder-pressure-levels.sh`. + +### 3. Fine-tune + +```bash +bash run-ace-fine-tune-decoder-pressure-levels.sh +``` + +Freezes the best trained checkpoint and trains a secondary MLP decoder for 65 pressure-level +diagnostic variables (TMP, Q, UGRD, VGRD, h at 13 pressure levels plus near-surface fields) across +4 new random seeds. Config: `ace-fine-tune-pressure-level-separate-decoder-config.yaml`. + +### 4. Evaluate fine-tuned seeds + +Re-run both evaluator scripts from step 2. The scripts already include checkpoint IDs for both +the trained and fine-tuned ensemble members, enabling direct comparison. After evaluating seeds +similarly as before (though there is little variability due frozen prognostic state), the best +checkpoint ID is used in `run-ace-inference.sh`. + +### 5. Run inference + +```bash +bash run-ace-inference.sh +``` + +Runs 15 parallel 46-year simulations (1978-10-01 to 2024-12-31) using the best fine-tuned +checkpoint: + +- 5 initial conditions (IC1–IC5) from the AIMIP IC dataset +- 3 SST scenarios: baseline, +2 K, +4 K + +Configs: `ace-aimip-inference-config.yaml`, `ace-aimip-inference-p2k-config.yaml`, +`ace-aimip-inference-p4k-config.yaml`. diff --git a/configs/baselines/era5-aimip/ace-aimip-inference-config.yaml b/configs/baselines/era5-aimip/ace-aimip-inference-config.yaml new file mode 100644 index 000000000..360807c14 --- /dev/null +++ b/configs/baselines/era5-aimip/ace-aimip-inference-config.yaml @@ -0,0 +1,671 @@ +checkpoint_path: /ckpt.tar +data_writer: + save_monthly_files: false + save_prediction_files: false + files: + - label: ta_day_ACE2-ERA5_aimip_r1i1p1f1_gr_19781001-19791231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip_r1i1p1f1_gr_20240101-20241231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip_r1i1p1f1_gr_19781001-19791231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip_r1i1p1f1_gr_20240101-20241231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip_r1i1p1f1_gr_19781001-19791231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip_r1i1p1f1_gr_20240101-20241231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - eastward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip_r1i1p1f1_gr_19781001-19791231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip_r1i1p1f1_gr_20240101-20241231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - northward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip_r1i1p1f1_gr_19781001-19791231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip_r1i1p1f1_gr_20240101-20241231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - PRESsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - PRESsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - surface_temperature + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - surface_temperature + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - TMP2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - TMP2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - Q2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - Q2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - UGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - UGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - VGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - VGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip_r1i1p1f1_gn_19781001-19791231 + names: + - PRATEsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip_r1i1p1f1_gn_20240101-20241231 + names: + - PRATEsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip_r1i1p1f1_gr_197810-202412 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip_r1i1p1f1_gr_197810-202412 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip_r1i1p1f1_gr_197810-202412 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip_r1i1p1f1_gr_197810-202412 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: zg_Amon_ACE2-ERA5_aimip_r1i1p1f1_gr_197810-202412 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ps_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - PRESsfc + time_coarsen: + method: monthly_mean + save_reference: false + - label: ts_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - surface_temperature + time_coarsen: + method: monthly_mean + save_reference: false + - label: tas_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - TMP2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: huss_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - Q2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: uas_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - UGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: vas_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - VGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: pr_Amon_ACE2-ERA5_aimip_r1i1p1f1_gn_197810-202412 + names: + - PRATEsfc + time_coarsen: + method: monthly_mean + save_reference: false +experiment_dir: /results +forward_steps_in_memory: 40 +n_forward_steps: 67576 +initial_condition: + path: /climate-default/2025-09-12-aimip-evaluation/aimip-evaluation-ics-v3/1978-09-30_IC0.nc + start_indices: + times: + - '1978-09-30T18:00:00' +forcing_loader: + dataset: + data_path: /climate-default/2025-09-12-aimip-evaluation + file_pattern: aimip-era5-1deg-forcing-1978-2024-repeat-first.zarr + engine: zarr + num_data_workers: 8 +logging: + project: ace \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-aimip-inference-p2k-config.yaml b/configs/baselines/era5-aimip/ace-aimip-inference-p2k-config.yaml new file mode 100644 index 000000000..88f266d21 --- /dev/null +++ b/configs/baselines/era5-aimip/ace-aimip-inference-p2k-config.yaml @@ -0,0 +1,676 @@ +checkpoint_path: /ckpt.tar +data_writer: + save_monthly_files: false + save_prediction_files: false + files: + - label: ta_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_19781001-19791231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_20240101-20241231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_19781001-19791231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_20240101-20241231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_19781001-19791231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_20240101-20241231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - eastward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_19781001-19791231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_20240101-20241231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - northward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_19781001-19791231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_20240101-20241231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - PRESsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - PRESsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - surface_temperature + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - surface_temperature + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - TMP2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - TMP2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - Q2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - Q2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - UGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - UGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - VGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - VGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_19781001-19791231 + names: + - PRATEsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_20240101-20241231 + names: + - PRATEsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_197810-202412 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_197810-202412 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_197810-202412 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_197810-202412 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: zg_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gr_197810-202412 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ps_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - PRESsfc + time_coarsen: + method: monthly_mean + save_reference: false + - label: ts_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - surface_temperature + time_coarsen: + method: monthly_mean + save_reference: false + - label: tas_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - TMP2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: huss_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - Q2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: uas_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - UGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: vas_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - VGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: pr_Amon_ACE2-ERA5_aimip-p2k_r1i1p1f1_gn_197810-202412 + names: + - PRATEsfc + time_coarsen: + method: monthly_mean + save_reference: false +experiment_dir: /results +forward_steps_in_memory: 40 +n_forward_steps: 67576 +initial_condition: + path: /climate-default/2025-09-12-aimip-evaluation/aimip-evaluation-ics-v3/1978-09-30_IC0.nc + start_indices: + times: + - '1978-09-30T18:00:00' +forcing_loader: + dataset: + data_path: /climate-default/2025-09-12-aimip-evaluation + file_pattern: aimip-era5-1deg-forcing-1978-2024-repeat-first.zarr + engine: zarr + num_data_workers: 8 + perturbations: + sst: + - type: constant + config: + amplitude: 2.0 +logging: + project: ace \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-aimip-inference-p4k-config.yaml b/configs/baselines/era5-aimip/ace-aimip-inference-p4k-config.yaml new file mode 100644 index 000000000..d301e044e --- /dev/null +++ b/configs/baselines/era5-aimip/ace-aimip-inference-p4k-config.yaml @@ -0,0 +1,676 @@ +checkpoint_path: /ckpt.tar +data_writer: + save_monthly_files: false + save_prediction_files: false + files: + - label: ta_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_19781001-19791231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_20240101-20241231 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_19781001-19791231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_20240101-20241231 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: hus_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_19781001-19791231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_20240101-20241231 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ua_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - eastward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_19781001-19791231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_20240101-20241231 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: va_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - northward_wind_8 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_19781001-19791231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: zg_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_20240101-20241231 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - PRESsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ps_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - PRESsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - surface_temperature + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ts_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - surface_temperature + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - TMP2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: tas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - TMP2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - Q2m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: huss_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - Q2m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - UGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: uas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - UGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - VGRD10m + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: vas_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - VGRD10m + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_19781001-19791231 + names: + - PRATEsfc + time_selection: + start_time: '1978-10-01T00:00:00' + stop_time: '1979-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: pr_day_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_20240101-20241231 + names: + - PRATEsfc + time_selection: + start_time: '2024-01-01T00:00:00' + stop_time: '2024-12-31T18:00:00' + time_coarsen: + coarsen_factor: 4 + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_197810-202412 + names: + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ta_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_197810-202412 + names: + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: hus_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_197810-202412 + names: + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ua_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_197810-202412 + names: + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: va_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + time_coarsen: + method: monthly_mean + save_reference: false + - label: zg_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gr_197810-202412 + names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + time_coarsen: + method: monthly_mean + save_reference: false + - label: ps_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - PRESsfc + time_coarsen: + method: monthly_mean + save_reference: false + - label: ts_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - surface_temperature + time_coarsen: + method: monthly_mean + save_reference: false + - label: tas_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - TMP2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: huss_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - Q2m + time_coarsen: + method: monthly_mean + save_reference: false + - label: uas_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - UGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: vas_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - VGRD10m + time_coarsen: + method: monthly_mean + save_reference: false + - label: pr_Amon_ACE2-ERA5_aimip-p4k_r1i1p1f1_gn_197810-202412 + names: + - PRATEsfc + time_coarsen: + method: monthly_mean + save_reference: false +experiment_dir: /results +forward_steps_in_memory: 40 +n_forward_steps: 67576 +initial_condition: + path: /climate-default/2025-09-12-aimip-evaluation/aimip-evaluation-ics-v3/1978-09-30_IC0.nc + start_indices: + times: + - '1978-09-30T18:00:00' +forcing_loader: + dataset: + data_path: /climate-default/2025-09-12-aimip-evaluation + file_pattern: aimip-era5-1deg-forcing-1978-2024-repeat-first.zarr + engine: zarr + num_data_workers: 8 + perturbations: + sst: + - type: constant + config: + amplitude: 4.0 +logging: + project: ace \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-evaluator-seed-selection-config.yaml b/configs/baselines/era5-aimip/ace-evaluator-seed-selection-config.yaml new file mode 100644 index 000000000..14f5d3deb --- /dev/null +++ b/configs/baselines/era5-aimip/ace-evaluator-seed-selection-config.yaml @@ -0,0 +1,33 @@ +experiment_dir: /results +n_forward_steps: 7300 +forward_steps_in_memory: 10 +checkpoint_path: /ckpt.tar +logging: + log_to_screen: true + log_to_wandb: true + log_to_file: true + project: ace + entity: ai2cm +loader: + start_indices: + times: + - '1980-01-01T00:00:00' + - '1985-01-01T00:00:00' + - '1990-01-01T00:00:00' + - '1995-01-01T00:00:00' + - '2000-01-01T00:00:00' + - '2005-01-01T00:00:00' + - '2010-01-01T00:00:00' + dataset: + merge: + - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + - data_path: /climate-default/2025-11-10-era5-1deg-pressure-level-1940-2022 + file_pattern: 2025-11-10-era5-1deg-pressure-level-1940-2022.zarr + engine: "zarr" + num_data_workers: 4 +aggregator: + log_zonal_mean_images: true + log_histograms: true +data_writer: + save_prediction_files: false + save_monthly_files: false \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-evaluator-seed-selection-single-config.yaml b/configs/baselines/era5-aimip/ace-evaluator-seed-selection-single-config.yaml new file mode 100644 index 000000000..058a5b183 --- /dev/null +++ b/configs/baselines/era5-aimip/ace-evaluator-seed-selection-single-config.yaml @@ -0,0 +1,26 @@ +experiment_dir: /results +n_forward_steps: 52960 +forward_steps_in_memory: 50 +checkpoint_path: /ckpt.tar +logging: + log_to_screen: true + log_to_wandb: true + log_to_file: true + project: ace + entity: ai2cm +loader: + start_indices: + times: + - '1978-10-01T00:00:00' + dataset: + merge: + - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + - data_path: /climate-default/2025-11-10-era5-1deg-pressure-level-1940-2022 + file_pattern: 2025-11-10-era5-1deg-pressure-level-1940-2022.zarr + engine: "zarr" + num_data_workers: 16 +aggregator: + log_zonal_mean_images: false +data_writer: + save_prediction_files: false + save_monthly_files: false \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-fine-tune-pressure-level-separate-decoder-config.yaml b/configs/baselines/era5-aimip/ace-fine-tune-pressure-level-separate-decoder-config.yaml new file mode 100644 index 000000000..d2ae103b3 --- /dev/null +++ b/configs/baselines/era5-aimip/ace-fine-tune-pressure-level-separate-decoder-config.yaml @@ -0,0 +1,323 @@ +experiment_dir: /results +save_checkpoint: true +validate_using_ema: true +max_epochs: 120 +n_forward_steps: 2 +seed: 0 +ema: + decay: 0.999 +inference: + n_forward_steps: 7300 + forward_steps_in_memory: 40 + loader: + start_indices: + times: + - '2009-01-01T00:00:00' + - '2009-02-15T00:00:00' + - '2009-04-01T00:00:00' + - '2009-05-15T00:00:00' + - '2009-07-01T00:00:00' + - '2009-08-15T00:00:00' + - '2009-10-01T00:00:00' + - '2009-11-15T00:00:00' + dataset: + merge: + - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + - data_path: /climate-default/2025-11-10-era5-1deg-pressure-level-1940-2022 + file_pattern: 2025-11-10-era5-1deg-pressure-level-1940-2022.zarr + engine: "zarr" + num_data_workers: 4 + aggregator: + log_histograms: true + time_mean_reference_data: /statsdata/time-mean.nc +logging: + log_to_screen: true + log_to_wandb: true + log_to_file: true + project: ace + entity: ai2cm +train_loader: + batch_size: 16 + num_data_workers: 4 + prefetch_factor: 4 + dataset: + merge: + - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + subset: + start_time: '1979-01-01' + stop_time: '2008-12-31' + - data_path: /climate-default/2025-11-10-era5-1deg-pressure-level-1940-2022 + subset: + start_time: '1979-01-01' + stop_time: '2008-12-31' + file_pattern: 2025-11-10-era5-1deg-pressure-level-1940-2022.zarr + engine: "zarr" +validation_loader: + batch_size: 128 + num_data_workers: 4 + prefetch_factor: 4 + dataset: + merge: + - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + subset: + start_time: '2009-01-01' + stop_time: '2014-12-31' + step: 3 + - data_path: /climate-default/2025-11-10-era5-1deg-pressure-level-1940-2022 + file_pattern: 2025-11-10-era5-1deg-pressure-level-1940-2022.zarr + engine: "zarr" + subset: + start_time: '2009-01-01' + stop_time: '2014-12-31' + step: 3 +optimization: + use_gradient_accumulation: true + enable_automatic_mixed_precision: false + lr: 0.0001 + optimizer_type: AdamW + kwargs: + fused: true + weight_decay: 0.01 + scheduler: + schedulers: + - type: LinearLR + step_each_iteration: true + kwargs: + start_factor: 0.01 + end_factor: 1 + total_iters: 13695 # 5 epochs + - type: ConstantLR + step_each_iteration: true + milestones: [13695] +stepper_training: + loss: + type: MSE + weights: + air_temperature_0: 0.5 + air_temperature_1: 0.5 + eastward_wind_0: 0.5 + northward_wind_0: 0.5 + specific_total_water_0: 0.5 + specific_total_water_1: 0.25 + specific_total_water_2: 0.5 + PRATEsfc: 0.5 + DLWRFsfc: 2 + ULWRFsfc: 5 + USWRFsfc: 2 + DSWRFsfc: 2 + USWRFtoa: 2 + tendency_of_total_water_path_due_to_advection: 0.5 + parameter_init: + weights_path: /base_weights/ckpt.tar + parameters: + - frozen: + include: + - "*" +stepper: + step: + type: single_module + config: + builder: + type: NoiseConditionedSFNO + config: + embed_dim: 384 + noise_embed_dim: 0 + filter_type: linear + use_mlp: true + num_layers: 8 + operator_type: dhconv + separable: false + spectral_layers: 3 + spectral_transform: sht + normalization: + network: + global_means_path: /statsdata/centering.nc + global_stds_path: /statsdata/scaling-full-field.nc + residual: + global_means_path: /statsdata/centering.nc + global_stds_path: /statsdata/scaling-residual.nc + secondary_decoder: + secondary_diagnostic_names: + - h50 + - h100 + - h150 + - h200 + - h250 + - h300 + - h400 + - h500 + - h600 + - h700 + - h850 + - h925 + - h1000 + - TMP50 + - TMP100 + - TMP150 + - TMP200 + - TMP250 + - TMP300 + - TMP400 + - TMP500 + - TMP600 + - TMP700 + - TMP850 + - TMP925 + - TMP1000 + - TMP2m + - Q50 + - Q100 + - Q150 + - Q200 + - Q250 + - Q300 + - Q400 + - Q500 + - Q600 + - Q700 + - Q850 + - Q925 + - Q1000 + - Q2m + - UGRD50 + - UGRD100 + - UGRD150 + - UGRD200 + - UGRD250 + - UGRD300 + - UGRD400 + - UGRD500 + - UGRD600 + - UGRD700 + - UGRD850 + - UGRD925 + - UGRD1000 + - UGRD10m + - VGRD50 + - VGRD100 + - VGRD150 + - VGRD200 + - VGRD250 + - VGRD300 + - VGRD400 + - VGRD500 + - VGRD600 + - VGRD700 + - VGRD850 + - VGRD925 + - VGRD1000 + - VGRD10m + network: + type: MLP + config: + hidden_dim: 256 + depth: 2 + ocean: + surface_temperature_name: surface_temperature + ocean_fraction_name: ocean_fraction + corrector: + conserve_dry_air: true + moisture_budget_correction: advection_and_precipitation + force_positive_names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - PRATEsfc + - ULWRFsfc + - ULWRFtoa + - DLWRFsfc + - DSWRFsfc + - USWRFsfc + - USWRFtoa + next_step_forcing_names: + - DSWRFtoa + in_names: + - land_fraction + - ocean_fraction + - sea_ice_fraction + - DSWRFtoa + - HGTsfc + - PRESsfc + - surface_temperature + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + out_names: + - PRESsfc + - surface_temperature + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - ULWRFsfc + - ULWRFtoa + - DLWRFsfc + - DSWRFsfc + - USWRFsfc + - USWRFtoa + - tendency_of_total_water_path_due_to_advection \ No newline at end of file diff --git a/configs/baselines/era5-aimip/ace-train-config.yaml b/configs/baselines/era5-aimip/ace-train-config.yaml new file mode 100644 index 000000000..ad712e412 --- /dev/null +++ b/configs/baselines/era5-aimip/ace-train-config.yaml @@ -0,0 +1,211 @@ +experiment_dir: /results +save_checkpoint: true +validate_using_ema: true +max_epochs: 120 +n_forward_steps: 2 +seed: 0 +ema: + decay: 0.999 +inference: + n_forward_steps: 7300 + forward_steps_in_memory: 40 + loader: + start_indices: + times: + - '2009-01-01T00:00:00' + - '2009-02-15T00:00:00' + - '2009-04-01T00:00:00' + - '2009-05-15T00:00:00' + - '2009-07-01T00:00:00' + - '2009-08-15T00:00:00' + - '2009-10-01T00:00:00' + - '2009-11-15T00:00:00' + dataset: + data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + num_data_workers: 16 + aggregator: + log_histograms: true + time_mean_reference_data: /statsdata/time-mean.nc +logging: + log_to_screen: true + log_to_wandb: true + log_to_file: true + project: ace + entity: ai2cm +train_loader: + batch_size: 16 + num_data_workers: 16 + prefetch_factor: 4 + dataset: + data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + subset: + start_time: '1979-01-01' + stop_time: '2008-12-31' +validation_loader: + batch_size: 128 + num_data_workers: 16 + prefetch_factor: 4 + dataset: + data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs + subset: + start_time: '2009-01-01' + stop_time: '2014-12-31' + step: 3 +optimization: + use_gradient_accumulation: true + enable_automatic_mixed_precision: false + lr: 0.0001 + optimizer_type: AdamW + kwargs: + fused: true + weight_decay: 0.01 +stepper_training: + loss: + type: MSE + weights: + air_temperature_0: 0.5 + air_temperature_1: 0.5 + eastward_wind_0: 0.5 + northward_wind_0: 0.5 + specific_total_water_0: 0.5 + specific_total_water_1: 0.25 + specific_total_water_2: 0.5 + PRATEsfc: 0.5 + DLWRFsfc: 2 + ULWRFsfc: 5 + USWRFsfc: 2 + DSWRFsfc: 2 + USWRFtoa: 2 + tendency_of_total_water_path_due_to_advection: 0.5 +stepper: + step: + type: single_module + config: + builder: + type: NoiseConditionedSFNO + config: + embed_dim: 384 + noise_embed_dim: 0 + filter_type: linear + use_mlp: true + num_layers: 8 + operator_type: dhconv + separable: false + spectral_layers: 3 + spectral_transform: sht + normalization: + network: + global_means_path: /statsdata/centering.nc + global_stds_path: /statsdata/scaling-full-field.nc + residual: + global_means_path: /statsdata/centering.nc + global_stds_path: /statsdata/scaling-residual.nc + ocean: + surface_temperature_name: surface_temperature + ocean_fraction_name: ocean_fraction + corrector: + conserve_dry_air: true + moisture_budget_correction: advection_and_precipitation + force_positive_names: + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - PRATEsfc + - ULWRFsfc + - ULWRFtoa + - DLWRFsfc + - DSWRFsfc + - USWRFsfc + - USWRFtoa + next_step_forcing_names: + - DSWRFtoa + in_names: + - land_fraction + - ocean_fraction + - sea_ice_fraction + - DSWRFtoa + - HGTsfc + - PRESsfc + - surface_temperature + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + out_names: + - PRESsfc + - surface_temperature + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + - specific_total_water_0 + - specific_total_water_1 + - specific_total_water_2 + - specific_total_water_3 + - specific_total_water_4 + - specific_total_water_5 + - specific_total_water_6 + - specific_total_water_7 + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - ULWRFsfc + - ULWRFtoa + - DLWRFsfc + - DSWRFsfc + - USWRFsfc + - USWRFtoa + - tendency_of_total_water_path_due_to_advection diff --git a/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection-single.sh b/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection-single.sh new file mode 100755 index 000000000..b87efdee1 --- /dev/null +++ b/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection-single.sh @@ -0,0 +1,77 @@ +#!/bin/bash + +set -e + +JOB_NAME_BASE="ace-aimip-evaluator-1979-2014" +JOB_GROUP="ace-aimip" +SEED_CHECKPOINT_IDS=("01K9B1MR70QWN90KNY7NM22K5M" \ + "01K9B1MT4QY1ZEZPPS53G2SXPK" \ + "01K9B1MVP3VS3NEABHT0W151AX" \ + "01K9B1MXD6V26S8BQH5CKY514C" \ + ) +FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS=("01KAKXY0EK24K7BZK2N8SPJ5SJ"\ + "01KAVVAKANNYY096MYCGSZ7RMQ" \ + "01KAVVGKY28P5N1VA883C63EBY" \ + "01KAVVN8YPPB3P6ZSD0BGCCVX7" +) +CONFIG_FILENAME="ace-evaluator-seed-selection-single-config.yaml" +SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository +CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME +BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name') +REPO_ROOT=$(git rev-parse --show-toplevel) + +cd $REPO_ROOT # so config path is valid no matter where we are running this script + +python -m fme.ace.validate_config --config_type evaluator $CONFIG_PATH + +launch_job () { + + JOB_NAME=$1 + MODEL_CHECKPOINT_DATASET=$2 + shift 2 + OVERRIDE="$@" + echo $OVERRIDE + + cd $REPO_ROOT && gantry run \ + --name $JOB_NAME \ + --task-name $JOB_NAME \ + --description 'Run ACE2-ERA5 evaluation' \ + --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \ + --workspace ai2/ace \ + --priority high \ + --not-preemptible \ + --cluster ai2/saturn-cirrascale \ + --cluster ai2/ceres-cirrascale \ + --cluster ai2/titan-cirrascale \ + --cluster ai2/jupiter-cirrascale-2 \ + --env WANDB_USERNAME=$BEAKER_USERNAME \ + --env WANDB_NAME=$JOB_NAME \ + --env WANDB_JOB_TYPE=inference \ + --env WANDB_RUN_GROUP=$JOB_GROUP \ + --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \ + --env-secret WANDB_API_KEY=wandb-api-key-ai2cm-sa \ + --dataset-secret google-credentials:/tmp/google_application_credentials.json \ + --dataset $MODEL_CHECKPOINT_DATASET:training_checkpoints/best_inference_ckpt.tar:/ckpt.tar \ + --gpus 1 \ + --shared-memory 200GiB \ + --weka climate-default:/climate-default \ + --budget ai2/climate \ + --system-python \ + --install "pip install --no-deps ." \ + -- python -I -m fme.ace.evaluator $CONFIG_PATH --override $OVERRIDE + +} + +# pre-trained +for (( i=0; i<${#SEED_CHECKPOINT_IDS[@]}; i++ )); do + JOB_NAME="$JOB_NAME_BASE-RS$i" + echo "Launching job $JOB_NAME checkpoint ID: ${SEED_CHECKPOINT_IDS[$i]}" + launch_job "$JOB_NAME" "${SEED_CHECKPOINT_IDS[$i]}" +done + +# fine-tuned with separate decoder +for (( i=0; i<${#FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[@]}; i++ )); do + JOB_NAME="$JOB_NAME_BASE-RS3-pressure-level-fine-tuned-separate-decoder-RS$i" + echo "Launching job for fine-tuned with separate decoder seed $i checkpoint ID: ${FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[$i]}" + launch_job "$JOB_NAME" "${FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[$i]}" +done diff --git a/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection.sh b/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection.sh new file mode 100755 index 000000000..75862fb85 --- /dev/null +++ b/configs/baselines/era5-aimip/run-ace-evaluator-seed-selection.sh @@ -0,0 +1,73 @@ +#!/bin/bash + +set -e + +JOB_NAME_BASE="ace-aimip-evaluator-seed-selection" +JOB_GROUP="ace-aimip" +SEED_CHECKPOINT_IDS=("01K9B1MR70QWN90KNY7NM22K5M" \ + "01K9B1MT4QY1ZEZPPS53G2SXPK" \ + "01K9B1MVP3VS3NEABHT0W151AX" \ + "01K9B1MXD6V26S8BQH5CKY514C" \ + ) +FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS=("01KAKXY0EK24K7BZK2N8SPJ5SJ"\ + "01KAVVAKANNYY096MYCGSZ7RMQ" \ + "01KAVVGKY28P5N1VA883C63EBY" \ + "01KAVVN8YPPB3P6ZSD0BGCCVX7" +) +CONFIG_FILENAME="ace-evaluator-seed-selection-config.yaml" +SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository +CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME +BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name') +REPO_ROOT=$(git rev-parse --show-toplevel) + +cd $REPO_ROOT # so config path is valid no matter where we are running this script + +python -m fme.ace.validate_config --config_type evaluator $CONFIG_PATH + +launch_job () { + local JOB_NAME=$1 + local SEED_CHECKPOINT_ID=$2 + + cd $REPO_ROOT && gantry run \ + --name $JOB_NAME \ + --task-name $JOB_NAME \ + --description 'Run ACE2-ERA5 evaluator for AIMIP seed selection' \ + --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \ + --workspace ai2/ace \ + --priority high \ + --not-preemptible \ + --cluster ai2/saturn-cirrascale \ + --cluster ai2/ceres-cirrascale \ + --cluster ai2/titan-cirrascale \ + --cluster ai2/jupiter-cirrascale-2 \ + --env WANDB_USERNAME=$BEAKER_USERNAME \ + --env WANDB_NAME=$JOB_NAME \ + --env WANDB_JOB_TYPE=inference \ + --env WANDB_RUN_GROUP=$JOB_GROUP \ + --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \ + --env-secret WANDB_API_KEY=wandb-api-key-ai2cm-sa \ + --dataset-secret google-credentials:/tmp/google_application_credentials.json \ + --dataset $SEED_CHECKPOINT_ID:training_checkpoints/best_inference_ckpt.tar:/ckpt.tar \ + --gpus 1 \ + --shared-memory 100GiB \ + --weka climate-default:/climate-default \ + --budget ai2/climate \ + --system-python \ + --install "pip install --no-deps ." \ + -- python -I -m fme.ace.evaluator $CONFIG_PATH + + } + +# pre-trained +for (( i=0; i<${#SEED_CHECKPOINT_IDS[@]}; i++ )); do + JOB_NAME="$JOB_NAME_BASE-RS$i" + echo "Launching job for seed $i checkpoint ID: ${SEED_CHECKPOINT_IDS[$i]}" + launch_job "$JOB_NAME" "${SEED_CHECKPOINT_IDS[$i]}" +done + +# fine-tuned with separate decoder +for (( i=0; i<${#FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[@]}; i++ )); do + JOB_NAME="$JOB_NAME_BASE-RS3-pressure-level-fine-tuned-separate-decoder-RS$i" + echo "Launching job for fine-tuned with separate decoder seed $i checkpoint ID: ${FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[$i]}" + launch_job "$JOB_NAME" "${FINE_TUNED_SEPARATE_DECODER_CHECKPOINT_IDS[$i]}" +done diff --git a/configs/baselines/era5-aimip/run-ace-fine-tune-decoder-pressure-levels.sh b/configs/baselines/era5-aimip/run-ace-fine-tune-decoder-pressure-levels.sh new file mode 100755 index 000000000..17e6307aa --- /dev/null +++ b/configs/baselines/era5-aimip/run-ace-fine-tune-decoder-pressure-levels.sh @@ -0,0 +1,59 @@ +#!/bin/bash + +set -e + +JOB_NAME_BASE="ace-aimip-fine-tune-decoder-pressure-levels" +JOB_GROUP="ace-aimip" +PRESSURE_LEVEL_SEPARATE_DECODER_LR_WARMUP_CONFIG_FILENAME="ace-fine-tune-pressure-level-separate-decoder-config.yaml" +SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository +PRESSURE_LEVEL_SEPARATE_DECODER_LR_WARMUP_CONFIG_PATH=$SCRIPT_PATH/$PRESSURE_LEVEL_SEPARATE_DECODER_LR_WARMUP_CONFIG_FILENAME +EXISTING_RESULTS_DATASET="01K9B1MXD6V26S8BQH5CKY514C" # best checkpoint is ace-aimip-train-rs3 +BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name') +REPO_ROOT=$(git rev-parse --show-toplevel) +N_GPUS=4 + +cd $REPO_ROOT # so config path is valid no matter where we are running this script + +python -m fme.ace.validate_config --config_type train $PRESSURE_LEVEL_SEPARATE_DECODER_LR_WARMUP_CONFIG_PATH + +launch_job () { + + JOB_NAME=$1 + CONFIG_FILENAME=$2 + shift 2 + OVERRIDE="$@" + + gantry run \ + --name $JOB_NAME \ + --task-name $JOB_NAME \ + --description 'Fine-tune ACE decoder outputs on AIMIP period' \ + --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \ + --workspace ai2/climate-titan \ + --priority urgent \ + --preemptible \ + --cluster ai2/titan-cirrascale \ + --env WANDB_USERNAME=$BEAKER_USERNAME \ + --env WANDB_NAME=$JOB_NAME \ + --env WANDB_JOB_TYPE=training \ + --env WANDB_RUN_GROUP=$JOB_GROUP \ + --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \ + --env-secret WANDB_API_KEY=wandb-api-key-ai2cm-sa \ + --dataset-secret google-credentials:/tmp/google_application_credentials.json \ + --dataset brianhenn/era5-1deg-8layer-pressure-level-stats-1990-2019-v2:/statsdata \ + --dataset $EXISTING_RESULTS_DATASET:training_checkpoints/best_inference_ckpt.tar:/base_weights/ckpt.tar \ + --gpus $N_GPUS \ + --shared-memory 400GiB \ + --weka climate-default:/climate-default \ + --budget ai2/climate \ + --system-python \ + --install "pip install --no-deps ." \ + -- torchrun --nproc_per_node $N_GPUS -m fme.ace.train $CONFIG_FILENAME --override $OVERRIDE + +} + +# fine tune with separate decoder for pressure levels, with LR warmup +for SEED in 0 1 2 3; do + JOB_NAME="${JOB_NAME_BASE}-separate-decoder-lr-warmup-RS${SEED}" + OVERRIDE="seed=${SEED}" + launch_job $JOB_NAME $PRESSURE_LEVEL_SEPARATE_DECODER_LR_WARMUP_CONFIG_PATH $OVERRIDE +done diff --git a/configs/baselines/era5-aimip/run-ace-inference.sh b/configs/baselines/era5-aimip/run-ace-inference.sh new file mode 100755 index 000000000..9f5720f3e --- /dev/null +++ b/configs/baselines/era5-aimip/run-ace-inference.sh @@ -0,0 +1,89 @@ +#!/bin/bash + +set -e + +JOB_NAME_BASE="ace-aimip-inference-oct-1978-2024" +JOB_GROUP="ace-aimip" +# this is from ace-aimip-fine-tune-decoder-pressure-levels-separate-decoder-lr-warmup-RS0 +EXISTING_RESULTS_DATASET="01KAKXY0EK24K7BZK2N8SPJ5SJ" +BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name') + +SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository +REPO_ROOT=$(git rev-parse --show-toplevel) + +AIMIP_INFERENCE_CONFIG_FILENAME="ace-aimip-inference-config.yaml" +AIMIP_INFERENCE_BASE_CONFIG_PATH=$SCRIPT_PATH/$AIMIP_INFERENCE_CONFIG_FILENAME +AIMIP_INFERENCE_P2K_CONFIG_FILENAME="ace-aimip-inference-p2k-config.yaml" +AIMIP_INFERENCE_BASE_P2K_CONFIG_PATH=$SCRIPT_PATH/$AIMIP_INFERENCE_P2K_CONFIG_FILENAME +AIMIP_INFERENCE_P4K_CONFIG_FILENAME="ace-aimip-inference-p4k-config.yaml" +AIMIP_INFERENCE_BASE_P4K_CONFIG_PATH=$SCRIPT_PATH/$AIMIP_INFERENCE_P4K_CONFIG_FILENAME + +cd $REPO_ROOT # so config path is valid no matter where we are running this script + +python -m fme.ace.validate_config --config_type inference $AIMIP_INFERENCE_BASE_CONFIG_PATH +python -m fme.ace.validate_config --config_type inference $AIMIP_INFERENCE_BASE_P2K_CONFIG_PATH +python -m fme.ace.validate_config --config_type inference $AIMIP_INFERENCE_BASE_P4K_CONFIG_PATH + +launch_job () { + + local JOB_NAME=$1 + local TEMPLATE_CONFIG=$2 + local IC=$3 + local OVERRIDE=$4 + + cd $REPO_ROOT && gantry run \ + --name $JOB_NAME \ + --task-name $JOB_NAME \ + --description 'Run ACE2-ERA5 inference' \ + --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \ + --workspace ai2/ace \ + --priority high \ + --not-preemptible \ + --cluster ai2/ceres-cirrascale \ + --cluster ai2/titan-cirrascale \ + --cluster ai2/saturn-cirrascale \ + --cluster ai2/jupiter-cirrascale-2 \ + --env WANDB_USERNAME=$BEAKER_USERNAME \ + --env WANDB_NAME=$JOB_NAME \ + --env WANDB_JOB_TYPE=inference \ + --env WANDB_RUN_GROUP=$JOB_GROUP \ + --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \ + --env-secret WANDB_API_KEY=wandb-api-key-ai2cm-sa \ + --dataset-secret google-credentials:/tmp/google_application_credentials.json \ + --dataset $EXISTING_RESULTS_DATASET:training_checkpoints/best_inference_ckpt.tar:/ckpt.tar \ + --gpus 1 \ + --shared-memory 50GiB \ + --weka climate-default:/climate-default \ + --budget ai2/climate \ + --system-python \ + --install "pip install --no-deps ." \ + -- bash -c "sed 's/_r[0-9]i/_r${IC}i/g' ${TEMPLATE_CONFIG} > /tmp/ic-config.yaml && python -I -m fme.ace.inference /tmp/ic-config.yaml --override ${OVERRIDE}" + +} + +# launch 46-year (1979-2024) with spinup from 1978-10-01 +# use 5 different initial conditions files +for IC in {1..5}; do + JOB_NAME="${JOB_NAME_BASE}-IC${IC}" + IC_PATH="/climate-default/2025-09-12-aimip-evaluation/aimip-evaluation-ics-v3/1978-09-30_IC$(( IC - 1 )).nc" # files are 0-indexed + OUTPUT_PATH="/climate-default/2025-11-25-ace-aimip-inference-results/${JOB_NAME}" + OVERRIDE="initial_condition.path=${IC_PATH} experiment_dir=${OUTPUT_PATH}" + echo "Launching job $JOB_NAME with override: $OVERRIDE" + launch_job "$JOB_NAME" "$AIMIP_INFERENCE_BASE_CONFIG_PATH" "$IC" "$OVERRIDE" +done + +# same as above but use SST perturbed by +2K and +4K +for PERTURBATION in p2k p4k; do + case $PERTURBATION in + p2k) TEMPLATE_CONFIG=$AIMIP_INFERENCE_BASE_P2K_CONFIG_PATH ;; + p4k) TEMPLATE_CONFIG=$AIMIP_INFERENCE_BASE_P4K_CONFIG_PATH ;; + esac + for IC in {1..5}; do + JOB_NAME="${JOB_NAME_BASE}-${PERTURBATION}-IC${IC}" + IC_PATH="/climate-default/2025-09-12-aimip-evaluation/aimip-evaluation-ics-v3/1978-09-30_IC$(( IC - 1 )).nc" # files are 0-indexed + OUTPUT_PATH="/climate-default/2025-11-25-ace-aimip-inference-results/${JOB_NAME}" + OVERRIDE="initial_condition.path=${IC_PATH} experiment_dir=${OUTPUT_PATH}" + echo "Launching job: $JOB_NAME with perturbation: $PERTURBATION and IC: $IC" + launch_job "$JOB_NAME" "$TEMPLATE_CONFIG" "$IC" "$OVERRIDE" + done +done diff --git a/configs/baselines/era5-aimip/run-ace-train.sh b/configs/baselines/era5-aimip/run-ace-train.sh new file mode 100755 index 000000000..c2dc5ae2d --- /dev/null +++ b/configs/baselines/era5-aimip/run-ace-train.sh @@ -0,0 +1,56 @@ +#!/bin/bash + +set -e + +JOB_NAME_BASE="ace-aimip-train" +JOB_GROUP="ace-aimip" +CONFIG_FILENAME="ace-train-config.yaml" +SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository +CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME +BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name') +REPO_ROOT=$(git rev-parse --show-toplevel) +N_GPUS=4 + +cd $REPO_ROOT # so config path is valid no matter where we are running this script + +python -m fme.ace.validate_config --config_type train $CONFIG_PATH + +launch_job () { + + JOB_NAME=$1 + shift + OVERRIDE="$@" + + gantry run \ + --name $JOB_NAME \ + --task-name $JOB_NAME \ + --description 'Run ACE2-ERA5 training on AIMIP period' \ + --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \ + --workspace ai2/ace \ + --priority high \ + --preemptible \ + --cluster ai2/titan-cirrascale \ + --env WANDB_USERNAME=$BEAKER_USERNAME \ + --env WANDB_NAME=$JOB_NAME \ + --env WANDB_JOB_TYPE=training \ + --env WANDB_RUN_GROUP=$JOB_GROUP \ + --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \ + --env-secret WANDB_API_KEY=wandb-api-key-ai2cm-sa \ + --dataset-secret google-credentials:/tmp/google_application_credentials.json \ + --dataset oliverwm/era5-1deg-8layer-stats-1990-2019-v2:/statsdata \ + --gpus $N_GPUS \ + --shared-memory 400GiB \ + --weka climate-default:/climate-default \ + --budget ai2/climate \ + --system-python \ + --install "pip install --no-deps ." \ + -- torchrun --nproc_per_node $N_GPUS -m fme.ace.train $CONFIG_PATH --override $OVERRIDE + +} + +# random seed ensemble +for SEED in 0 1 2 3; do + JOB_NAME="${JOB_NAME_BASE}-rs${SEED}" + OVERRIDE="seed=${SEED}" + launch_job "$JOB_NAME" "$OVERRIDE" +done