fix: save eval predictions to zarr (multi-GPU safe, correct time semantics)

- Multi-GPU safe: rank-0 writes metadata template (compute=False), then
  barrier(), then all ranks write via region='auto'
- Correct start_time: analysis_time = batch_times[:,0] - step_length
- Rescale predictions to original scale in test_step before writing
- Drop raw 'time' coord; explicit int64 encoding for time variables
- Cache WeatherDataset per (split, category) to avoid O(N*T) re-instantiation
- Add --save-eval-to-zarr-path CLI flag
- Add tests/test_zarr_eval.py single-GPU integration test

Refs mllam#104
Part of mllam#138

Author: NullPointer-cell

neural_lam/models/ar_model.py

@@ -6,9 +6,11 @@
 # Third-party
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import pytorch_lightning as pl
 import torch
 import xarray as xr
+from loguru import logger
 
 # First-party
 from neural_lam.utils import get_integer_time
@@ -160,6 +162,8 @@ def __init__(
             self._datastore.step_length
         )
 
+        self._weather_dataset_cache: Dict[tuple, Any] = {}
+
     def _create_dataarray_from_tensor(
         self,
         tensor: torch.Tensor,
@@ -191,7 +195,12 @@ def _create_dataarray_from_tensor(
         # TODO: creating an instance of WeatherDataset here on every call is
         # not how this should be done but whether WeatherDataset should be
         # provided to ARModel or where to put plotting still needs discussion
-        weather_dataset = WeatherDataset(datastore=self._datastore, split=split)
+        cache_key = (split, category)
+        if cache_key not in self._weather_dataset_cache:
+            self._weather_dataset_cache[cache_key] = WeatherDataset(
+                datastore=self._datastore, split=split
+            )
+        weather_dataset = self._weather_dataset_cache[cache_key]
         time = np.array(time.cpu(), dtype="datetime64[ns]")
         da = weather_dataset.create_dataarray_from_tensor(
             tensor=tensor, time=time, category=category
@@ -381,6 +390,113 @@ def on_validation_epoch_end(self):
         for metric_list in self.val_metrics.values():
             metric_list.clear()
 
+    def _save_predictions_to_zarr(
+        self,
+        batch_times: torch.Tensor,
+        batch_predictions: torch.Tensor,
+        batch_idx: int,
+        zarr_output_path: str,
+    ):
+        """
+        Save state predictions for a single batch to a Zarr dataset.
+
+        Predictions must already be in the original (un-standardized) scale
+        before calling this method. The resulting dataset contains a variable
+        named ``state`` with coordinates
+        ``(start_time, elapsed_forecast_duration, grid_index, state_feature)``
+        (or ``x``/``y`` instead of ``grid_index`` for regular-grid datastores).
+
+        Multi-GPU safe: rank 0 writes the full-extent metadata template on the
+        first batch, all ranks then synchronise via ``barrier()``, and every
+        rank writes its slice using ``region="auto"``.
+
+        Parameters
+        ----------
+        batch_times : torch.Tensor
+            Forecast times, shape ``(B, pred_steps)`` as int64 nanoseconds
+            since epoch.  These are the *predicted* times, so
+            ``batch_times[:, 0]`` is ``analysis_time + step_length``.
+        batch_predictions : torch.Tensor
+            Predictions in the **original** data scale, shape
+            ``(B, pred_steps, num_grid_nodes, d_f)``.
+        batch_idx : int
+            Index of the current batch in the epoch.
+        zarr_output_path : str
+            Filesystem path where the Zarr store will be written.
+        """
+        batch_predictions = batch_predictions.cpu()
+        batch_times_cpu = batch_times.cpu()
+        batch_size = batch_predictions.shape[0]
+
+        step_ns = int(pd.Timedelta(self._datastore.step_length).value)
+
+        analysis_times_ns = (
+            batch_times_cpu[:, 0].numpy().astype("int64") - step_ns
+        )  # (B,)
+        analysis_times = analysis_times_ns.astype("datetime64[ns]")
+
+        pred_steps = batch_times_cpu.shape[1]
+        offsets_ns = np.arange(1, pred_steps + 1, dtype="int64") * step_ns
+        elapsed = offsets_ns.astype("timedelta64[ns]")
+
+        time_encoding = {
+            "start_time": {
+                "units": "Seconds since 1970-01-01 00:00:00",
+                "dtype": "int64",
+            },
+            "elapsed_forecast_duration": {
+                "units": "seconds",
+                "dtype": "int64",
+            },
+        }
+
+        das = []
+        for i in range(batch_size):
+            da_i = self._create_dataarray_from_tensor(
+                tensor=batch_predictions[i],
+                time=batch_times_cpu[i],
+                split="test",
+                category="state",
+            )
+            da_i = da_i.assign_coords(
+                elapsed_forecast_duration=(
+                    "time",
+                    elapsed,
+                )
+            )
+            da_i = da_i.swap_dims({"time": "elapsed_forecast_duration"})
+            da_i = da_i.drop_vars("time", errors="ignore")
+            da_i.name = "state"
+            das.append(da_i)
+
+        da_batch = xr.concat(das, dim="start_time")
+        da_batch = da_batch.assign_coords(
+            start_time=("start_time", analysis_times)
+        )
+        da_batch = da_batch.chunk({"start_time": batch_size})
+        ds_batch = da_batch.to_dataset(name="state")
+
+        if self.trainer.is_global_zero and batch_idx == 0:
+            logger.info(f"Creating Zarr store at {zarr_output_path}")
+            all_times = self._datastore.get_dataarray(
+                category="state", split="test"
+            ).coords["time"].values
+            template_da = da_batch.reindex(
+                start_time=all_times, fill_value=np.nan
+            )
+            template_ds = template_da.to_dataset(name="state")
+            template_ds.to_zarr(
+                zarr_output_path,
+                compute=False,
+                mode="w",
+                encoding=time_encoding,
+                consolidated=True,
+            )
+
+        self.trainer.strategy.barrier()
+
+        ds_batch.to_zarr(zarr_output_path, region="auto")
+
     # pylint: disable-next=unused-argument
     def test_step(self, batch, batch_idx):
         """
@@ -445,6 +561,17 @@ def test_step(self, batch, batch_idx):
         self.spatial_loss_maps.append(log_spatial_losses)
         # (B, N_log, num_grid_nodes)
 
+        # Save predictions to Zarr if requested
+        if getattr(self.args, "save_eval_to_zarr_path", None):
+            # Rescale from standardized space to original data scale
+            pred_rescaled = prediction * self.state_std + self.state_mean
+            self._save_predictions_to_zarr(
+                batch_times=batch_times,
+                batch_predictions=pred_rescaled,
+                batch_idx=batch_idx,
+                zarr_output_path=self.args.save_eval_to_zarr_path,
+            )
+
         # Plot example predictions (on rank 0 only)
         if (
             self.trainer.is_global_zero

neural_lam/train_model.py

@@ -170,6 +170,15 @@ def main(input_args=None):
         default=1,
         help="Number of example predictions to plot during evaluation",
     )
+    parser.add_argument(
+        "--save-eval-to-zarr-path",
+        type=str,
+        default=None,
+        dest="save_eval_to_zarr_path",
+        help="If set, save evaluation predictions (in original data scale) "
+        "to a Zarr store at this path. Multi-GPU safe: rank 0 writes the "
+        "metadata template, then all ranks write their slices in parallel.",
+    )
 
     # Logger Settings
     parser.add_argument(
@@ -307,13 +316,24 @@ def main(input_args=None):
         datastore=datastore, args=args, run_name=run_name
     )
 
-    checkpoint_callback = pl.callbacks.ModelCheckpoint(
+    val_checkpoint = pl.callbacks.ModelCheckpoint(
         dirpath=f"saved_models/{run_name}",
         filename="min_val_loss",
         monitor="val_mean_loss",
         mode="min",
-        save_last=True,
+        save_top_k=1,
     )
+    
+    latest_checkpoint = pl.callbacks.ModelCheckpoint(
+        dirpath=f"saved_models/{run_name}",
+        filename="last",
+        monitor=None,
+        save_top_k=1,
+        every_n_epochs=1,
+        save_on_train_epoch_end=True,
+        enable_version_counter=False,
+    )
+
     trainer = pl.Trainer(
         max_epochs=args.epochs,
         deterministic=True,
@@ -323,7 +343,7 @@ def main(input_args=None):
         devices=devices,
         logger=training_logger,
         log_every_n_steps=1,
-        callbacks=[checkpoint_callback],
+        callbacks=[val_checkpoint, latest_checkpoint],
         check_val_every_n_epoch=args.val_interval,
         precision=args.precision,
     )

tests/test_zarr_eval.py

@@ -0,0 +1,128 @@
+# Standard library
+from pathlib import Path
+
+# Third-party
+import numpy as np
+import pytest
+import pytorch_lightning as pl
+import torch
+import xarray as xr
+
+# First-party
+from neural_lam import config as nlconfig
+from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.models.graph_lam import GraphLAM
+from neural_lam.weather_dataset import WeatherDataModule
+from tests.conftest import init_datastore_example
+
+
+class ModelArgs:
+    """Minimal args object for testing zarr eval saving."""
+
+    output_std = False
+    loss = "mse"
+    restore_opt = False
+    n_example_pred = 0  # skip plotting to keep the test fast
+    graph = "1level"
+    hidden_dim = 4
+    hidden_layers = 1
+    processor_layers = 2
+    mesh_aggr = "sum"
+    lr = 1.0e-3
+    val_steps_to_log = [1, 2]
+    metrics_watch = []
+    num_past_forcing_steps = 1
+    num_future_forcing_steps = 1
+    save_eval_to_zarr_path = None  # overridden per test
+
+
+def run_zarr_eval(datastore, zarr_path, tmp_path):
+    """
+    Run one test epoch using the given datastore and save predictions to
+    *zarr_path*. Returns the opened ``xr.Dataset``.
+    """
+    device_name = "cuda" if torch.cuda.is_available() else "cpu"
+
+    trainer = pl.Trainer(
+        max_epochs=1,
+        deterministic=True,
+        accelerator=device_name,
+        devices=1,
+        log_every_n_steps=1,
+        logger=False,
+        enable_checkpointing=False,
+    )
+
+    graph_name = "1level"
+    graph_dir_path = Path(datastore.root_path) / "graph" / graph_name
+    if not graph_dir_path.exists():
+        create_graph_from_datastore(
+            datastore=datastore,
+            output_root_path=str(graph_dir_path),
+            n_max_levels=1,
+        )
+
+    data_module = WeatherDataModule(
+        datastore=datastore,
+        ar_steps_train=3,
+        ar_steps_eval=3,
+        standardize=True,
+        batch_size=2,
+        num_workers=0,
+        num_past_forcing_steps=1,
+        num_future_forcing_steps=1,
+    )
+
+    model_args = ModelArgs()
+    model_args.save_eval_to_zarr_path = str(zarr_path)
+
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        )
+    )
+
+    model = GraphLAM(args=model_args, datastore=datastore, config=config)
+    trainer.test(model=model, datamodule=data_module)
+
+    return xr.open_zarr(str(zarr_path))
+
+
+def test_zarr_eval_single_gpu(tmp_path):
+    """
+    Single-GPU integration test: run eval on DummyDatastore and assert that
+    the Zarr output has the expected structure and correct time semantics.
+    """
+    datastore = init_datastore_example("dummydata")
+    zarr_path = tmp_path / "eval_preds.zarr"
+
+    ds = run_zarr_eval(datastore, zarr_path, tmp_path)
+
+    # 1. Store must exist and contain "state"
+    assert zarr_path.exists(), "Zarr store was not created"
+    assert "state" in ds.data_vars
+
+    # 2. Required dimensions
+    required_dims = {"start_time", "elapsed_forecast_duration"}
+    missing = required_dims - set(ds.dims)
+    assert not missing, f"Missing dimensions: {missing}"
+
+    # 3. Raw 'time' coord must be absent
+    assert "time" not in ds.coords
+
+    # 4. start_time must be analysis_time (= first forecast time - step_length)
+    step_ns = int(
+        ds["elapsed_forecast_duration"].values[0] / np.timedelta64(1, "ns")
+    )
+    for t0 in ds["start_time"].values:
+        first_fcst_abs = t0 + ds["elapsed_forecast_duration"].values[0]
+        expected_t0 = first_fcst_abs - np.timedelta64(step_ns, "ns")
+        assert t0 == expected_t0, (
+            f"start_time {t0} != expected {expected_t0}"
+        )
+
+    # 5. Values must be finite (confirms rescaling happened)
+    state_sample = float(
+        ds["state"].isel(start_time=0, elapsed_forecast_duration=0).mean()
+    )
+    assert np.isfinite(state_sample), "Zarr output contains NaN/Inf"