Refactor/private generate method

configs/experiments/2026-02-10-downsc-add-pressfc/eval-100-to-3km-prmsl-output.yaml

@@ -0,0 +1,207 @@
+experiment_dir: /results
+n_samples: 2
+patch:
+  divide_generation: true
+  composite_prediction: true
+  coarse_horizontal_overlap: 1
+model:
+  checkpoint_path: /checkpoints/best_histogram_tail.ckpt #ema_ckpt.tar #best.ckpt
+  model_updates:
+   num_diffusion_generation_steps: 18
+   #sigma_max: 2000.0
+   #sigma_min: 0.002
+   #churn: 2.0
+data:
+  coarse:
+  - merge:
+    - data_path: /climate-default/2025-09-25-downscaling-data-X-SHiELD-AMIP-downscaling
+      file_pattern: 100km.zarr
+      engine: zarr
+      subset:
+        #start_time: '2023-01-01T00:00:00'
+        start: 0
+        stop: 2
+    - data_path:  /climate-default/2026-02-03-downscaling-data-X-SHiELD-AMIP-pressure-data
+      file_pattern: pressfc_renamed_to_prmsl_100km.zarr
+      engine: zarr
+      subset:
+        start: 0
+        stop: 2
+        #start_time: '2023-01-01T00:00:00'
+  fine:
+  - merge:
+    - data_path: /climate-default/2025-09-25-downscaling-data-X-SHiELD-AMIP-downscaling
+      file_pattern: 3km.zarr
+      engine: zarr
+      subset:
+        #start_time: '2023-01-01T00:00:00'
+        start: 0
+        stop: 2
+    - data_path:  /climate-default/2026-02-03-downscaling-data-X-SHiELD-AMIP-pressure-data
+      file_pattern: instantaneous_surface_and_sea_level_pressure_3km.zarr
+      engine: zarr
+      subset:
+        #start_time: '2023-01-01T00:00:00'
+        start: 0
+        stop: 2
+  # CONUS
+  lat_extent:
+    start: 0
+    stop: 16
+    # start: 22.0
+    # stop: 50.0
+  lon_extent:
+    #start: 230.0
+    #stop: 295.0
+    start: 0
+    stop: 16
+  batch_size: 2
+  num_data_workers: 2
+  strict_ensemble: false
+logging:
+  log_to_screen: true
+  log_to_wandb: true
+  log_to_file: true
+  project: multivariate-downscaling
+  entity: ai2cm
+events:
+- name: NE_US_Quebec_20230206
+  date: 2023-02-06T00:00
+  lat_extent:
+    start: 36
+    stop: 52
+  lon_extent:
+    start: 283
+    stop: 299
+  save_generated_samples: true
+  n_samples: 16
+- name: WA_AR_20230101
+  date: 2023-01-01T06:00
+  lat_extent:
+    start: 36
+    stop: 52
+  lon_extent:
+    start: 228
+    stop: 244
+  save_generated_samples: true
+  n_samples: 16
+- name: WPacific_hurricane_20230425
+  date: 2023-04-25T18:00
+  lat_extent:
+    start: 7
+    stop: 23
+  lon_extent:
+    start: 130.0
+    stop: 146.0
+  save_generated_samples: true
+  n_samples: 16
+- name: WPacific_hurricane_landfall_china_20230510
+  date: 2023-05-10T12:00
+  lat_extent:
+    start: 7
+    stop: 23
+  lon_extent:
+    start: 104
+    stop: 120
+  save_generated_samples: true
+  n_samples: 16
+- name: extratropical_cyclone_US_20230403
+  date: 2023-04-03T12:00
+  lat_extent:
+    start: 34
+    stop: 50
+  lon_extent:
+    start: 254
+    stop: 270
+  save_generated_samples: true
+  n_samples: 16
+- name: santa_ana_winds_20231221
+  date: 2023-12-21T06:00
+  lat_extent:
+    start: 26
+    stop: 42
+  lon_extent:
+    start: 234
+    stop: 250
+  save_generated_samples: true
+  n_samples: 16
+- name: alpine_foehn_20230330
+  date: 2023-03-30T18:00
+  lat_extent:
+    start: 37
+    stop: 53
+  lon_extent:
+    start: 2
+    stop: 18
+  save_generated_samples: true
+  n_samples: 16
+- name: hindu_kush_20230122
+  date: 2023-01-22T06:00
+  lat_extent:
+    start: 28
+    stop: 44
+  lon_extent:
+    start: 60
+    stop: 76
+  save_generated_samples: true
+  n_samples: 16
+- name: WPac_tc_20230426T06
+  date: 2023-04-26T06:00
+  lat_extent:
+    start: 8
+    stop: 24
+  lon_extent:
+    start: 130
+    stop: 146
+  save_generated_samples: true
+  n_samples: 16
+- name: Phl_tc_landfall_20230514T06
+  date: 2023-05-14T06:00
+  lat_extent:
+    start: 4
+    stop: 20
+  lon_extent:
+    start: 117
+    stop: 133
+  save_generated_samples: true
+  n_samples: 16
+- name: Phl_tc_landfall_20230517T18
+  date: 2023-05-17T18:00
+  lat_extent:
+    start: 7
+    stop: 23
+  lon_extent:
+    start: 133
+    stop: 149
+  save_generated_samples: true
+  n_samples: 16
+- name: Taiwan_tc_landfall_20230707T18
+  date: 2023-07-07T18:00
+  lat_extent:
+    start: 14
+    stop: 30
+  lon_extent:
+    start: 115
+    stop: 131
+  save_generated_samples: true
+  n_samples: 16
+- name: Japan_tc_landfall_20230919T18
+  date: 2023-09-19T18:00
+  lat_extent:
+    start: 22
+    stop: 38
+  lon_extent:
+    start: 123
+    stop: 139
+  save_generated_samples: true
+  n_samples: 16
+- name: Phl_tc_landfall_20231027T00
+  date: 2023-10-27T00:00
+  lat_extent:
+    start: 8
+    stop: 24
+  lon_extent:
+    start: 115
+    stop: 131
+  save_generated_samples: true
+  n_samples: 16

configs/experiments/2026-02-10-downsc-add-pressfc/eval-global.yaml

@@ -0,0 +1,50 @@
+experiment_dir: /results
+n_samples: 2
+patch:
+  divide_generation: true
+  composite_prediction: true
+  coarse_horizontal_overlap: 1
+model:
+  checkpoint_path: /checkpoints/best.ckpt
+data:
+  #topography: /climate-default/2025-09-25-downscaling-data-X-SHiELD-AMIP-downscaling/3km.zarr
+  fine:
+  - merge:
+    - data_path: /climate-default/2025-09-25-downscaling-data-X-SHiELD-AMIP-downscaling
+      file_pattern: 3km.zarr
+      engine: zarr
+      subset:
+        start_time: '2023-01-01T00:00:00'
+        step: 29
+    - data_path:  /climate-default/2026-02-03-downscaling-data-X-SHiELD-AMIP-pressure-data
+      file_pattern: instantaneous_surface_and_sea_level_pressure_3km.zarr
+      engine: zarr
+      subset:
+        start_time: '2023-01-01T00:00:00'
+        step: 29
+  coarse:
+  - merge:
+    - data_path: /climate-default/2025-09-25-downscaling-data-X-SHiELD-AMIP-downscaling
+      file_pattern: 100km.zarr
+      engine: zarr
+      subset:
+        start_time: '2023-01-01T00:00:00'
+        step: 29
+    - data_path:  /climate-default/2026-02-03-downscaling-data-X-SHiELD-AMIP-pressure-data
+      file_pattern: pressfc_renamed_to_prmsl_100km.zarr
+      engine: zarr
+      subset:
+        start_time: '2023-01-01T00:00:00'
+        step: 29
+  lat_extent:
+    start: -66
+    stop: 70.0
+  batch_size: 6
+  num_data_workers: 2
+  strict_ensemble: false
+logging:
+  log_to_screen: true
+  log_to_wandb: true
+  log_to_file: true
+  project: andrep-downscaling
+  entity: ai2cm

configs/experiments/2026-02-10-downsc-add-pressfc/eval.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+
+set -e
+
+#JOB_NAME="eval-xshield-amip-100km-to-3km-0.5sigmaexp-tropics-events"
+JOB_NAME="eval-xshield-amip-100km-to-3km-loguni-multivariate-global"
+
+CONFIG_FILENAME="eval-global.yaml"
+
+SCRIPT_PATH=$(echo "$(git rev-parse --show-prefix)" | sed 's:/*$::')
+CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME
+
+ # since we use a service account API key for wandb, we use the beaker username to set the wandb username
+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
+
+N_NODES=1
+NGPU=4
+
+IMAGE="$(cat latest_deps_only_image.txt)"
+
+EXISTING_RESULTS_DATASET=01KK07E72Z2H4CSP1EWXCT25FB
+wandb_group=""
+
+#--not-preemptible \
+
+gantry run \
+    --name $JOB_NAME \
+    --description 'Run 100km to 3km evaluation on coarsened X-SHiELD' \
+    --workspace ai2/climate-titan \
+    --priority urgent \
+    --cluster ai2/jupiter \
+    --cluster ai2/titan \
+    --beaker-image $IMAGE \
+    --env WANDB_USERNAME=$BEAKER_USERNAME \
+    --env WANDB_NAME=$JOB_NAME \
+    --env WANDB_JOB_TYPE=inference \
+    --env WANDB_RUN_GROUP=$wandb_group \
+    --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \
+    --env-secret WANDB_API_KEY=wandb-api-key-annak \
+    --dataset-secret google-credentials:/tmp/google_application_credentials.json \
+    --dataset $EXISTING_RESULTS_DATASET:checkpoints:/checkpoints \
+    --weka climate-default:/climate-default \
+    --gpus $NGPU \
+    --shared-memory 400GiB \
+    --budget ai2/climate \
+    --no-conda \
+    --install "pip install --no-deps ." \
+    --allow-dirty \
+    -- torchrun --nproc_per_node $NGPU -m fme.downscaling.evaluator $CONFIG_PATH

configs/experiments/2026-02-10-downsc-add-pressfc/resume.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+# uses the augusta cluster which doesn't have weka access but has GCS access and is
+# typically more available than cirrascale clusters
+
+set -e
+
+# recommended but not required to change this
+
+JOB_NAME="xshield-downscaling-100km-to-3km-0weight-prate-tropics-resume"
+CONFIG_FILENAME="train-100-to-3km-prmsl-clamp-loss-weight.yaml"
+
+SCRIPT_PATH=$(echo "$(git rev-parse --show-prefix)" | sed 's:/*$::')
+CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME
+wandb_group=""
+
+ # since we use a service account API key for wandb, we use the beaker username to set the wandb username
+BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name')
+REPO_ROOT=$(git rev-parse --show-toplevel)
+N_GPUS=8 # TODO: change to 8 after testing
+
+cd $REPO_ROOT  # so config path is valid no matter where we are running this script
+
+IMAGE=$(cat $REPO_ROOT/latest_deps_only_image.txt)
+
+PREVIOUS_RESULTS_DATASET="01KK9ZRC9M9MF7T45T4XAP7GF0"
+
+
+gantry run \
+    --name $JOB_NAME \
+    --description 'Run downscaling 100km to 3km multivar training' \
+    --workspace ai2/climate-titan \
+    --priority low \
+    --preemptible \
+    --cluster ai2/titan \
+    --cluster ai2/jupiter \
+    --beaker-image $IMAGE \
+    --env WANDB_USERNAME=$BEAKER_USERNAME \
+    --env WANDB_NAME=$JOB_NAME \
+    --env WANDB_JOB_TYPE=training \
+    --env WANDB_RUN_GROUP=$wandb_group \
+    --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \
+    --env-secret WANDB_API_KEY=wandb-api-key-annak \
+    --dataset $PREVIOUS_RESULTS_DATASET:/previous_results \
+    --dataset-secret google-credentials:/tmp/google_application_credentials.json \
+    --weka climate-default:/climate-default \
+    --gpus $N_GPUS \
+    --shared-memory 400GiB \
+    --budget ai2/climate \
+    --no-conda \
+    --install "pip install --no-deps ." \
+    --allow-dirty \
+    -- torchrun --nproc_per_node $N_GPUS -m fme.downscaling.train $CONFIG_PATH