Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow creating DataArrays with nD coordinate variables #8126

Merged
merged 10 commits into from
Sep 22, 2023
9 changes: 1 addition & 8 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,21 +119,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
Expand Down
32 changes: 29 additions & 3 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,9 +409,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(
Expand Down Expand Up @@ -7085,3 +7082,32 @@ 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:
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@benbovy can you review these tests please

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gentle ping @benbovy

# 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)),
},
)
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 "x" not in da4.xindexes
assert "x" in da4.coords
Loading