diff --git a/src/samudra/configs/samudra_om4/model.yaml b/src/samudra/configs/samudra_om4/model.yaml index 0bcb2c0f..b408ca66 100644 --- a/src/samudra/configs/samudra_om4/model.yaml +++ b/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: diff --git a/src/samudra/viz/core.py b/src/samudra/viz/core.py index ead349cf..0b858460 100644 --- a/src/samudra/viz/core.py +++ b/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, diff --git a/tests/test_viz_core.py b/tests/test_viz_core.py new file mode 100644 index 00000000..70934899 --- /dev/null +++ b/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", + )