MCPcopy Index your code
hub / github.com/pydata/xarray / test_expand_dims

Method test_expand_dims

xarray/tests/test_dataarray.py:2230–2286  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

2228 array.expand_dims({"d": 4}, e=4)
2229
2230 def test_expand_dims(self) -> None:
2231 array = DataArray(
2232 np.random.randn(3, 4),
2233 dims=["x", "dim_0"],
2234 coords={"x": np.linspace(0.0, 1.0, 3)},
2235 attrs={"key": "entry"},
2236 )
2237 # pass only dim label
2238 actual = array.expand_dims(dim="y")
2239 expected = DataArray(
2240 np.expand_dims(array.values, 0),
2241 dims=["y", "x", "dim_0"],
2242 coords={"x": np.linspace(0.0, 1.0, 3)},
2243 attrs={"key": "entry"},
2244 )
2245 assert_identical(expected, actual)
2246 roundtripped = actual.squeeze("y", drop=True)
2247 assert_identical(array, roundtripped)
2248
2249 # pass multiple dims
2250 actual = array.expand_dims(dim=["y", "z"])
2251 expected = DataArray(
2252 np.expand_dims(np.expand_dims(array.values, 0), 0),
2253 dims=["y", "z", "x", "dim_0"],
2254 coords={"x": np.linspace(0.0, 1.0, 3)},
2255 attrs={"key": "entry"},
2256 )
2257 assert_identical(expected, actual)
2258 roundtripped = actual.squeeze(["y", "z"], drop=True)
2259 assert_identical(array, roundtripped)
2260
2261 # pass multiple dims and axis. Axis is out of order
2262 actual = array.expand_dims(dim=["z", "y"], axis=[2, 1])
2263 expected = DataArray(
2264 np.expand_dims(np.expand_dims(array.values, 1), 2),
2265 dims=["x", "y", "z", "dim_0"],
2266 coords={"x": np.linspace(0.0, 1.0, 3)},
2267 attrs={"key": "entry"},
2268 )
2269 assert_identical(expected, actual)
2270 # make sure the attrs are tracked
2271 assert actual.attrs["key"] == "entry"
2272 roundtripped = actual.squeeze(["z", "y"], drop=True)
2273 assert_identical(array, roundtripped)
2274
2275 # Negative axis and they are out of order
2276 actual = array.expand_dims(dim=["y", "z"], axis=[-1, -2])
2277 expected = DataArray(
2278 np.expand_dims(np.expand_dims(array.values, -1), -1),
2279 dims=["x", "dim_0", "z", "y"],
2280 coords={"x": np.linspace(0.0, 1.0, 3)},
2281 attrs={"key": "entry"},
2282 )
2283 assert_identical(expected, actual)
2284 assert actual.attrs["key"] == "entry"
2285 roundtripped = actual.squeeze(["y", "z"], drop=True)
2286 assert_identical(array, roundtripped)
2287

Callers

nothing calls this directly

Calls 6

expand_dimsMethod · 0.95
DataArrayClass · 0.90
assert_identicalFunction · 0.90
linspaceMethod · 0.80
expand_dimsMethod · 0.45
squeezeMethod · 0.45

Tested by

no test coverage detected