Function test_cumulative_integrate

xarray/tests/test_dataset.py:8094–8151

(dask)

@requires_scipy
@pytest.mark.parametrize("dask", [True, False])
def test_cumulative_integrate(dask) -> None:
    rs = np.random.default_rng(43)
    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,
            "x2": (("x",), rs.random(8)),
            "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.cumulative_integrate("x")

    from scipy.integrate import cumulative_trapezoid

    expected_x = xr.DataArray(
        cumulative_trapezoid(da.compute(), da["x"], axis=0, initial=0.0),  # type: ignore[call-overload,unused-ignore]
        dims=["x", "y"],
        coords=da.coords,
    )
    assert_allclose(expected_x, actual.compute())
    assert_equal(
        ds["var"].cumulative_integrate("x"),
        ds.cumulative_integrate("x")["var"],
    )

    # make sure result is also a dask array (if the source is dask array)
    assert isinstance(actual.data, type(da.data))

    # along y
    actual = da.cumulative_integrate("y")
    expected_y = xr.DataArray(
        cumulative_trapezoid(da, da["y"], axis=1, initial=0.0),  # type: ignore[call-overload,unused-ignore]
        dims=["x", "y"],
        coords=da.coords,
    )
    assert_allclose(expected_y, actual.compute())
    assert_equal(actual, ds.cumulative_integrate("y")["var"])
    assert_equal(
        ds["var"].cumulative_integrate("y"),
        ds.cumulative_integrate("y")["var"],
    )

    # along x and y
    actual = da.cumulative_integrate(("y", "x"))
    assert actual.ndim == 2

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