Add ConvGRU ensemble training pipeline

[franchg]

Summary

• Transfers the ConvGRU ensemble model from ConvGRU-Ensemble into mlcast as a trainable model, ready for the hackathon
• Adds Fiddle-based configuration using the @auto_config pattern (as discussed in Designing configuration infrastructure for mlcast python package #5 and used in weatherduck)
• Replaces the skeleton data layer with a sampled-CSV-based pipeline compatible with mlcast-dataset-sampler output

New files

• losses.py — CRPS, afCRPS, MaskedLoss, and build_loss() factory
• utils.py — rain rate ↔ normalized reflectivity conversions (Marshall-Palmer)
• models/convgru.py — RadarLightningModel with ensemble support, TensorBoard image logging, and inference API
• configs.py — @auto_config experiment factory following the weatherduck pattern
• __main__.py — CLI entry point (python -m mlcast --zarr-path ... --csv-path ...)

Updated files

• modules/convgru_modules.py — added ensemble generation via noisy decoder inputs
• data/zarr_dataset.py — SampledRadarDataset loading datacubes from Zarr using CSV (t, x, y) coordinates
• data/zarr_datamodule.py — RadarDataModule with chronological train/val/test splits and augmentation
• pyproject.toml — added fiddle, pandas, pytorch-lightning, torchvision dependencies

Usage (hackathon)

# 1. Data is already on Leonardo as Zarr
# 2. Run sampler to get CSV:
#    mlcast-dataset-sampler filter-nan ... && mlcast-dataset-sampler sample ...
# 3. Train:
python -m mlcast --zarr-path /path/to/data.zarr --csv-path /path/to/sampled.csv --variable-name RR

# Or programmatically with config overrides:
from mlcast.configs import convgru_experiment
import fiddle as fdl

cfg = convgru_experiment.as_buildable(zarr_path="...", csv_path="...", variable_name="RR")
cfg.data.batch_size = 32
cfg.trainer.max_epochs = 50
fdl.build(cfg).run()

Closes #7 (partially — brings ensemble ConvGRU into mlcast)
Implements #5 (Fiddle configuration)

Test plan

• Install with uv pip install -e . and verify imports work
• Run python -m mlcast --help to verify CLI
• Train on a small Zarr + CSV on Leonardo to verify end-to-end pipeline
• Verify TensorBoard logging works

[leifdenby]

Ok, @franchg I think I'm done with my refactor 🥳 I hope you like what I've done to your code 🎁 ➡️ https://github.com/leifdenby/mlcast/tree/feat/convgru-ensemble-training

I've (with my agent's help) written up this list of changes (below). I went through this by hand and added/removed detail so I think it gives quite a complete overview. As well as reading the list below (hopefully it makes sense...) if you could also read through the updated README that would be great. I tried to take care to make it clear how the CLI now works, what the code structure is, and what ConvGRU (our first architecture) is. In the README I have also added an example of how to couple in a new architecture not yet in the mlcast codebase (by borrowing the HalfUNet from mfai).

If you are happy with these change my suggestion is that:

• we merge my commits into your branch, I've made a PR here Feat/convgru ensemble training franchg/mlcast#1
• we squash merge this PR into main
• we tag this as v0.1.0 of mlcast and I write a changelog entry explaining the functionality that is there now.

We can then build on this when we discover things I have overlooked :)

Changes

What Changed

• reworked the config/orchestration layer so dataset, model, trainer, logger, and remote-data choices are expressed as explicit Fiddle config nodes plus fiddlers, by splitting the old config surface into config/base.py, fiddlers.py, consistency_checks.py, loader.py, and orchestrator.py.
  • The idea here is to have a "base" config (ConvGRU training to a path that doesn't exist, maybe this shoudl change?), and use --config set: to set specific config parameters or apply fiddlers where multiple parameters need to be changed together (for example list of variables to load from zarr and number if input channels in the model)
  • once the config is built we then apply consistency checks that tries to ensure that the config is self-consistent (is this dataset compatible with this model?), and then finally built the Experiment object that actually calls the Training.fit() training loop start method.
  • I also made it so that the base config can be replaced either through 1) naming an alternative config (we might want to create base configs for different models in future maybe) or 2) passing in a path to a .yaml-file which is then loaded and parsed as a fiddle Config object.
• separated the generic training wrapper from the ConvGRU architecture by refactoring the old Lightning model into a NowcastLightningModule that accepts an injected network, and by moving ConvGRU-specific logic into the model layer.
• pushed more of the temporal contract into the dataset layer by moving forecast_steps out of the Lightning module and making input_steps and forecast_steps the primary dataset parameters, with steps derived from them.
• made the dataset sample interface more explicit by changing dataset outputs from a single time tensor to a structured sample with input, target, and target_mask.
• on variable selection: generalized variable selection by switching dataset lookup to CF standard_name via cf_xarray, instead of relying on a broader set of source-specific variable names.
  • made normalization follow the same abstraction by introducing normalization registries keyed by standard_name, so preprocessing is aligned with how variables are selected.
  • improved support for multi-variable inputs by deriving channel-related behavior from standard_names and by updating the custom-network examples to compute dimensions from config instead of assuming a single variable.
• added a second sampling strategy by introducing SourceDataRandomSamplingDataset alongside the precomputed sampler (useful when not wanting to generate a sampler CSV, and makes space for switching other samplers in future), and then made the data module consume a dataset_factory so either sampler can be injected through config.
• fixed target masking semantics by computing masks only on the forecast target window and per timestep/per channel, before NaNs are filled, instead of collapsing validity across time (I think there was a bug here before, the pytorch.Dataset did reduction of time to create a mask whereas the nn.Module.forward() call for the forecasting part did time-based indexing)
• introduced use of jaxtyping and beartype to communicate what the shapes of tensors mean, and to check that calls are consistent on this. This will help (I hope) to make the interface to our nn.Module.forward(...) methods clear so that it will be simpler to add new architectures in future.
• moved some model-adjacent responsibilities out of the core architecture by separating visualization logic and simplifying loss construction and naming.
• expanded test coverage across config, data, normalization, losses, CLI, orchestrator, and README snippets, so the refactor is checked through the interfaces people actually use.
• expanded the documentation substantially by rewriting the README, adding architecture diagrams, adding a config diagram generator/hook, documenting custom model integration, and tightening the install instructions.
• added MLflow as a config-level logging option by introducing a use_mlflow_logger fiddler plus a dedicated callback for MLflow tags, system metrics, and run URLs, rather than wiring MLflow directly into the training loop.

Why

All these changes were motivate by a desire to:

• make the system easier to reason about as a config-driven training stack, so the important choices live in the config graph instead of being partly hidden in training code.
• reduce coupling between dataset shape semantics, forecast horizon handling, model architecture, and Lightning orchestration, because those concerns looked like they were bleeding into each other.
• make custom architectures easier to plug in, especially non-ConvGRU models, without forcing them to conform to assumptions that I think really belong to one specific model family.
• make the temporal contract clearer by distinguishing “what the model sees” (input_steps) from “what the model predicts” (forecast_steps) at the dataset/config boundary.
• make variable handling more robust by using CF standard_name as the common key for selection, validation, and normalization.
• make data access more portable by supporting both local and anonymous S3-backed Zarr stores through the same dataset/config interface.
• make experimentation easier by allowing sampler choice, logger choice, and some data/model behaviors to be swapped through fiddlers instead of code edits.
• make masking behavior more correct and less surprising, especially for future probabilistic or multi-step work where per-timestep target validity matters.
• make the data layer more extensible by giving both samplers the same interface and shared base behavior, which should lower the cost of adding more dataset variants later.
• make the project easier for other people to pick up and modify by strengthening tests, docs, examples, and generated diagrams around the refactored interfaces.

What This Now Enables

With this, it is now:

• easier to plug in non-ConvGRU networks, because the Lightning wrapper is more generic, the validation logic is less tied to ConvGRU-only parameters and input_steps is directly available on the dataset config.
• more natural to support multiple variables end-to-end, because variable lookup, normalization, and example model wiring all derive from standard_names.
• possible to swap in different sampling approaches, because the data module now consumes an injected dataset_factory.

In addition the support for mlflow and remote S3 Zarr stores were convenient for my own experimentation, but I thought that enabling this flexiblity should make it easier to add more choices here in future.

[leifdenby]

@franchg just a heads-up, I have continue working from this branch in https://github.com/leifdenby/mlcast/tree/feat/dmi-convgru-training (because I think you are continue with this branch yourself and have already merged my other branch).

I found some issues, made some additions on my branch and changed a few things:

• bug: issue is that DMI radar is stored as float64 (which the other datasets are not), a these need converting to float32 before returning from the pytorch Dataset class, wasn't handled previously
• bug: .gitignore was excluding the mlcast.data submodule so that uv add {mlcast-repo-path} install wouldn't include mlcast.data submodule
• addition: out DMI data is stored as reflectivity factor, and i've tried to add that (but I don't know if it makes sense)
• changed: the exiting way of defining splits was to pass train/val split as fractions to the data-module, from this the data-module computed the time-index range for each split and passe these to the dataset class __init__. I think passing time-indexes rather than time-values to the dataset class is a bit inflexible since we humans generally think in timespans rather than indices from the dataset start. So, my idea is to support both like so:

from mlcast.config import training_experiment
# Returns a fdl.Config graph — nothing is instantiated yet.
cfg = training_experiment.as_buildable()

## Ratio mode — chronological split by fraction of the time axis:
# Implicit test split (remainder = 15%)
cfg.data.splits = {"train": 0.70, "val": 0.15}
# Explicit test fraction
cfg.data.splits = {"train": 0.70, "val": 0.15, "test": 0.15}

## Datetime mode — explicit date boundaries:
# All three splits
cfg.data.splits = {
    "train": ("2016-02-29", "2022-12-31"),
    "val":   ("2023-01-01", "2023-12-31"),
    "test":  ("2024-01-01", "2025-12-31"),
}
# No test split (test_dataset will be None)
cfg.data.splits = {
    "train": ("2016-02-29", "2022-12-31"),
    "val":   ("2023-01-01", "2023-12-31"),
    "test":  None,
}

the data module then passes these split definitions to each dataset instances created, and internally in the dataset __init__ it is computed what time-indexes the start and end times provided match with, and then use these to subset the sampling CSVs. This way we don't tie ourself to providing time-indexes to the dataset classes as well.

[leifdenby]

@franchg and I discussed this PR yesterday and agreed to merge it in and the continue work by creating further PRs. The purpose of doing this is to give more visibility to the work we have done here, and make it easier for others to build on this work.

We expect to make further PRs to address the following:

• @franchg is refactoring the source dataset sampler (https://github.com/mlcast-community/mlcast-dataset-sampler) to 1) use a faster implementation of the intensity based sampling (by retaining statistics of extracted sequences found during the nan-filtering step) and 2) to switch to parquet-based storage for the sampling information so that meta information (about the parameters of the sampling process) can be included directly in the parquet files
• I will be adding a PR to introduce the generalisation of how dataset splits are defined by 1) allowing for splitting (in future) on other coordinates besides time by making it explicit in the argument that defines the splits what coordinate to be split on, and allowing for range-based defined splits (tuples like so (2020-01-01T12:00, 2021-01-01T12:00)

And additional PRs that fix the following bugs I identified earlier:

• issue is that DMI radar is stored as float64 (which the other datasets are not), a these need converting to float32 before returning from the pytorch Dataset class, wasn't handled previously
• .gitignore was excluding the mlcast.data submodule so that uv add {mlcast-repo-path} install wouldn't include mlcast.data submodule