MCPcopy
hub / github.com/pydata/xarray / test_set_dims_without_broadcast

Method test_set_dims_without_broadcast

xarray/tests/test_variable.py:1669–1735  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

1667 assert_identical(actual, expected)
1668
1669 def test_set_dims_without_broadcast(self):
1670 class ArrayWithoutBroadcastTo(NDArrayMixin, indexing.ExplicitlyIndexed):
1671 def __init__(self, array):
1672 self.array = array
1673
1674 # Broadcasting with __getitem__ is "easier" to implement
1675 # especially for dims of 1
1676 def __getitem__(self, key):
1677 return self.array[key]
1678
1679 def __array_function__(self, *args, **kwargs):
1680 raise NotImplementedError(
1681 "Not we don't want to use broadcast_to here "
1682 "https://github.com/pydata/xarray/issues/9462"
1683 )
1684
1685 arr = ArrayWithoutBroadcastTo(np.zeros((3, 4)))
1686 # We should be able to add a new axis without broadcasting
1687 assert arr[np.newaxis, :, :].shape == (1, 3, 4)
1688 with pytest.raises(NotImplementedError):
1689 np.broadcast_to(arr, (1, 3, 4))
1690
1691 v = Variable(["x", "y"], arr)
1692 v_expanded = v.set_dims(["z", "x", "y"])
1693 assert v_expanded.dims == ("z", "x", "y")
1694 assert v_expanded.shape == (1, 3, 4)
1695
1696 v_expanded = v.set_dims(["x", "z", "y"])
1697 assert v_expanded.dims == ("x", "z", "y")
1698 assert v_expanded.shape == (3, 1, 4)
1699
1700 v_expanded = v.set_dims(["x", "y", "z"])
1701 assert v_expanded.dims == ("x", "y", "z")
1702 assert v_expanded.shape == (3, 4, 1)
1703
1704 # Explicitly asking for a shape of 1 triggers a different
1705 # codepath in set_dims
1706 # https://github.com/pydata/xarray/issues/9462
1707 v_expanded = v.set_dims(["z", "x", "y"], shape=(1, 3, 4))
1708 assert v_expanded.dims == ("z", "x", "y")
1709 assert v_expanded.shape == (1, 3, 4)
1710
1711 v_expanded = v.set_dims(["x", "z", "y"], shape=(3, 1, 4))
1712 assert v_expanded.dims == ("x", "z", "y")
1713 assert v_expanded.shape == (3, 1, 4)
1714
1715 v_expanded = v.set_dims(["x", "y", "z"], shape=(3, 4, 1))
1716 assert v_expanded.dims == ("x", "y", "z")
1717 assert v_expanded.shape == (3, 4, 1)
1718
1719 v_expanded = v.set_dims({"z": 1, "x": 3, "y": 4})
1720 assert v_expanded.dims == ("z", "x", "y")
1721 assert v_expanded.shape == (1, 3, 4)
1722
1723 v_expanded = v.set_dims({"x": 3, "z": 1, "y": 4})
1724 assert v_expanded.dims == ("x", "z", "y")
1725 assert v_expanded.shape == (3, 1, 4)
1726

Callers

nothing calls this directly

Calls 4

set_dimsMethod · 0.95
VariableClass · 0.90
broadcast_toMethod · 0.80

Tested by

no test coverage detected