Function get_example_data

xarray/tests/test_interp.py:39–70

(case: int)

def get_example_data(case: int) -> xr.DataArray:
    if case == 0:
        # 2D
        x = np.linspace(0, 1, 100)
        y = np.linspace(0, 0.1, 30)
        return xr.DataArray(
            np.sin(x[:, np.newaxis]) * np.cos(y),
            dims=["x", "y"],
            coords={"x": x, "y": y, "x2": ("x", x**2)},
        )
    elif case == 1:
        # 2D chunked single dim
        return get_example_data(0).chunk({"y": 3})
    elif case == 2:
        # 2D chunked both dims
        return get_example_data(0).chunk({"x": 25, "y": 3})
    elif case == 3:
        # 3D
        x = np.linspace(0, 1, 100)
        y = np.linspace(0, 0.1, 30)
        z = np.linspace(0.1, 0.2, 10)
        return xr.DataArray(
            np.sin(x[:, np.newaxis, np.newaxis]) * np.cos(y[:, np.newaxis]) * z,
            dims=["x", "y", "z"],
            coords={"x": x, "y": y, "x2": ("x", x**2), "z": z},
        )
    elif case == 4:
        # 3D chunked single dim
        # chunksize=5 lets us check whether we rechunk to 1 with quintic
        return get_example_data(3).chunk({"z": 5})
    else:
        raise ValueError("case must be 1-4")


@pytest.fixture