From 78b3a8c0fa9c1ec433d649412347ef20af3cda6c Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Thu, 13 Aug 2026 14:50:25 +0200 Subject: [PATCH 1/9] 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/9] 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/9] 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 From 82630027c5b5d3dbf88fb8eeb10d135228efa6be Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 12:54:06 +0200 Subject: [PATCH 4/9] Apply suggestions from code review Co-authored-by: SarahAlidoost <55081872+SarahAlidoost@users.noreply.github.com> --- scripts/README.md | 8 ++++---- scripts/run_best_tuned_model.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/scripts/README.md b/scripts/README.md index 8bdfeeb..739642d 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -2,17 +2,17 @@ ## 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. +- `data_preparation.*`: Scripts for preparing the data for training, tuning and evaluation and saving them to Zarr storage with specific chunking strategies. - `tuning.*`: Scripts for hyperparameter tuning. - `run_best_tuned_model.*`: Scripts for running the best tuned model on the test set. ## Experiments -### Tuning experiments +### Tuning experiments for SST variable - datasplit: train set = 2020, validation set = 2021, test set = 2022 -- path of tuning results: `/work//eso4clima/tune/`. -- test loss: 0.036662004509047774 +- path of tuning results: `/work//eso4clima/tune/sst_01`. +- test loss: 0.036662004509047774 (K) - hyperparameters of the best model: ``` {'patch_size': 8, diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index 7c2964d..e939eea 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -151,7 +151,7 @@ def main() -> None: prediction_config = PredictionConfig( calculate_residuals=True, - return_numpy=True, + return_numpy=False, save_predictions=False, return_loss=True, device="cpu", From 9bf05f39c11f91b49f94b0403cf51cef36a9f115 Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 13:37:27 +0200 Subject: [PATCH 5/9] add getting bets hyperparameters --- scripts/run_best_tuned_model.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index e939eea..4cf0fef 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -145,6 +145,12 @@ def main() -> None: analysis = tune.ExperimentAnalysis(str(experiment_path)) best_result = analysis.get_best_trial("loss", "min") + + # Get the best hyperparameters + best_hyperparameters = best_result.get_best_config(metric="loss", mode="min") + print(f"Best config: {best_hyperparameters}") + + # Get the best model best_checkpoint = best_result.checkpoint model_path = Path(best_checkpoint.path) / "checkpoint.pt" print(f"Best checkpoint path: {model_path}") From 02917c5c1a8cdfecfb5d233e2614f298054bf100 Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 16:26:32 +0200 Subject: [PATCH 6/9] change variable name --- scripts/run_best_tuned_model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index 4cf0fef..3c13177 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -72,14 +72,14 @@ def main() -> None: if not lsm_file_path.exists(): raise FileNotFoundError(f"LSM file does not exist: {lsm_file_path}") - daily_files = list( + input_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: + if not input_files: raise FileNotFoundError( f"No daily test files found for year '{args.year}' in '{test_data_dir}'" ) @@ -88,11 +88,11 @@ def main() -> None: 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 daily files ({len(input_files)}): {input_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 + input_files, combine="by_coords", parallel=False ) monthly_data_test = xr.open_mfdataset( monthly_files, combine="by_coords", parallel=False From 02b3e51b032d4feeee7491fe7bd7b7253bb6576d Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 16:29:13 +0200 Subject: [PATCH 7/9] use cuda --- scripts/run_best_tuned_model.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index 3c13177..86e6a0b 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -160,7 +160,7 @@ def main() -> None: return_numpy=False, save_predictions=False, return_loss=True, - device="cpu", + device="cuda", verbose=False, ) @@ -170,7 +170,7 @@ def main() -> None: num_workers=0, pin_memory=False, persistent_workers=False, - device="cpu", + device="cuda", multiprocessing_context=None, ) @@ -182,8 +182,7 @@ def main() -> None: run_dir=run_dir, ) - print("Test loss:") - print(test_loss) + print("Test loss:", test_loss) if __name__ == "__main__": From 4f20ce70a56a5a262aa3de4a4d8aca9e8b799db9 Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 16:32:00 +0200 Subject: [PATCH 8/9] update configs in script --- scripts/run_best_tuned_model.py | 6 +++--- scripts/run_best_tuned_model.slurm | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index 86e6a0b..c199ccf 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -120,12 +120,12 @@ def main() -> None: lsm_mask = xr.open_dataset(lsm_file_path) num_patches = (10, 10) - patch_size = (1, 4, 4) + patch_size = (1, 40, 40) 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) + stride = (20, 20) dataset_test = STDataset( input_da=input_da, @@ -148,7 +148,7 @@ def main() -> None: # Get the best hyperparameters best_hyperparameters = best_result.get_best_config(metric="loss", mode="min") - print(f"Best config: {best_hyperparameters}") + print(f"Best hyperparameters: {best_hyperparameters}") # Get the best model best_checkpoint = best_result.checkpoint diff --git a/scripts/run_best_tuned_model.slurm b/scripts/run_best_tuned_model.slurm index aa3f3b4..95279fb 100644 --- a/scripts/run_best_tuned_model.slurm +++ b/scripts/run_best_tuned_model.slurm @@ -5,7 +5,7 @@ #SBATCH --cpus-per-task=128 #SBATCH --time=02:00:00 #SBATCH --account=bd0854 -#SBATCH --partition=compute +#SBATCH --partition=gpu #SBATCH --output=climanet_eval_%j.out #SBATCH --error=climanet_eval_%j.err From a843779a9ca2ac8e0bfbfc9f2d001301caa2f13d Mon Sep 17 00:00:00 2001 From: Ou Ku Date: Wed, 26 Aug 2026 16:37:26 +0200 Subject: [PATCH 9/9] reove the data preparation part --- scripts/run_best_tuned_model.py | 25 +++---------------------- 1 file changed, 3 insertions(+), 22 deletions(-) diff --git a/scripts/run_best_tuned_model.py b/scripts/run_best_tuned_model.py index c199ccf..a162e6f 100644 --- a/scripts/run_best_tuned_model.py +++ b/scripts/run_best_tuned_model.py @@ -6,7 +6,7 @@ from climanet.dataset import DataLoaderConfig, STDataset from climanet.predict import PredictionConfig, predict_monthly_var -from climanet.utils import data_preparation, read_st_data +from climanet.utils import read_st_data def build_parser() -> argparse.ArgumentParser: @@ -26,7 +26,7 @@ def build_parser() -> argparse.ArgumentParser: "--test-data-dir", type=Path, required=True, - help="Directory containing the test NetCDF files.", + help="Directory containing the test Zarr files.", ) parser.add_argument( "--lsm-file-path", @@ -91,28 +91,9 @@ def main() -> None: print(f"Using daily files ({len(input_files)}): {input_files[:3]} ...") print(f"Using monthly files ({len(monthly_files)}): {monthly_files[:3]} ...") - daily_data_test = xr.open_mfdataset( - input_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, + data_path=test_data_dir, var_name=args.var_name, ) )