From 78b3a8c0fa9c1ec433d649412347ef20af3cda6c Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Thu, 13 Aug 2026 14:50:25 +0200 Subject: [PATCH 1/3] add scripts to run best tuned model on test dataset --- scripts/run_best_tuned_model.py | 186 +++++++++++++++++++++++++++++ scripts/run_best_tuned_model.slurm | 24 ++++ 2 files changed, 210 insertions(+) create mode 100644 scripts/run_best_tuned_model.py create mode 100644 scripts/run_best_tuned_model.slurm diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py new file mode 100644 index 0000000..8393b54 --- /dev/null +++ b/scripts/run_best_tuned_model.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 + +import argparse +from pathlib import Path + +import xarray as xr +from ray import tune + +from climanet.dataset import DataLoaderConfig, STDataset +from climanet.predict import PredictionConfig, predict_monthly_var +from climanet.utils import data_preparation, read_st_data + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Load the best Ray Tune checkpoint, prepare the test data, and evaluate the " + "trained model on the 2023 test period." + ) + ) + parser.add_argument( + "--experiment-path", + type=Path, + required=True, + help="Path to the Ray Tune experiment directory containing the checkpoint.", + ) + parser.add_argument( + "--test-data-dir", + type=Path, + required=True, + help="Directory containing the test NetCDF files.", + ) + parser.add_argument( + "--lsm-file-path", + type=Path, + required=True, + help="Path to the land-sea mask NetCDF file.", + ) + parser.add_argument( + "--run-dir", + type=Path, + default=Path("./run_dir_tune_test").resolve(), + help="Directory used for the evaluation run and saved logs.", + ) + parser.add_argument( + "--var-name", + type=str, + default="tos", + help="Variable name to evaluate in the NetCDF files.", + ) + parser.add_argument( + "--year", + type=str, + default="2022", + help="Year pattern to include in the test files (e.g. 2022).", + ) + return parser + + +def main() -> None: + args = build_parser().parse_args() + experiment_path = args.experiment_path.resolve() + test_data_dir = args.test_data_dir.resolve() + lsm_file_path = args.lsm_file_path.resolve() + run_dir = args.run_dir.resolve() + run_dir.mkdir(parents=True, exist_ok=True) + + if not experiment_path.exists(): + raise FileNotFoundError( + f"Experiment directory does not exist: {experiment_path}" + ) + if not test_data_dir.exists(): + raise FileNotFoundError(f"Test data directory does not exist: {test_data_dir}") + if not lsm_file_path.exists(): + raise FileNotFoundError(f"LSM file does not exist: {lsm_file_path}") + + daily_files = list( + test_data_dir.glob(f"{args.year}*_hr_ERA5dc_masked_{args.var_name}*.nc") + ) + monthly_files = list( + test_data_dir.glob(f"{args.year}*_mon_ERA5dc_masked_{args.var_name}*.nc") + ) + + if not daily_files: + raise FileNotFoundError( + f"No daily test files found for year '{args.year}' in '{test_data_dir}'" + ) + if not monthly_files: + raise FileNotFoundError( + f"No monthly test files found for year '{args.year}' in '{test_data_dir}'" + ) + + print(f"Using daily files ({len(daily_files)}): {daily_files[:3]} ...") + print(f"Using monthly files ({len(monthly_files)}): {monthly_files[:3]} ...") + + daily_data_test = xr.open_mfdataset( + daily_files, combine="by_coords", parallel=False + ) + monthly_data_test = xr.open_mfdataset( + monthly_files, combine="by_coords", parallel=False + ) + + test_data_zarr_dir = run_dir / "test_data_zarr" + test_data_zarr_dir.mkdir(parents=True, exist_ok=True) + + _ = data_preparation( + daily_data_test[args.var_name], + monthly_data_test[args.var_name], + calculate_residuals=True, + is_hourly=True, + save_to_zarr=True, + run_dir=test_data_zarr_dir, + ) + + input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = ( + read_st_data( + data_path=test_data_zarr_dir, + var_name=args.var_name, + ) + ) + + lsm_mask = xr.open_dataset(lsm_file_path) + + num_patches = (10, 10) + patch_size = (1, 4, 4) + spatial_patch_size = ( + patch_size[1] * num_patches[0], + patch_size[2] * num_patches[1], + ) + stride = (spatial_patch_size[0] // 5, spatial_patch_size[1] // 5) + + dataset_test = STDataset( + input_da=input_da, + input_da_nan_mask=input_da_nan_mask, + monthly_da=monthly_da, + padded_days_mask=padded_days_mask, + time_features=time_features, + land_mask=lsm_mask["lsm"], + patch_size=(1, *spatial_patch_size), + stride=stride, + sh_embed_dim=96, + sh_order_L=10, + verbose=True, + load_lazy=False, + ) + print(f"Created test dataset with {len(dataset_test)} patches.") + + analysis = tune.ExperimentAnalysis(str(experiment_path)) + best_result = analysis.get_best_trial("loss", "min") + best_checkpoint = best_result.checkpoint + model_path = Path(best_checkpoint.path) / "checkpoint.pt" + print(f"Best checkpoint path: {model_path}") + + prediction_config = PredictionConfig( + calculate_residuals=True, + return_numpy=True, + save_predictions=False, + return_loss=True, + device="cpu", + verbose=False, + ) + + dataloader_config = DataLoaderConfig( + batch_size=10, + shuffle=True, + num_workers=0, + pin_memory=False, + persistent_workers=False, + device="cpu", + multiprocessing_context=None, + ) + + test_loss = predict_monthly_var( + model=model_path, + dataset=dataset_test, + dataloader_config=dataloader_config, + prediction_config=prediction_config, + run_dir=run_dir, + ) + + print("Test loss:") + print(test_loss) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_best_tuned_model.slurm b/scripts/run_best_tuned_model.slurm new file mode 100644 index 0000000..aa3f3b4 --- /dev/null +++ b/scripts/run_best_tuned_model.slurm @@ -0,0 +1,24 @@ +#!/bin/bash +#SBATCH --job-name=climanet_eval +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=128 +#SBATCH --time=02:00:00 +#SBATCH --account=bd0854 +#SBATCH --partition=compute +#SBATCH --output=climanet_eval_%j.out +#SBATCH --error=climanet_eval_%j.err + +set -euo pipefail + +source /home/b/b383704/eso4clima/ClimaNet/.venv/bin/activate + +python -u /home/b/b383704/eso4clima/run_best_tuned_model/run_best_tuned_model.py \ + --experiment-path /work/bd0854/eso4clima/tune/sst_01 \ + --test-data-dir /work/bd0854/b380103/eso4clima/output/sst/concatenated/ \ + --lsm-file-path /home/b/b383704/eso4clima/data/era5_lsm_bool.nc \ + --run-dir /home/b/b383704/eso4clima/run_best_tuned_model/run_dir \ + --var-name tos \ + --year 2022 + +printf "\nFinished evaluation run.\n" From 66560b08fb2ce9efba7892f4e60854dfe6d8db3a Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Thu, 13 Aug 2026 15:23:54 +0200 Subject: [PATCH 2/3] doc best hypterparameters --- scripts/README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 scripts/README.md diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000..8bdfeeb --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,27 @@ +# Scripts + +## Structure + +- `data_preparation.*`: Scripts for preparing the data for training, tuning and evaluation. Mainly converting large netCDF files to Zarr storage with specific chunking strategies. This allows executing training for larger-than-memory datasets. +- `tuning.*`: Scripts for hyperparameter tuning. +- `run_best_tuned_model.*`: Scripts for running the best tuned model on the test set. + +## Experiments + +### Tuning experiments + +- datasplit: train set = 2020, validation set = 2021, test set = 2022 +- path of tuning results: `/work//eso4clima/tune/`. +- test loss: 0.036662004509047774 +- hyperparameters of the best model: + ``` + {'patch_size': 8, + 'overlap': 1, + 'embed_dim': 64, + 'dropout': 0.2, + 'hidden': 32, + 'spatial_depth': 3, + 'spatial_heads': 2, + 'optimizer_lr': 0.001787422899066508, + 'batch_config': {'batch_size': 100, 'accumulation_steps': 2}} + ``` From 1e93be51c12b68dba4563b70ed5f186605c29d89 Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Fri, 14 Aug 2026 13:02:40 +0200 Subject: [PATCH 3/3] remove shebang --- scripts/run_best_tuned_model.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index 8393b54..7c2964d 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python3 - import argparse from pathlib import Path