Function test_integrate

xarray/tests/test_dataset.py:8040–8089

(dask)

@pytest.mark.parametrize("dask", [True, False])
def test_integrate(dask) -> 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,
            "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.integrate("x")
    # coordinate that contains x should be dropped.
    expected_x = xr.DataArray(
        trapezoid(da.compute(), da["x"], axis=0),
        dims=["y"],
        coords={k: v for k, v in da.coords.items() if "x" not in v.dims},
    )
    assert_allclose(expected_x, actual.compute())
    assert_equal(ds["var"].integrate("x"), ds.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.integrate("y")
    expected_y = xr.DataArray(
        trapezoid(da, da["y"], axis=1),
        dims=["x"],
        coords={k: v for k, v in da.coords.items() if "y" not in v.dims},
    )
    assert_allclose(expected_y, actual.compute())
    assert_equal(actual, ds.integrate("y")["var"])
    assert_equal(ds["var"].integrate("y"), ds.integrate("y")["var"])

    # along x and y
    actual = da.integrate(("y", "x"))
    assert actual.ndim == 0

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


@requires_scipy