Improvement to Samudra v2: Larger effective resolutions via increased dilation schedule

src/samudra/configs/samudra_om4/model.yaml

@@ -11,7 +11,8 @@ pad: "circular"
 
 unet:
   ch_width: [280, 380, 480, 520]
-  dilation: [1, 2, 4, 8]
+  # Larger dilations make effective receptive fields [3, 17, 33, 33] instead of [3, 5, 9, 17] in Samudra v2.
+  dilation: [1, 8, 16, 16]
   n_layers: [1, 1, 1, 1]
 
   core_block:

src/samudra/viz/core.py

@@ -3967,8 +3967,24 @@ def postprocess_for_plot(
     coords = ds_groundtruth.coords
 
     for key in pred_dict.keys():
+        prediction = pred_dict[key]["ds_prediction"]
+        # Rollouts written by the current writer retain native ``y``/``x``
+        # dimensions plus two-dimensional ``lat``/``lon`` coordinates.  Older
+        # rollouts use ``lat``/``lon`` as the dimensions themselves.  Normalize
+        # the current layout before assigning the ground-truth coordinates;
+        # otherwise the final lon->x rename conflicts with the retained x
+        # coordinate when old and new rollouts are visualized together.
+        if "y" in prediction.dims or "x" in prediction.dims:
+            if not {"y", "x"}.issubset(prediction.dims):
+                raise ValueError(
+                    f"Prediction {key!r} has only one native spatial dimension: "
+                    f"{tuple(prediction.dims)}"
+                )
+            prediction = prediction.drop_vars(["lat", "lon"], errors="ignore")
+            prediction = prediction.rename({"y": "lat", "x": "lon"})
+
         pred_dict[key]["ds_prediction"] = _postprocess_for_plot(
-            pred_dict[key]["ds_prediction"],
+            prediction,
             areacello_values,
             areacello_spherical_values,
             dz,

tests/test_viz_core.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import xarray as xr
+
+from samudra.viz.core import postprocess_for_plot
+
+
+def _groundtruth() -> xr.Dataset:
+    shape = (2, 1, 2, 3)
+    return xr.Dataset(
+        {
+            "thetao": (("time", "lev", "lat", "lon"), np.ones(shape)),
+            "wetmask": (("lev", "lat", "lon"), np.ones(shape[1:], dtype=bool)),
+            "areacello_spherical": (("lat", "lon"), np.ones(shape[2:])),
+        },
+        coords={
+            "time": np.arange(shape[0]),
+            "lev": [2.5],
+            "lat": [-45.0, 45.0],
+            "lon": [0.0, 120.0, 240.0],
+        },
+    )
+
+
+def test_postprocess_for_plot_accepts_old_and_current_rollout_grids():
+    groundtruth = _groundtruth()
+    values = np.ones((2, 1, 2, 3))
+    old = xr.Dataset(
+        {"thetao": (("time", "lev", "lat", "lon"), values)},
+        coords=groundtruth.coords,
+    )
+    current = xr.Dataset(
+        {"thetao": (("time", "lev", "y", "x"), values)},
+        coords={
+            "time": groundtruth.time,
+            "lev": groundtruth.lev,
+            "y": np.arange(2),
+            "x": np.arange(3),
+            "lat": (("y", "x"), np.broadcast_to([[-45.0], [45.0]], (2, 3))),
+            "lon": (("y", "x"), np.broadcast_to([[0.0, 120.0, 240.0]], (2, 3))),
+        },
+    )
+    predictions = {
+        "old": {"ds_prediction": old},
+        "current": {"ds_prediction": current},
+    }
+
+    result, predictions = postprocess_for_plot(
+        groundtruth,
+        xr.DataArray(np.ones((2, 3)), dims=("lat", "lon")),
+        np.array([5.0]),
+        predictions,
+    )
+
+    assert result.thetao.dims == ("time", "lev", "y", "x")
+    for prediction in predictions.values():
+        assert prediction["ds_prediction"].thetao.dims == (
+            "time",
+            "lev",
+            "y",
+            "x",
+        )