From a34496650e32a4eaf2804f797f31444564152ad0 Mon Sep 17 00:00:00 2001
From: pragnyanramtha <pragnyanramtha@users.noreply.github.com>
Date: Wed, 13 May 2026 16:39:40 +0000
Subject: [PATCH] Add dataarray size details to merge errors

---
 mllam_data_prep/create_dataset.py | 15 +++++++++++
 tests/test_dataset.py             | 45 +++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py
index 3b7cfb1..7e51992 100644
--- a/mllam_data_prep/create_dataset.py
+++ b/mllam_data_prep/create_dataset.py
@@ -65,6 +65,17 @@ def _check_dataset_attributes(ds, expected_attributes, dataset_name):
         )
 
 
+def _format_dataarray_sizes_by_target(dataarrays_by_target):
+    lines = []
+    for target, dataarrays in dataarrays_by_target.items():
+        lines.append(f"{target}:")
+        for i, da in enumerate(dataarrays, start=1):
+            dims_str = ", ".join(f"{dim}={da.sizes[dim]}" for dim in da.dims)
+            lines.append(f"  {i}. {da.name}: {dims_str}")
+
+    return "\n".join(lines)
+
+
 def _merge_dataarrays_by_target(dataarrays_by_target):
     attrs_to_keep = ["source_dataset"]
     dataarrays = []
@@ -111,9 +122,13 @@ def _summarize(da):
                 return f"{da.name} ({dims})\n{da.coords}"
 
             coord_summaries = "\n".join([_summarize(da) for da in dataarrays])
+            dataarray_size_summary = _format_dataarray_sizes_by_target(
+                dataarrays_by_target
+            )
             raise InvalidConfigException(
                 f"Couldn't merge together the dataarrays for all targets ({', '.join(dataarrays_by_target.keys())}). "
                 "This is likely because the dataarrays have different dimensions or coordinates. "
+                f"Dataarray sizes by target:\n{dataarray_size_summary}\n"
                 f"Dataarray coords:\n{coord_summaries}"
                 "Maybe you need to give the 'feature' dimension a unique name for each target variable?"
             ) from ex
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index e78a93b..f9ba439 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,8 +1,10 @@
 """Tests for the output dataset created by `mllam-data-prep`."""
 import pytest
+import xarray as xr
 import yaml
 
 import mllam_data_prep as mdp
+from mllam_data_prep.create_dataset import _merge_dataarrays_by_target
 
 with open("example.danra.yaml", "r") as file:
     BASE_CONFIG = file.read()
@@ -164,6 +166,49 @@ def update_config(config: str, update: str):
     return modified_config
 
 
+def test_merge_error_includes_dataarray_sizes_by_target():
+    state = xr.DataArray(
+        [[[1], [2]]],
+        dims=("state_feature", "time", "grid_index"),
+        coords={
+            "state_feature": ["z"],
+            "time": [0, 1],
+            "grid_index": [0],
+        },
+        name="z",
+        attrs={
+            "source_dataset": "state_source",
+            "variables_mapping_dim": "state_feature",
+        },
+    )
+    forcing = xr.DataArray(
+        [[[1], [2], [3]]],
+        dims=("forcing_feature", "time", "grid_index"),
+        coords={
+            "forcing_feature": ["toa_radiation"],
+            "time": [0, 1, 2],
+            "grid_index": [0],
+        },
+        name="toa_radiation",
+        attrs={
+            "source_dataset": "forcing_source",
+            "variables_mapping_dim": "forcing_feature",
+        },
+    )
+
+    with pytest.raises(mdp.InvalidConfigException) as exc_info:
+        _merge_dataarrays_by_target({"state": [state], "forcing": [forcing]})
+
+    message = str(exc_info.value)
+    assert "Dataarray sizes by target:" in message
+    assert "state:" in message
+    assert "  1. z: state_feature=1, time=2, grid_index=1" in message
+    assert "forcing:" in message
+    assert (
+        "  1. toa_radiation: forcing_feature=1, time=3, grid_index=1" in message
+    )
+
+
 @pytest.mark.parametrize(
     "base_config, new_inputs_section",
     [