def test_concat_2d__variables(ds: Dataset): # type: ignore[no-untyped-def]
# Select data variables expected to remain after extract
data_vars = [v for v in ds if ds[v].dims[0] == "dim-0-0"]
# Sum the number of columns for the kept variables
n_cols = _col_shape_sum(ds[data_vars])
res = concat_2d(ds, dims=("dim-0-0", "dim-1"))
# Ensure that 'dim-0-1' is lost and that the result
# contains the expected number of columns
assert res.dims == ("dim-0-0", "dim-1")
assert res.shape == (ds.attrs["n"], n_cols)
def test_concat_2d__values(n: int) -> None:
x, y = np.arange(n), np.arange(n * n).reshape(n, n)
z = np.copy(y)
ds = xr.Dataset(
dict(
x=(("dim0"), x),
y=(("dim0", "dim1"), y),
z=(("dim1", "dim2"), z),
))
actual = concat_2d(ds, dims=("dim0", "dim1"))
expected = np.concatenate([x.reshape(-1, 1), y], axis=1)
np.testing.assert_equal(actual, expected)
def test_concat_2d__raise_on_gte_3d():
ds = xr.Dataset(dict(x=(("dim0", "dim1", "dim2"), np.empty((1, 1, 1)))))
with pytest.raises(ValueError,
match="All variables must have <= 2 dimensions"):
concat_2d(ds, dims=("dim0", "dim1"))