[WIP] Standard baseline configs

icenet_mp/callbacks/ema_weight_averaging_callback.py

@@ -1,3 +1,6 @@
+from typing import Any
+
+from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import WeightAveraging
 from torch.optim.swa_utils import get_ema_multi_avg_fn
 
@@ -26,6 +29,18 @@ def __init__(
         self.every_n_epochs = every_n_epochs
         self.every_n_steps = every_n_steps
 
+    def on_train_batch_end(
+        self, trainer: Trainer, pl_module: LightningModule, *args: Any, **kwargs: Any
+    ) -> None:
+        """Ignore the update if the module has no parameters."""
+        if next(pl_module.parameters(), None) is not None:
+            super().on_train_batch_end(trainer, pl_module, *args, **kwargs)
+
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Ignore the update if the module has no parameters."""
+        if next(pl_module.parameters(), None) is not None:
+            super().on_train_epoch_end(trainer, pl_module)
+
     def should_update(
         self, step_idx: int | None = None, epoch_idx: int | None = None
     ) -> bool:

icenet_mp/callbacks/metric_summary_callback.py

@@ -25,6 +25,8 @@ def log_per_epoch_metrics(
         """Log per-epoch metrics to W&B."""
         for name, metric in metrics.items():
             # Compute the metric value (e.g., SIEError) across all batches
+            if not metric._update_called:
+                continue
             values: Tensor = metric.compute()
 
             # Log the mean value of the metric across all days
@@ -56,6 +58,8 @@ def log_per_run_metrics(
         values_per_forecast_day: dict[str, dict[str, Tensor]] = defaultdict(dict)
         for stage, metric_collection in metrics.items():
             for metric_name, metric in metric_collection.items():
+                if not metric._update_called:
+                    continue
                 metric_tensor: Tensor = metric.compute()
                 if metric_tensor.reshape(-1).shape[0] > 1:
                     values_per_forecast_day[metric_name][stage] = metric_tensor

icenet_mp/compatibility/__init__.py

@@ -12,7 +12,7 @@
 
 def configure_external_libraries() -> None:
     """Configure any external libraries used by the application."""
-    log.info("Configuring external libraries...")
+    log.debug("Configuring external libraries...")
     patch_parameter_deepcopy()
     register_accelerators()
     register_animation_backends()

icenet_mp/config/base.yaml

@@ -4,7 +4,7 @@ defaults:
   - evaluate: default
   - loggers:
     - wandb
-  - model: naive_unet_naive
+  - model: 00_naive_unet_naive
   - predict: sic-icenet-2d
   - random: default
   - train: default

icenet_mp/config/model/naive_unet_naive.yaml → icenet_mp/config/model/00_naive_unet_naive.yaml

@@ -15,8 +15,8 @@ encoders:
 
 processor:
   _target_: icenet_mp.models.processors.UNetProcessor
-  start_out_channels: 100 # Reduce number of channels to support 21 day forecasts
+  start_out_channels: 128
 
 decoder:
   _target_: icenet_mp.models.decoders.NaiveLinearDecoder
-  bounded: false # Whether to bound the output between 0 and 1
+  bounded: true

icenet_mp/config/model/cnn_unet_cnn.yaml → icenet_mp/config/model/02_cnn_unet_cnn.yaml

@@ -6,16 +6,31 @@ encoders:
   latent_space: [144, 144]
   era5:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    kernel_size: 5
+    n_layers: 1
+    scale_factor: 3
   float-argo:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    kernel_size: 5
+    n_layers: 1
+    scale_factor: 3
   sic-icenet:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    kernel_size: 5
+    n_layers: 1
+    scale_factor: 3
   sic-ssmis:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    n_layers: 1
+    scale_factor: 3
 
 processor:
   _target_: icenet_mp.models.processors.UNetProcessor
+  start_out_channels: 128
 
 decoder:
   _target_: icenet_mp.models.decoders.CNNDecoder
-  bounded: false # Whether to bound the output between 0 and 1
+  bounded: true
+  kernel_size: 5
+  n_layers: 1
+  scale_factor: 3

icenet_mp/config/model/cnn_vit_cnn.yaml → icenet_mp/config/model/03_cnn_vit_cnn.yaml

@@ -3,19 +3,30 @@ _target_: icenet_mp.models.EncodeProcessDecode
 name: cnn-vit-cnn
 
 encoders:
-  latent_space: [144, 144]
+  latent_space: [216, 216]
   era5:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    n_layers: 1
   float-argo:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    n_layers: 1
   sic-icenet:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    n_layers: 1
   sic-ssmis:
     _target_: icenet_mp.models.encoders.CNNEncoder
+    n_layers: 1
 
 processor:
   _target_: icenet_mp.models.processors.VitProcessor
+  depth: 6        # number of transformer blocks
+  dropout: 0.1    # dropout rate
+  emb_dim: 256    # embedding dimension
+  heads: 8        # number of attention heads
+  mlp_dim: 1024   # dimension of the MLP in the transformer block
+  patch_size: 24  # 216 / 24 = 9 patches per dimension (81 in total)
 
 decoder:
   _target_: icenet_mp.models.decoders.CNNDecoder
-  bounded: false # Whether to bound the output between 0 and 1
+  bounded: true
+  n_layers: 1

icenet_mp/config/model/ddpm.yaml → icenet_mp/config/model/04_ddpm.yaml

@@ -2,3 +2,6 @@ _target_: icenet_mp.models.ddpm.DDPM
 
 # Run DDPM model with default settings
 name: ddpm
+dropout_rate: 0.05
+start_out_channels: 128
+timesteps: 500

icenet_mp/config/model/piecewise_unet_piecewise.yaml → icenet_mp/config/model/05_piecewise_unet_piecewise.yaml

@@ -18,3 +18,4 @@ processor:
 
 decoder:
   _target_: icenet_mp.models.decoders.PiecewiseDecoder
+  restrict_range: sigmoid

icenet_mp/config/train/scheduler/ddpm.yaml

@@ -5,4 +5,3 @@ defaults:
 
 scheduler_parameters:
   eta_min: 5e-6
-  T_max: 100

icenet_mp/data_processors/data_downloader.py

@@ -1,10 +1,12 @@
 import logging
 import shutil
+from contextlib import suppress
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 from anemoi.datasets import open_dataset
+from anemoi.datasets.commands.cleanup import Cleanup
 from anemoi.datasets.commands.finalise import Finalise
 from anemoi.datasets.commands.init import Init
 from anemoi.datasets.commands.inspect import InspectZarr
@@ -14,6 +16,7 @@
 from zarr.errors import PathNotFoundError
 
 from icenet_mp.types import (
+    AnemoiCleanupArgs,
     AnemoiDatasetStatus,
     AnemoiFinaliseArgs,
     AnemoiInitArgs,
@@ -46,6 +49,14 @@ def __init__(
         self.recipe = Recipe(**anemoi_config)
         self.preprocessor = cls_preprocessor(anemoi_config)
 
+    def artifacts(self) -> list[Path]:
+        """Return a list of temporary artifacts created during the download and finalise process."""
+        return [
+            path
+            for path in self.path_dataset.parent.glob(f"{self.path_dataset.stem}.*")
+            if path != self.path_dataset
+        ]
+
     def check_status(self) -> AnemoiDatasetStatus:
         """Return the status of the dataset."""
         try:
@@ -85,7 +96,9 @@ def check_status(self) -> AnemoiDatasetStatus:
         return AnemoiDatasetStatus(
             copy_in_progress=copy_in_progress,
             download_complete=download_complete,
-            is_finalised=download_complete and statistics_ready,
+            is_finalised=download_complete
+            and statistics_ready
+            and not self.artifacts(),
         )
 
     def create(self, *, overwrite: bool = False) -> None:
@@ -169,22 +182,33 @@ def finalise(self, *, overwrite: bool, status: AnemoiDatasetStatus) -> None:
             )
             logger.info("Finalised dataset %s at %s.", self.name, self.path_dataset)
 
-        # create active grid cell and land masks for the SSMIS dataset
+        # Create active grid cell and land masks if appropriate
         self.generate_masks(overwrite=overwrite)
 
+        # Cleanup any temporary artifacts created during the download and finalise process
+        if self.artifacts():
+            with suppress(ValueError):
+                Cleanup().run(AnemoiCleanupArgs(path=str(self.path_dataset)))
+            if remaining := self.artifacts():
+                logger.warning("Residual artifacts for dataset %s:", self.name)
+                for artifact in remaining:
+                    logger.warning("... %s", artifact)
+            else:
+                logger.info("Cleaned up temporary artifacts for dataset %s.", self.name)
+
     def generate_masks(self, *, overwrite: bool) -> None:
         """Generate land and active grid cell masks for the SSMIS dataset."""
-        # if there is an SSMIS dataset, create the masks
+        # Create the masks if this is an SSMIS dataset
         if "ssmis" not in self.name:
-            logger.info("Not SSMIS dataset, skipping mask creation.")
             return
+        logger.debug("Generating land and active grid cell masks for SSMIS dataset.")
 
         self.path_masks.mkdir(parents=True, exist_ok=True)
         land_mask_path = self.path_masks / "land_mask.npy"
         active_mask_path = self.path_masks / "active_mask.npy"
 
         if land_mask_path.exists() and active_mask_path.exists() and not overwrite:
-            logger.info("Both masks already exist, skipping creation.")
+            logger.debug("Both masks already exist, skipping creation.")
             return
 
         # Unpack status flags into a binary array, skipping any missing dates
@@ -205,7 +229,7 @@ def generate_masks(self, *, overwrite: bool) -> None:
 
         # land mask: land = 0, sea = 1
         if land_mask_path.exists() and not overwrite:
-            logger.info("Land mask already exists, skipping creation.")
+            logger.debug("Land mask already exists, skipping creation.")
             land_mask = np.load(land_mask_path)
         else:
             land_mask = np.squeeze(binary[..., [7]]).sum(axis=0)
@@ -219,7 +243,7 @@ def generate_masks(self, *, overwrite: bool) -> None:
 
         # active mask: active grid cells = 1, inactive = 0
         if active_mask_path.exists() and not overwrite:
-            logger.info("Active mask already exists, skipping creation.")
+            logger.debug("Active mask already exists, skipping creation.")
         else:
             # Identify grid cells that are inactive for all time steps
             inactive_count = np.squeeze(binary[..., [0]]).sum(axis=0)