Function test_apply_missing_dims

xarray/tests/test_computation.py:260–419

()

def test_apply_missing_dims() -> None:
    ## Single arg

    def add_one(a, core_dims, on_missing_core_dim):
        return apply_ufunc(
            lambda x: x + 1,
            a,
            input_core_dims=core_dims,
            output_core_dims=core_dims,
            on_missing_core_dim=on_missing_core_dim,
        )

    array = np.arange(6).reshape(2, 3)
    variable = xr.Variable(["x", "y"], array)
    variable_no_y = xr.Variable(["x", "z"], array)

    ds = xr.Dataset({"x_y": variable, "x_z": variable_no_y})

    # Check the standard stuff works OK
    assert_identical(
        add_one(ds[["x_y"]], core_dims=[["y"]], on_missing_core_dim="raise"),
        ds[["x_y"]] + 1,
    )

    # `raise` — should raise on a missing dim
    with pytest.raises(ValueError):
        add_one(ds, core_dims=[["y"]], on_missing_core_dim="raise")

    # `drop` — should drop the var with the missing dim
    assert_identical(
        add_one(ds, core_dims=[["y"]], on_missing_core_dim="drop"),
        (ds + 1).drop_vars("x_z"),
    )

    # `copy` — should not add one to the missing with `copy`
    copy_result = add_one(ds, core_dims=[["y"]], on_missing_core_dim="copy")
    assert_identical(copy_result["x_y"], (ds + 1)["x_y"])
    assert_identical(copy_result["x_z"], ds["x_z"])

    ## Multiple args

    def sum_add(a, b, core_dims, on_missing_core_dim):
        return apply_ufunc(
            lambda a, b, axis=None: a.sum(axis) + b.sum(axis),
            a,
            b,
            input_core_dims=core_dims,
            on_missing_core_dim=on_missing_core_dim,
        )

    # Check the standard stuff works OK
    assert_identical(
        sum_add(
            ds[["x_y"]],
            ds[["x_y"]],
            core_dims=[["x", "y"], ["x", "y"]],
            on_missing_core_dim="raise",
        ),