Function test_differentiate

xarray/tests/test_dataset.py:7913–7957

(dask, edge_order)

@pytest.mark.parametrize("dask", [True, False])
@pytest.mark.parametrize("edge_order", [1, 2])
def test_differentiate(dask, edge_order) -> None:
    rs = np.random.default_rng(42)
    coord = [0.2, 0.35, 0.4, 0.6, 0.7, 0.75, 0.76, 0.8]

    da = xr.DataArray(
        rs.random((8, 6)),
        dims=["x", "y"],
        coords={"x": coord, "z": 3, "x2d": (("x", "y"), rs.random((8, 6)))},
    )
    if dask and has_dask:
        da = da.chunk({"x": 4})

    ds = xr.Dataset({"var": da})

    # along x
    actual = da.differentiate("x", edge_order)
    expected_x = xr.DataArray(
        np.gradient(da, da["x"], axis=0, edge_order=edge_order),
        dims=da.dims,
        coords=da.coords,
    )
    assert_equal(expected_x, actual)
    assert_equal(
        ds["var"].differentiate("x", edge_order=edge_order),
        ds.differentiate("x", edge_order=edge_order)["var"],
    )
    # coordinate should not change
    assert_equal(da["x"], actual["x"])

    # along y
    actual = da.differentiate("y", edge_order)
    expected_y = xr.DataArray(
        np.gradient(da, da["y"], axis=1, edge_order=edge_order),
        dims=da.dims,
        coords=da.coords,
    )
    assert_equal(expected_y, actual)
    assert_equal(actual, ds.differentiate("y", edge_order=edge_order)["var"])
    assert_equal(
        ds["var"].differentiate("y", edge_order=edge_order),
        ds.differentiate("y", edge_order=edge_order)["var"],
    )

    with pytest.raises(ValueError):
        da.differentiate("x2d")


@pytest.mark.parametrize("dask", [True, False])