forecast_local.sh

#!/bin/bash

# Initialize conda:
export OMP_NUM_THREADS=4

# Config options - should remain fixed for now
MODEL="standard_Aurora"
DATASET="cerra_with_boundaries"

# Paths - change these to your paths
# The checkpoint_path currently points to the rollout trained model with boundaries
# The output_dir should be changed to your desired output path
DATASET_PATH="/mnt/ssd/datasets/cerra_2022.zarr"
BOUNDARY_DATASET_PATH="/mnt/ssd/datasets/hres_forecasts_2022_projected.zarr/"
CHECKPOINT_PATH="/mnt/ssd/trained_aurora_models_on_cerra/Aurora_large_6h_rollout_boundaries_final/aurora-Aurora-rollout_long-step15000-best.ckpt"
OUTPUT_DIR="/mnt/ssd/datasets/forecast_aurora_6h_rollout_long_w_boundaries_final.zarr/"

rm -rf "$OUTPUT_DIR"

# There are three models available: The 6-hour base model, the rollout-trained model, and the quantized rollout-trained model.
# The default settings in this script use the rollout model.
# To use the base model
# 1. Set CHECKPOINT_PATH to the base model checkpoint
# 2. Change DATASET to "cerra"
# 3. Set task.max_rollout_steps=1
# 4. Set task.lead_times="[6]"
# 5. Remove the dataset.common.boundary_path and dataset.common.boundary_size lines
#
#
# To use the quantized model
# 1. Set CHECKPOINT_PATH to the quantized model checkpoint
# 2. Set task=forecast_quantized instead of task=forecast

# task.use_wb2_format is used to control whether to save in the normal CERRA format or in
# a format compatible with WeatherBench2 (WB2). When we plan to evaluate using our own
# compute_forecast_rmse.py script, we set this to False. When we plan to evaluate using
# the WB2 evaluation scripts, we set this to True.
# Comment the dataset.common.boundary_path line out if not using boundaries.
# To use a HuggingFace checkpoint instead of a local one, set task.load_from_hf to True,
# and adjust task.hf_repo and task.hf_checkpoint accordingly
python main.py --config-name forecast \
    task=forecast \
    task.load_from_hf=False \
    task.distributed=False \
    task.model_name=Aurora \
    task.checkpoint_path=$CHECKPOINT_PATH \
    task.output_dir=$OUTPUT_DIR \
    task.use_wb2_format=True \
    task.lead_times="[6,12,18,24,30]" \
    task.save_fp64=False \
    model=$MODEL \
    model.patch_size=8 \
    model.use_lora=False \
    dataset=$DATASET \
    dataset.common.lead_time_hours=6 \
    dataset.common.data_path=$DATASET_PATH \
    dataset.common.boundary_path=$BOUNDARY_DATASET_PATH \
    dataset.common.boundary_size="[168,171,168,171]" \
    task.max_rollout_steps=5 \
    dataset.val.start_time="2022-01-01T00:00:00" \
    dataset.val.end_time="2022-12-31T21:00:00" \
    dataloader.num_workers=4