Skip to content
forked from pydata/xarray

Commit

Permalink
Propagate MultiIndex variables in broadcast
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Apr 12, 2022
1 parent c12cb3d commit 76e7747
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 2 deletions.
4 changes: 2 additions & 2 deletions xarray/core/alignment.py
Original file line number Diff line number Diff line change
Expand Up @@ -927,8 +927,8 @@ def _get_broadcast_dims_map_common_coords(args, exclude):
for dim in arg.dims:
if dim not in common_coords and dim not in exclude:
dims_map[dim] = arg.sizes[dim]
if dim in arg.coords:
common_coords[dim] = arg.coords[dim].variable
if dim in arg.xindexes:
common_coords.update(arg.xindexes[dim].create_variables())

return dims_map, common_coords

Expand Down
12 changes: 12 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -2331,6 +2331,18 @@ def test_broadcast_misaligned(self):
assert_identical(expected_x2, x2)
assert_identical(expected_y2, y2)

def test_broadcast_multi_index(self):
# GH6430
ds = Dataset(
{"foo": (("x", "y", "z"), np.ones((3, 4, 2)))},
{"x": ["a", "b", "c"], "y": [1, 2, 3, 4]},
)
stacked = ds.stack(space=["x", "y"])
broadcasted, _ = broadcast(stacked, stacked.space)

assert broadcasted.xindexes["x"] is broadcasted.xindexes["space"]
assert broadcasted.xindexes["y"] is broadcasted.xindexes["space"]

def test_variable_indexing(self):
data = create_test_data()
v = data["var1"]
Expand Down

0 comments on commit 76e7747

Please sign in to comment.