Allow creating DataArrays with nD coordinate variables (#8126)

* Allow creating DataArrays with nD coordinate variables Closes #2233 Closes #8106 * more test more test# make_aggs.bash * Fix test * Apply suggestions from code review Co-authored-by: Michael Niklas <[email protected]> * Update test --------- Co-authored-by: Michael Niklas <[email protected]> Co-authored-by: Maximilian Roos <[email protected]>
pydata · Sep 22, 2023 · cdf0726 · cdf0726
1 parent 3ace2fb
commit cdf0726
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 11 deletions.
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -128,21 +128,14 @@ def _check_coords_dims(shape, coords, dims):
                 f"dimensions {dims}"
             )
 
-        for d, s in zip(v.dims, v.shape):
+        for d, s in v.sizes.items():
             if s != sizes[d]:
                 raise ValueError(
                     f"conflicting sizes for dimension {d!r}: "
                     f"length {sizes[d]} on the data but length {s} on "
                     f"coordinate {k!r}"
                 )
 
-        if k in sizes and v.shape != (sizes[k],):
-            raise ValueError(
-                f"coordinate {k!r} is a DataArray dimension, but "
-                f"it has shape {v.shape!r} rather than expected shape {sizes[k]!r} "
-                "matching the dimension size"
-            )
-
 
 def _infer_coords_and_dims(
     shape, coords, dims

diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
@@ -38,6 +38,7 @@
 from xarray.core.indexes import Index, PandasIndex, filter_indexes_from_coords
 from xarray.core.types import QueryEngineOptions, QueryParserOptions
 from xarray.core.utils import is_scalar
+from xarray.testing import _assert_internal_invariants
 from xarray.tests import (
     InaccessibleArray,
     ReturnItem,
@@ -415,9 +416,6 @@ def test_constructor_invalid(self) -> None:
         with pytest.raises(ValueError, match=r"conflicting MultiIndex"):
             DataArray(np.random.rand(4, 4), [("x", self.mindex), ("level_1", range(4))])
 
-        with pytest.raises(ValueError, match=r"matching the dimension size"):
-            DataArray(data, coords={"x": 0}, dims=["x", "y"])
-
     def test_constructor_from_self_described(self) -> None:
         data = [[-0.1, 21], [0, 2]]
         expected = DataArray(
@@ -7112,3 +7110,35 @@ def test_error_on_ellipsis_without_list(self) -> None:
         da = DataArray([[1, 2], [1, 2]], dims=("x", "y"))
         with pytest.raises(ValueError):
             da.stack(flat=...)  # type: ignore
+
+
+def test_nD_coord_dataarray() -> None:
+    # should succeed
+    da = DataArray(
+        np.ones((2, 4)),
+        dims=("x", "y"),
+        coords={
+            "x": (("x", "y"), np.arange(8).reshape((2, 4))),
+            "y": ("y", np.arange(4)),
+        },
+    )
+    _assert_internal_invariants(da, check_default_indexes=True)
+
+    da2 = DataArray(np.ones(4), dims=("y"), coords={"y": ("y", np.arange(4))})
+    da3 = DataArray(np.ones(4), dims=("z"))
+
+    _, actual = xr.align(da, da2)
+    assert_identical(da2, actual)
+
+    expected = da.drop_vars("x")
+    _, actual = xr.broadcast(da, da2)
+    assert_identical(expected, actual)
+
+    actual, _ = xr.broadcast(da, da3)
+    expected = da.expand_dims(z=4, axis=-1)
+    assert_identical(actual, expected)
+
+    da4 = DataArray(np.ones((2, 4)), coords={"x": 0}, dims=["x", "y"])
+    _assert_internal_invariants(da4, check_default_indexes=True)
+    assert "x" not in da4.xindexes
+    assert "x" in da4.coords