======================================================
EAGLE-AI TUTORIAL COMMAND CHEAT SHEET
======================================================

These commands are designed to be run within a Google Colab notebook environment. Running these commands on a local machine or HPC system may require modifications to paths, environment setup, and dependency management.

# ------------------------------------------------------
# 1. Set up your environment
# ------------------------------------------------------

# Install required Python packages
!pip install anemoi-datasets==0.5.28 anemoi-graphs==0.8.2 anemoi-models==0.11.2 anemoi-training==0.8.3 anemoi-inference==0.8.3 anemoi-utils==0.4.40 anemoi-transform==0.1.19 trimesh ufs2arco==0.18.0 eagle-tools==0.6.2 mpi4py "torch<2.7" torchvision

# Clone tutorial repository
!git clone https://github.com/NOAA-EPIC/AMS-2026-short-course.git


# ------------------------------------------------------
# 2. Prepare training data
# ------------------------------------------------------

# Run data preparation using ufs2arco
!ufs2arco /content/AMS-2026-short-course/config/data/replay.yaml

# Inspect generated dataset
import xarray as xr

ufs2arco_ds = xr.open_dataset("/content/AMS-2026-short-course/config/data/replay.zarr")
ufs2arco_ds


# ------------------------------------------------------
# 3. Train the AI model
# ------------------------------------------------------

# Set environment variables for reproducibility and runtime
import os
os.environ["ANEMOI_BASE_SEED"] = "42"
os.environ["SLURM_JOB_ID"] = "0"

# Navigate to training configuration directory
%cd /content/AMS-2026-short-course/config/train/

# Run training
!anemoi-training train --config-name=config

# If running on CPU, use reduced steps
#!anemoi-training train --config-name=config ++training.max_steps=2


# ------------------------------------------------------
# 4. Generate a forecast
# ------------------------------------------------------

# Update inference config with latest checkpoint
import os

checkpoint_folder = "/content/AMS-2026-short-course/config/train/training-output/checkpoint/"
files = os.listdir(checkpoint_folder)
run_id = files[0]

inference_yaml_path = "/content/AMS-2026-short-course/config/inference/inference_config.yaml"
with open(inference_yaml_path, "r") as f:
    lines = f.readlines()

with open(inference_yaml_path, "w") as f:
    for line in lines:
        if line.strip().startswith("checkpoint_path"):
            f.write(f"checkpoint_path: ../train/training-output/checkpoint/{run_id}/inference-last.ckpt")
        else:
            f.write(line)

# Navigate to inference directory
%cd /content/AMS-2026-short-course/config/inference/

# Run inference
!eagle-tools inference inference_config.yaml


# ------------------------------------------------------
# 5. Post-process the output
# ------------------------------------------------------

# Navigate to postprocess directory
%cd /content/AMS-2026-short-course/config/postprocess

# Run postprocessing for verification compatibility
!eagle-tools prewxv postprocess.yaml

# Load postprocessed dataset
import xarray as xr
ds_post = xr.open_dataset("/content/AMS-2026-short-course/config/postprocess/postprocessed_files/global.2022-01-05T00.48h.nc")
ds_post


# ------------------------------------------------------
# 6. Verify and evaluate the model
# ------------------------------------------------------

# Install conda in Colab
!pip install -q condacolab
import condacolab
condacolab.install()

# Create wxvx environment
!conda create -y -n wxvx -c ufs-community -c paul.madden wxvx python=3.13

# Install time utility
!apt install -y time

# Navigate to verification directory
%cd /content/AMS-2026-short-course/config/verification

# Set backend for plotting
import os
os.environ["MPLBACKEND"] = "agg"

# Run verification (wxvx)
!conda run -n wxvx wxvx -c wxvx_config.yaml -t plots


# ------------------------------------------------------
# Output
# ------------------------------------------------------
# View plots here:
# AMS-2026-short-course/config/verification/run/plots/20220105/00