From e4bebd7a52e11e0fd4861cf48e8a4392a68062c0 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 23 Sep 2026 16:07:51 +0200 Subject: [PATCH 01/17] check nan of pred in compute_masked_loss in utils --- climanet/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/climanet/utils.py b/climanet/utils.py index 3278665..875c24d 100644 --- a/climanet/utils.py +++ b/climanet/utils.py @@ -271,12 +271,12 @@ def compute_masked_loss( """Compute L1 loss masked to ocean pixels only.""" ocean = (~land_mask).to(pred.device).unsqueeze(1) - # Mask for valid (non-NaN) target values - valid = ~torch.isnan(target) + # Mask for valid (non-NaN) values + valid = ~torch.isnan(target) & ~torch.isnan(pred) target = torch.nan_to_num(target, nan=0.0) - - mask = ocean & valid + pred = torch.nan_to_num(pred, nan=0.0) loss = torch.nn.functional.l1_loss(pred, target, reduction="none") + mask = ocean & valid loss = loss * mask num = loss.sum(dim=(-2, -1)) From 06643dc493c9e33c9d54a8ff98e3670217920ff9 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 23 Sep 2026 16:08:38 +0200 Subject: [PATCH 02/17] change model to predict residuals --- climanet/st_encoder_decoder.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index 32f76c4..256bea4 100644 --- a/climanet/st_encoder_decoder.py +++ b/climanet/st_encoder_decoder.py @@ -742,9 +742,24 @@ def forward( # decoder input shape is (B, M*Hp*Wp, C), C: embedding dimension # decoder output shape is (B, M, H, W) if self.use_checkpoint: - monthly_pred = checkpoint( + monthly_res_pred = checkpoint( self.decoder, x, M, H, W, land_mask_patch, use_reentrant=False ) # (B, M, H, W) else: - monthly_pred = self.decoder(x, M, H, W, land_mask_patch) # (B, M, H, W) + monthly_res_pred = self.decoder(x, M, H, W, land_mask_patch) # (B, M, H, W) + + # Add residuals to monthly mean + masked_data = input_data.masked_fill(daily_mask, float("nan")) + mean_data = torch.nanmean(masked_data, dim=3) + all_missing = daily_mask.all(dim=3) + mean_data = mean_data.masked_fill(all_missing, float("nan")) + + # Apply land mask + mean_data = mean_data.masked_fill( + land_mask_patch[:, None, None, :, :], + float("nan"), + ) + mean_data = mean_data.squeeze(1) # (B, M, H, W) + monthly_pred = monthly_res_pred + mean_data + return monthly_pred From 7f02ef964a51391efe5d699b18618f33cf76c6bc Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 23 Sep 2026 16:09:59 +0200 Subject: [PATCH 03/17] improve decoder, add residual_scale to decoder --- climanet/st_encoder_decoder.py | 76 ++++++++++++++++++++++++---------- 1 file changed, 55 insertions(+), 21 deletions(-) diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index 256bea4..ab19573 100644 --- a/climanet/st_encoder_decoder.py +++ b/climanet/st_encoder_decoder.py @@ -6,6 +6,7 @@ import torch from torch import nn from torch.utils.checkpoint import checkpoint +import math class VideoEncoder(nn.Module): @@ -41,6 +42,7 @@ def __init__(self, in_chans=1, embed_dim=128, patch_size=(1, 4, 4)): 2 * in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) + def forward(self, x, mask): """Forward pass with masking support via an additional validity channel. Args: @@ -66,7 +68,6 @@ def forward(self, x, mask): x = x.contiguous(memory_format=torch.channels_last_3d) x = self.proj(x) x = x.flatten(2).transpose(1, 2) - return x @@ -381,6 +382,7 @@ def __init__( patch_w=4, hidden=128, dropout=0.0, + residual_scale=0.5, ): """ Args: @@ -392,46 +394,67 @@ def __init__( hidden: Hidden dimension in the decoder for mixing channel features. The default is 128, which can be tuned. dropout: Dropout rate for regularization in the refinement block. Default is 0.0. + residual_scale: Scaling factor for the residual connection in the refinement block. Default is 0.5. """ super().__init__() self.patch_h = patch_h self.patch_w = patch_w + self.residual_scale = residual_scale + + # Number of progressive x2 upsampling stages. + n_upsample = int(math.log2(patch_h)) # Mix channel features on the patch grid (Hp, Wp) # Input shape: (B, embed_dim, Hp, Wp) → Output shape: (B, hidden, Hp, Wp) # here kernel_size=1 means we are mixing features at each patch location # without spatial interaction in_channels, out_channels = embed_dim, hidden - self.proj = nn.Conv2d(in_channels, out_channels, kernel_size=1) + self.proj = nn.Sequential( + nn.Conv2d(embed_dim, hidden, kernel_size=1), + nn.GroupNorm(8, hidden), + nn.GELU(), + ) - # Upsample to full resolution - in_channels, out_channels = hidden, hidden // 2 - self.deconv = nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=(patch_h, patch_w), - stride=(patch_h, patch_w), - padding=0, + # Spatial processing before upsampling + self.spatial_mix = nn.Sequential( + nn.Conv2d(hidden, hidden, kernel_size=3, padding=1), + nn.GroupNorm(8, hidden), + nn.GELU(), ) + # Upsample to full resolution + channels = hidden + self.deconv = nn.ModuleList() + + for i in range(n_upsample): + # Gradually reduce channels as spatial resolution increases + next_channels = max(hidden // (2 ** (i + 1)), 16) + block = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(channels, next_channels, kernel_size=3, padding=1), + nn.GroupNorm(min(8, next_channels), next_channels), + nn.GELU(), + nn.Dropout2d(dropout), + ) + self.deconv.append(block) + channels = next_channels + # Final conv head to get single channel output kernel_size=3 is the most # common choice for spatial convolutions; it's the smallest kernel that # captures spatial context in all directions - in_channels, out_channels = hidden // 2, hidden // 2 - # Refinement block: a small conv layers to smooth patch boundaries self.refine = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), - nn.GroupNorm(num_groups=8, num_channels=out_channels), - nn.GELU(), - nn.Dropout2d(dropout), - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), - nn.GroupNorm(num_groups=8, num_channels=out_channels), + nn.Conv2d(channels, channels, kernel_size=3, padding=1), + nn.GroupNorm(min(8, channels), channels), nn.GELU(), ) # Final conv head to map to single-channel output - self.head = nn.Conv2d(out_channels, 1, kernel_size=1) + self.head = nn.Conv2d(channels, 1, kernel_size=1) + + # Start from zero residual + nn.init.zeros_(self.head.weight) + nn.init.zeros_(self.head.bias) def forward(self, latent, M, out_H, out_W, land_mask=None): """Reconstruct 2D maps from latent patch tokens. @@ -459,8 +482,13 @@ def forward(self, latent, M, out_H, out_W, land_mask=None): # Apply 1x1 convolution to mix features out = self.proj(out) # (B*M, hidden, Hp, Wp) - # Use transposed convolution to upsample - out = self.deconv(out) # (B*M, hidden//2, H, W) + # Spatial interaction before upsampling + residual = self.spatial_mix(out) + out = out + residual + + # Progressive x2 upsampling + for block in self.deconv: + out = block(out) # Refinement CNN to smooth boundaries refine_input = out @@ -469,6 +497,10 @@ def forward(self, latent, M, out_H, out_W, land_mask=None): # Apply final conv head to get single channel output out = self.head(out) # (B*M, 1, H, W) + + # Apply residual scaling and non-linearity + out = self.residual_scale * torch.tanh(out) + out = out.view(B, M, out_H, out_W) # (B, M, H, W) # Mask out land areas if land_mask is provided @@ -589,6 +621,7 @@ def __init__( sh_dim=96, scale_dim=10, use_checkpoint=True, + decoder_residual_scale=0.5, ): """Initialize the Spatio-Temporal Model. @@ -633,6 +666,7 @@ def __init__( patch_w=patch_size[2], hidden=hidden, dropout=dropout, + residual_scale=decoder_residual_scale, ) self.patch_size = patch_size From 2c876b4adae75ca6a30faf8470598f04ca5c532c Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Tue, 29 Sep 2026 09:48:36 +0200 Subject: [PATCH 04/17] make dataloader shuffle false in predict --- climanet/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/climanet/predict.py b/climanet/predict.py index 09ecb26..e774a3a 100644 --- a/climanet/predict.py +++ b/climanet/predict.py @@ -132,7 +132,7 @@ def predict_monthly_var( dataloader = DataLoader( dataset, batch_size=dataloader_config.batch_size, - shuffle=dataloader_config.shuffle, + shuffle=False, # no shuffling during prediction/validation pin_memory=use_cuda, num_workers=dataloader_config.num_workers, # for data loading persistent_workers=dataloader_config.persistent_workers, # keep workers alive between epochs From f2cc65be341e82ba00544ec186856bc4ba63426d Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Tue, 29 Sep 2026 09:49:38 +0200 Subject: [PATCH 05/17] fix model copy in train --- climanet/train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/climanet/train.py b/climanet/train.py index 3dbd4a3..b9d3473 100644 --- a/climanet/train.py +++ b/climanet/train.py @@ -1,3 +1,4 @@ +import copy import tempfile from dataclasses import dataclass from pathlib import Path @@ -196,6 +197,7 @@ def train_monthly_model( ): gap = avg_val_loss - avg_train_loss print(f"Epoch {epoch}: gap between train and val loss: {gap:.6f}") + print(f"Epoch {epoch}: train loss = {avg_train_loss:.6f}, val loss = {avg_val_loss:.6f}") # Step scheduler scheduler.step(avg_epoch_loss) @@ -204,7 +206,7 @@ def train_monthly_model( # Consider improvement only if loss decreases more than a small threshold if avg_epoch_loss < best_loss - 1e-4: best_loss = avg_epoch_loss - best_state_dict = {k: v.detach() for k, v in model.state_dict().items()} + best_state_dict = copy.deepcopy(model.state_dict()) counter = 0 else: counter += 1 From df34ecac5bc06338e032fea264f64df99cfd486e Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 10:02:29 +0200 Subject: [PATCH 06/17] refactor train and predict --- climanet/predict.py | 26 ++------------------------ climanet/train.py | 25 ++----------------------- 2 files changed, 4 insertions(+), 47 deletions(-) diff --git a/climanet/predict.py b/climanet/predict.py index e774a3a..5dff1fc 100644 --- a/climanet/predict.py +++ b/climanet/predict.py @@ -8,7 +8,7 @@ from climanet.dataset import DataLoaderConfig from climanet.utils import ( - compute_masked_loss, + run_one_batch, load_model, setup_logging, ) @@ -74,28 +74,6 @@ def _save_netcdf( ds_pred.sel(time=[t]).to_netcdf(file_name) -def _move_batch_to_device(batch: dict, device: str): - use_cuda = device == "cuda" - return {k: v.to(device, non_blocking=use_cuda) for k, v in batch.items()} - - -def _run_one_batch(model: torch.nn.Module, batch: dict, device: str): - batch = _move_batch_to_device(batch, device) - pred = model( - batch["input_data"], - batch["input_data_mask"], - batch["input_data_timef"], - batch["land_mask"], - batch["geo_pos_embedding"], - batch["scale_feature"], - batch["padded_days_mask"], - ) # (B, M, H, W) - - # Compute masked loss - loss = compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"]) - return loss, pred - - def predict_monthly_var( model: torch.nn.Module | str, dataset: Dataset, @@ -158,7 +136,7 @@ def predict_monthly_var( idx = 0 average_loss = 0.0 for i, batch in enumerate(dataloader): - loss, predictions = _run_one_batch(model, batch, device) + loss, predictions = run_one_batch(model, batch, device) average_loss += loss.detach() all_predictions[idx : idx + predictions.size(0)] = predictions.detach() diff --git a/climanet/train.py b/climanet/train.py index b9d3473..8c15faa 100644 --- a/climanet/train.py +++ b/climanet/train.py @@ -11,7 +11,7 @@ from climanet.dataset import DataLoaderConfig from climanet.predict import PredictionConfig, predict_monthly_var from climanet.utils import ( - compute_masked_loss, + run_one_batch, save_model, setup_logging, ) @@ -37,27 +37,6 @@ class TrainConfig: store_logs: bool = True -def _move_batch_to_device(batch: dict, device: str): - use_cuda = device == "cuda" - return {k: v.to(device, non_blocking=use_cuda) for k, v in batch.items()} - - -def _run_one_batch(model: torch.nn.Module, batch: dict, device): - batch = _move_batch_to_device(batch, device) - pred = model( - batch["input_data"], - batch["input_data_mask"], - batch["input_data_timef"], - batch["land_mask"], - batch["geo_pos_embedding"], - batch["scale_feature"], - batch["padded_days_mask"], - ) # (B, M, H, W) - - # Compute masked loss - return compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"]) - - def _load_checkpoint(model, optimizer, loaded_checkpoint): with loaded_checkpoint.as_directory() as loaded_checkpoint_dir: loaded_checkpoint_dir = Path(loaded_checkpoint_dir).resolve() @@ -137,7 +116,7 @@ def train_monthly_model( optimizer.zero_grad() for i, batch in enumerate(dataloader): - loss = _run_one_batch(model, batch, device) + loss, _ = run_one_batch(model, batch, device, return_predictions=False) # Scale loss for gradient accumulation scaled_loss = loss * (1.0 / training_config.accumulation_steps) From dababfbc5390219df6f6c013e03a3169229b2746 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 10:03:12 +0200 Subject: [PATCH 07/17] replace l1_loss with mse to fix generalization --- climanet/utils.py | 52 ++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 49 insertions(+), 3 deletions(-) diff --git a/climanet/utils.py b/climanet/utils.py index 875c24d..22ac98a 100644 --- a/climanet/utils.py +++ b/climanet/utils.py @@ -265,6 +265,51 @@ def setup_logging(log_dir: str) -> SummaryWriter: return SummaryWriter(log_dir, filename_suffix=f"_UTC{timestamp_utc}") + +def move_batch_to_device(batch: dict, device: str): + use_cuda = device == "cuda" + return {k: v.to(device, non_blocking=use_cuda) for k, v in batch.items()} + + +def run_one_batch( + model: torch.nn.Module, batch: dict, device: str, return_predictions: bool = True + ): + batch = move_batch_to_device(batch, device) + monthly_res_pred = model( + batch["input_data"], + batch["input_data_mask"], + batch["input_data_timef"], + batch["land_mask"], + batch["geo_pos_embedding"], + batch["scale_feature"], + batch["padded_days_mask"], + ) # (B, M, H, W) + + # Add residuals to monthly mean + input_data = batch["input_data"] + daily_mask = batch["input_data_mask"] + land_mask_patch = batch["land_mask"] + masked_data = input_data.masked_fill(daily_mask, float("nan")) + mean_data = torch.nanmean(masked_data, dim=3) + all_missing = daily_mask.all(dim=3) + mean_data = mean_data.masked_fill(all_missing, float("nan")) + + # Apply land mask + mean_data = mean_data.masked_fill( + land_mask_patch[:, None, None, :, :], + float("nan"), + ) + mean_data = mean_data.squeeze(1) # (B, M, H, W) + + monthly_res = batch["monthly_data"] - mean_data + monthly_pred = monthly_res_pred + mean_data + + # Compute masked loss + loss = compute_masked_loss(monthly_res_pred, monthly_res, batch["land_mask"]) + + return (loss, monthly_pred) if return_predictions else (loss, None) + + def compute_masked_loss( pred: torch.Tensor, target: torch.Tensor, land_mask: torch.Tensor ) -> torch.Tensor: @@ -275,14 +320,15 @@ def compute_masked_loss( valid = ~torch.isnan(target) & ~torch.isnan(pred) target = torch.nan_to_num(target, nan=0.0) pred = torch.nan_to_num(pred, nan=0.0) - loss = torch.nn.functional.l1_loss(pred, target, reduction="none") + + loss = torch.nn.functional.mse_loss(pred, target, reduction="none") mask = ocean & valid loss = loss * mask - num = loss.sum(dim=(-2, -1)) denom = mask.sum(dim=(-2, -1)).clamp_min(1) + loss = (num / denom).mean() - return (num / denom).mean() + return loss def save_model( From f3185ba84662007d8c25862a503f85ca19217625 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 10:03:46 +0200 Subject: [PATCH 08/17] refactor the model --- climanet/st_encoder_decoder.py | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index ab19573..0f4e964 100644 --- a/climanet/st_encoder_decoder.py +++ b/climanet/st_encoder_decoder.py @@ -404,6 +404,7 @@ def __init__( # Number of progressive x2 upsampling stages. n_upsample = int(math.log2(patch_h)) + # Mix channel features on the patch grid (Hp, Wp) # Input shape: (B, embed_dim, Hp, Wp) → Output shape: (B, hidden, Hp, Wp) # here kernel_size=1 means we are mixing features at each patch location @@ -452,10 +453,6 @@ def __init__( # Final conv head to map to single-channel output self.head = nn.Conv2d(channels, 1, kernel_size=1) - # Start from zero residual - nn.init.zeros_(self.head.weight) - nn.init.zeros_(self.head.bias) - def forward(self, latent, M, out_H, out_W, land_mask=None): """Reconstruct 2D maps from latent patch tokens. Args: @@ -782,18 +779,4 @@ def forward( else: monthly_res_pred = self.decoder(x, M, H, W, land_mask_patch) # (B, M, H, W) - # Add residuals to monthly mean - masked_data = input_data.masked_fill(daily_mask, float("nan")) - mean_data = torch.nanmean(masked_data, dim=3) - all_missing = daily_mask.all(dim=3) - mean_data = mean_data.masked_fill(all_missing, float("nan")) - - # Apply land mask - mean_data = mean_data.masked_fill( - land_mask_patch[:, None, None, :, :], - float("nan"), - ) - mean_data = mean_data.squeeze(1) # (B, M, H, W) - monthly_pred = monthly_res_pred + mean_data - - return monthly_pred + return monthly_res_pred From 7a871a3a63a67a6655b7dcd1f258f63973d7580b Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 10:13:56 +0200 Subject: [PATCH 09/17] add decoder_residual_scale to tune --- climanet/tune.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/climanet/tune.py b/climanet/tune.py index 3ae2c5e..d12cd5d 100644 --- a/climanet/tune.py +++ b/climanet/tune.py @@ -67,7 +67,7 @@ def _train(tune_config, static_args): # Prepare training configuration training_config = TrainConfig( - calculate_residuals=True, + calculate_residuals=False, num_epoch=static_args["num_epoch"], patience=10, accumulation_steps=tune_config["batch_config"]["accumulation_steps"], @@ -84,11 +84,13 @@ def _train(tune_config, static_args): embed_dim = tune_config["embed_dim"] dropout = tune_config["dropout"] hidden = tune_config["hidden"] + decoder_residual_scale = tune_config["decoder_residual_scale"] model = SpatioTemporalModel( patch_size=(1, patch_size, patch_size), embed_dim=embed_dim, dropout=dropout, hidden=hidden, + decoder_residual_scale=decoder_residual_scale, ) # Train the model From 0686884385c0f6d7528cf914fb78fddf8ac74ed5 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 10:14:29 +0200 Subject: [PATCH 10/17] fix scripts accordingly --- scripts/prediction.py | 24 +----------------------- scripts/training.py | 8 ++++---- scripts/tuning.py | 7 ++++--- 3 files changed, 9 insertions(+), 30 deletions(-) diff --git a/scripts/prediction.py b/scripts/prediction.py index b7a4235..e8a4510 100644 --- a/scripts/prediction.py +++ b/scripts/prediction.py @@ -108,7 +108,7 @@ # Training configuration# create prediction config prediction_config = PredictionConfig( - calculate_residuals=True, + calculate_residuals=False, return_numpy=False, save_predictions=True, return_loss=False, @@ -125,25 +125,3 @@ prediction_config=prediction_config, run_dir=run_dir, ) - - # add residuals to the averaged monthly data - # load hourly data - files = sorted(raw_data_folder.glob(f"{predict_year}*_hr_ERA5dc_masked_{var_name}.nc")) - input_data = xr.open_mfdataset(files) - - # load prediction residuals - files = sorted(run_dir.glob(f"{predict_year}*_{var_name}_prediction_residual.nc")) - predictions_res = xr.open_mfdataset(files) - - # add residuals to the averaged monthly data - input_data_averaged = input_data.resample({"time": "MS"}).mean(skipna=True) - input_data_averaged["time"] = predictions_res["time"] - adjusted_data = input_data_averaged[var_name] + predictions_res - - # save the adjusted data to a new NetCDF file, one file per month - times = adjusted_data.coords["time"].values - - for t in times: - time_str = np.datetime_as_string(t, unit="M").replace("-", "") - file_name = f"{run_dir}/{time_str}_{var_name}_prediction.nc" - adjusted_data.sel(time=[t]).to_netcdf(file_name) diff --git a/scripts/training.py b/scripts/training.py index 86730db..6f20bcc 100644 --- a/scripts/training.py +++ b/scripts/training.py @@ -111,7 +111,7 @@ def _build_dataset( var_name=var_name, land_mask=lsm_mask, crop_size=dataset_crop_size, - stride=dataset_stride, + stride=None, # no overlap for validation model_patch_size=model_patch_size, ) @@ -120,7 +120,7 @@ def _build_dataset( use_cuda = device == "cuda" dataloader_config = DataLoaderConfig( batch_size=100, # adjust if OOM issue - shuffle=True, + shuffle=True, # this is only for training and not for validation num_workers=dataloader_num_workers, pin_memory=use_cuda, persistent_workers=True, @@ -150,10 +150,10 @@ def _build_dataset( # Training configuration training_config = TrainConfig( - calculate_residuals=True, + calculate_residuals=False, num_epoch=101, patience=10, - accumulation_steps=2, + accumulation_steps=1, optimizer_lr=best_config["optimizer_lr"], device=device, verbose=True, diff --git a/scripts/tuning.py b/scripts/tuning.py index b63abc0..789f975 100644 --- a/scripts/tuning.py +++ b/scripts/tuning.py @@ -65,7 +65,7 @@ "land_mask_data": land_mask_data, # dont use ray.put() "load_lazy": False, # one year fits in memory "crop_size": data_crop_size, - "stride": data_stride, + "stride": None, # no overlap for validation "var_name": var_name, } @@ -85,7 +85,7 @@ "dataloader_persistent_workers": True, "dataloader_multiprocessing_context": None, # load_lazy is False "num_epoch": 100, - "max_concurrent_trials": args.num_nodes * 2, # less than GPUs per node (4) avoid OOM + "max_concurrent_trials": args.num_nodes * 3, # less than GPUs per node (4) avoid OOM "experiment_name": "sst_01", } @@ -95,10 +95,11 @@ "embed_dim": tune.choice([32, 64, 128]), "dropout": tune.choice([0.0, 0.1, 0.2]), "hidden": tune.choice([32, 64, 128]), + "decoder_residual_scale": tune.choice([0.1, 0.5, 1.0]), "optimizer_lr": tune.loguniform(1e-3, 1e-1), "batch_config": tune.grid_search( [ - {"batch_size": 100, "accumulation_steps": 2}, + {"batch_size": 100, "accumulation_steps": 1}, ] ), } From 1a6810bcff3846817d15bd6f4093db2439dc3ab5 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 13:23:53 +0200 Subject: [PATCH 11/17] add optimizer_weight_decay to tune --- climanet/tune.py | 1 + 1 file changed, 1 insertion(+) diff --git a/climanet/tune.py b/climanet/tune.py index d12cd5d..c08bf4f 100644 --- a/climanet/tune.py +++ b/climanet/tune.py @@ -72,6 +72,7 @@ def _train(tune_config, static_args): patience=10, accumulation_steps=tune_config["batch_config"]["accumulation_steps"], optimizer_lr=tune_config["optimizer_lr"], + optimizer_weight_decay=tune_config["optimizer_weight_decay"], device=static_args["device"], verbose=False, verbose_epoch_interval=20, From 10d3ffe5ea70eadb09cc8e8ce5d2dc06c22edb27 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 13:24:12 +0200 Subject: [PATCH 12/17] fix scripts --- scripts/training.py | 3 +++ scripts/tuning.py | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/scripts/training.py b/scripts/training.py index 6f20bcc..fcfa259 100644 --- a/scripts/training.py +++ b/scripts/training.py @@ -132,12 +132,14 @@ def _build_dataset( embed_dim = best_config["embed_dim"] dropout = best_config["dropout"] hidden = best_config["hidden"] + decoder_residual_scale = best_config["decoder_residual_scale"] model = SpatioTemporalModel( patch_size=model_patch_size, embed_dim=embed_dim, dropout=dropout, hidden=hidden, + decoder_residual_scale=decoder_residual_scale, ) # move the model to GPU and configure compute resources @@ -155,6 +157,7 @@ def _build_dataset( patience=10, accumulation_steps=1, optimizer_lr=best_config["optimizer_lr"], + optimizer_weight_decay=best_config["optimizer_weight_decay"], device=device, verbose=True, verbose_epoch_interval=10, diff --git a/scripts/tuning.py b/scripts/tuning.py index 789f975..697be6b 100644 --- a/scripts/tuning.py +++ b/scripts/tuning.py @@ -81,7 +81,7 @@ "gpu_per_trial": 1, "run_dir": args.storage_path, "device": "cuda", - "dataloader_num_workers": 1, + "dataloader_num_workers": 4, "dataloader_persistent_workers": True, "dataloader_multiprocessing_context": None, # load_lazy is False "num_epoch": 100, @@ -91,12 +91,13 @@ # parameters to tune tune_config = { - "patch_size": tune.choice([4, 8, 10]), + "patch_size": tune.choice([4, 8]), "embed_dim": tune.choice([32, 64, 128]), "dropout": tune.choice([0.0, 0.1, 0.2]), "hidden": tune.choice([32, 64, 128]), "decoder_residual_scale": tune.choice([0.1, 0.5, 1.0]), "optimizer_lr": tune.loguniform(1e-3, 1e-1), + "optimizer_weight_decay": tune.loguniform(1e-5, 1e-1), "batch_config": tune.grid_search( [ {"batch_size": 100, "accumulation_steps": 1}, From 16ae16ebda9d7c00c75601d43b04a98f4ed5a28e Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Mon, 5 Oct 2026 13:26:29 +0200 Subject: [PATCH 13/17] set calculate_residuals False in data preparation --- scripts/data_preparation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/data_preparation.py b/scripts/data_preparation.py index 0439784..f9c5e57 100644 --- a/scripts/data_preparation.py +++ b/scripts/data_preparation.py @@ -67,7 +67,7 @@ input_data_year, monthly_data_year, run_dir=run_dir_year, - calculate_residuals=True, + calculate_residuals=False, is_hourly=True, save_to_zarr=True, ) From fd03fda94ec62f6078026388f1431b0d006d37a6 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 7 Oct 2026 11:53:43 +0200 Subject: [PATCH 14/17] initialize the conv head in the model --- climanet/st_encoder_decoder.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index 0f4e964..1b3ebaf 100644 --- a/climanet/st_encoder_decoder.py +++ b/climanet/st_encoder_decoder.py @@ -453,6 +453,11 @@ def __init__( # Final conv head to map to single-channel output self.head = nn.Conv2d(channels, 1, kernel_size=1) + # Initialize the final conv head weights and biases to zero for stable + # training + nn.init.zeros_(self.head.weight) + nn.init.zeros_(self.head.bias) + def forward(self, latent, M, out_H, out_W, land_mask=None): """Reconstruct 2D maps from latent patch tokens. Args: From 1517769ccd0dd5a691d10b22444ee160a34eb433 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 7 Oct 2026 11:55:55 +0200 Subject: [PATCH 15/17] normalize residuals before loss calculation --- climanet/dataset.py | 4 ++++ climanet/utils.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/climanet/dataset.py b/climanet/dataset.py index 4053cfc..83e6af0 100644 --- a/climanet/dataset.py +++ b/climanet/dataset.py @@ -44,6 +44,7 @@ def __init__( monthly_da: xr.DataArray, padded_days_mask: xr.DataArray, time_features: xr.DataArray, + res_std: xr.DataArray, land_mask: xr.DataArray = None, model_patch_size: tuple[int, int, int] = (1, 4, 4), # (Month, lat, lon) spatial_dims: tuple[str, str] = ("lat", "lon"), @@ -151,6 +152,7 @@ def __init__( self.geo_pos_t = calculate_sh_geo_pos_embeddings( self.lat_coords, self.lon_coords, self.sh_order_L, self.sh_embed_dim ) + self.res_std_t = torch.from_numpy(res_std.to_numpy()).contiguous() def _compute_crop_indices(self, M: int, H: int, W: int) -> list: """Generate crop start indices with coverage warning (overlap support).""" @@ -322,6 +324,7 @@ def __getitem__(self, idx): # compute geo_pos_embedding and scale_feature for each model patch within the dataset crop geo_pos_embedding_t = self._compute_patch_geo_pos_embedding(i, j) # (Hp*Wp, sh_embed_dim) scale_feature_t = self._compute_patch_scale_features(i, j) # (Hp*Wp, scale_dim) + res_std_t = self.res_std_t[m : m + pm, i : i + ph, j : j + pw] # Convert to dictionary return { @@ -338,6 +341,7 @@ def __getitem__(self, idx): "coords": torch.tensor([m, i, j]), "lat": lat_crop, # (pH,) "lon": lon_crop, # (pW,) + "res_std": res_std_t, # (pm, pH, pW) } def __getitems__(self, indices): diff --git a/climanet/utils.py b/climanet/utils.py index 22ac98a..c2c0b2c 100644 --- a/climanet/utils.py +++ b/climanet/utils.py @@ -305,7 +305,9 @@ def run_one_batch( monthly_pred = monthly_res_pred + mean_data # Compute masked loss - loss = compute_masked_loss(monthly_res_pred, monthly_res, batch["land_mask"]) + monthly_res_norm = monthly_res / batch["res_std"] + monthly_res_pred_norm = monthly_res_pred / batch["res_std"] + loss = compute_masked_loss(monthly_res_pred_norm, monthly_res_norm, batch["land_mask"]) return (loss, monthly_pred) if return_predictions else (loss, None) From b549d901c844d0561500989235f391e9e609f198 Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 7 Oct 2026 12:03:13 +0200 Subject: [PATCH 16/17] add res_std to training script --- scripts/training.py | 98 ++++++++++++++++++++++++--------------------- 1 file changed, 53 insertions(+), 45 deletions(-) diff --git a/scripts/training.py b/scripts/training.py index fcfa259..31a000e 100644 --- a/scripts/training.py +++ b/scripts/training.py @@ -10,42 +10,6 @@ from climanet.utils import configure_compute_resources, read_st_data, set_seed -def _build_dataset( - prepared_data_dir: Path, - years: list[int], - var_name: str, - land_mask: xr.DataArray, - crop_size: tuple[int, int, int], - stride: tuple[int, int], - model_patch_size: tuple[int, int, int], -) -> STDataset: - - data = [read_st_data(data_path=f"{prepared_data_dir}/{year}", var_name=var_name) for year in years] - input_das, input_da_nan_masks, monthly_das, padded_days_masks, time_features_list = zip(*data) - - input_da = xr.concat(input_das, dim="M") - input_da_nan_mask = xr.concat(input_da_nan_masks, dim="M") - monthly_da = xr.concat(monthly_das, dim="M") - padded_days_mask = xr.concat(padded_days_masks, dim="M") - time_features = xr.concat(time_features_list, dim="M") - - return 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=land_mask, - crop_size=crop_size, - stride=stride, - model_patch_size=model_patch_size, - sh_embed_dim=96, - sh_order_L=10, - verbose=False, - load_lazy=True, - ) - - if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -94,25 +58,69 @@ def _build_dataset( model_patch_size = (1, best_config["patch_size"], best_config["patch_size"]) train_years = [2018, 2019, 2020] - dataset_train = _build_dataset( - prepared_data_dir=prepared_data_dir, - years=train_years, - var_name=var_name, + data = [read_st_data(data_path=f"{prepared_data_dir}/{year}", var_name=var_name) for year in train_years] + input_das, input_da_nan_masks, monthly_das, padded_days_masks, time_features_list = zip(*data) + + input_da = xr.concat(input_das, dim="M") + input_da_nan_mask = xr.concat(input_da_nan_masks, dim="M") + monthly_da = xr.concat(monthly_das, dim="M") + padded_days_mask = xr.concat(padded_days_masks, dim="M") + time_features = xr.concat(time_features_list, dim="M") + + # calculate std of res + mean_monthly_da = input_da.where(~input_da_nan_mask).mean(dim="T", skipna=True) + mean_monthly_da["M"] = monthly_da["M"] + monthly_residuals = monthly_da - mean_monthly_da + monthly_std = monthly_residuals.groupby("M.month").std(dim=["M", "lat", "lon"], skipna=True) + res_std_train = ( + monthly_std + .rename({"month": "M"}) + .assign_coords(M=monthly_residuals["M"].isel(M=slice(0, 12)).values) + .broadcast_like(monthly_residuals.isel(M=slice(0, 12))) + ) + + dataset_train = 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, + res_std=res_std_train, land_mask=lsm_mask, crop_size=dataset_crop_size, stride=dataset_stride, model_patch_size=model_patch_size, + sh_embed_dim=96, + sh_order_L=10, + verbose=False, + load_lazy=True, ) validation_year = [2021] - dataset_validation = _build_dataset( - prepared_data_dir=prepared_data_dir, - years=validation_year, - var_name=var_name, + data = [read_st_data(data_path=f"{prepared_data_dir}/{year}", var_name=var_name) for year in validation_year] + input_das, input_da_nan_masks, monthly_das, padded_days_masks, time_features_list = zip(*data) + + input_da = xr.concat(input_das, dim="M") + input_da_nan_mask = xr.concat(input_da_nan_masks, dim="M") + monthly_da = xr.concat(monthly_das, dim="M") + padded_days_mask = xr.concat(padded_days_masks, dim="M") + time_features = xr.concat(time_features_list, dim="M") + + dataset_validation = 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, + res_std=res_std_train, # use the training residual std for validation land_mask=lsm_mask, crop_size=dataset_crop_size, - stride=None, # no overlap for validation + stride=None, # no overlap for validation model_patch_size=model_patch_size, + sh_embed_dim=96, + sh_order_L=10, + verbose=False, + load_lazy=True, ) # Build the dataloader config From 0e073124da1569aebe7473d884d3eadbfcea54ac Mon Sep 17 00:00:00 2001 From: SarahAlidoost Date: Wed, 7 Oct 2026 12:45:52 +0200 Subject: [PATCH 17/17] fix training script --- scripts/training.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/scripts/training.py b/scripts/training.py index 31a000e..2bf4c9b 100644 --- a/scripts/training.py +++ b/scripts/training.py @@ -72,12 +72,8 @@ mean_monthly_da["M"] = monthly_da["M"] monthly_residuals = monthly_da - mean_monthly_da monthly_std = monthly_residuals.groupby("M.month").std(dim=["M", "lat", "lon"], skipna=True) - res_std_train = ( - monthly_std - .rename({"month": "M"}) - .assign_coords(M=monthly_residuals["M"].isel(M=slice(0, 12)).values) - .broadcast_like(monthly_residuals.isel(M=slice(0, 12))) - ) + monthly_std = monthly_std.sel(month=monthly_residuals["M.month"]) + res_std_train = monthly_std.broadcast_like(monthly_residuals) dataset_train = STDataset( input_da=input_da, @@ -106,13 +102,14 @@ padded_days_mask = xr.concat(padded_days_masks, dim="M") time_features = xr.concat(time_features_list, dim="M") + res_std_val =res_std_train.isel(M=slice(0, 12)) dataset_validation = 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, - res_std=res_std_train, # use the training residual std for validation + res_std=res_std_val, # use the training residual std for validation land_mask=lsm_mask, crop_size=dataset_crop_size, stride=None, # no overlap for validation