Function test_tree_index_sel

xarray/tests/test_nd_point_index.py:49–119

()

@requires_scipy
def test_tree_index_sel() -> None:
    xx, yy = np.meshgrid([1.0, 2.0], [3.0, 4.0])
    ds = xr.Dataset(coords={"xx": (("y", "x"), xx), "yy": (("y", "x"), yy)}).set_xindex(
        ("xx", "yy"), NDPointIndex
    )

    # 1-dimensional labels
    actual = ds.sel(
        xx=xr.Variable("u", [1.1, 1.1, 1.1]),
        yy=xr.Variable("u", [3.1, 3.1, 3.1]),
        method="nearest",
    )
    expected = xr.Dataset(
        coords={"xx": ("u", [1.0, 1.0, 1.0]), "yy": ("u", [3.0, 3.0, 3.0])}
    )
    assert_identical(actual, expected)

    # 2-dimensional labels
    actual = ds.sel(
        xx=xr.Variable(("u", "v"), [[1.1, 1.1, 1.1], [1.9, 1.9, 1.9]]),
        yy=xr.Variable(("u", "v"), [[3.1, 3.1, 3.1], [3.9, 3.9, 3.9]]),
        method="nearest",
    )
    expected = xr.Dataset(
        coords={
            "xx": (("u", "v"), [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]),
            "yy": (("u", "v"), [[3.0, 3.0, 3.0], [4.0, 4.0, 4.0]]),
        },
    )
    assert_identical(actual, expected)

    # all scalar labels
    actual = ds.sel(xx=1.1, yy=3.1, method="nearest")
    expected = xr.Dataset(coords={"xx": 1.0, "yy": 3.0})
    assert_identical(actual, expected)

    # broadcast scalar to label shape and dimensions
    actual = ds.sel(xx=1.1, yy=xr.Variable("u", [3.1, 3.1, 3.1]), method="nearest")
    expected = ds.sel(
        xx=xr.Variable("u", [1.1, 1.1, 1.1]),
        yy=xr.Variable("u", [3.1, 3.1, 3.1]),
        method="nearest",
    )
    assert_identical(actual, expected)

    # broadcast orthogonal 1-dimensional labels
    actual = ds.sel(
        xx=xr.Variable("u", [1.1, 1.1]),
        yy=xr.Variable("v", [3.1, 3.1]),
        method="nearest",
    )
    expected = xr.Dataset(
        coords={
            "xx": (("u", "v"), [[1.0, 1.0], [1.0, 1.0]]),
            "yy": (("u", "v"), [[3.0, 3.0], [3.0, 3.0]]),
        },
    )
    assert_identical(actual, expected)