test_nearest_point_to.py

from typing import Any

import numpy as np
import pytest
import xarray as xr

from movement.roi.base import BaseRegionOfInterest
from movement.roi.line import LineOfInterest


@pytest.fixture
def points_of_interest() -> dict[str, np.ndarray]:
    return xr.DataArray(
        np.array(
            [
                [-0.5, 0.50],
                [0.00, 0.50],
                [0.40, 0.45],
                [2.00, 1.00],
                [0.40, 0.75],
                [0.95, 0.90],
                [0.80, 0.76],
            ]
        ),
        dims=["time", "space"],
        coords={"space": ["x", "y"]},
    )


@pytest.fixture()
def unit_line_in_x() -> LineOfInterest:
    return LineOfInterest([[0.0, 0.0], [1.0, 0.0]])


@pytest.mark.parametrize(
    ["region", "other_fn_args", "expected_output"],
    [
        pytest.param(
            "unit_square",
            {"boundary": True},
            np.array(
                [
                    [0.00, 0.50],
                    [0.00, 0.50],
                    [0.00, 0.45],
                    [1.00, 1.00],
                    [0.40, 1.00],
                    [1.00, 0.90],
                    [1.00, 0.76],
                ]
            ),
            id="Unit square, boundary only",
        ),
        pytest.param(
            "unit_square",
            {},
            np.array(
                [
                    [0.00, 0.50],
                    [0.00, 0.50],
                    [0.40, 0.45],
                    [1.00, 1.00],
                    [0.40, 0.75],
                    [0.95, 0.90],
                    [0.80, 0.76],
                ]
            ),
            id="Unit square, whole region",
        ),
        pytest.param(
            "unit_square_with_hole",
            {"boundary": True},
            np.array(
                [
                    [0.00, 0.50],
                    [0.00, 0.50],
                    [0.25, 0.45],
                    [1.00, 1.00],
                    [0.40, 0.75],
                    [1.00, 0.90],
                    [0.75, 0.75],
                ]
            ),
            id="Unit square w/ hole, boundary only",
        ),
        pytest.param(
            "unit_square_with_hole",
            {},
            np.array(
                [
                    [0.00, 0.50],
                    [0.00, 0.50],
                    [0.25, 0.45],
                    [1.00, 1.00],
                    [0.40, 0.75],
                    [0.95, 0.90],
                    [0.80, 0.76],
                ]
            ),
            id="Unit square w/ hole, whole region",
        ),
        pytest.param(
            "unit_line_in_x",
            {},
            np.array(
                [
                    [0.00, 0.00],
                    [0.00, 0.00],
                    [0.40, 0.00],
                    [1.00, 0.00],
                    [0.40, 0.00],
                    [0.95, 0.00],
                    [0.80, 0.00],
                ]
            ),
            id="Line, whole region",
        ),
        pytest.param(
            "unit_line_in_x",
            {"boundary": True},
            np.array(
                [
                    [0.00, 0.00],
                    [0.00, 0.00],
                    [0.00, 0.00],
                    [1.00, 0.00],
                    [0.00, 0.00],
                    [1.00, 0.00],
                    [1.00, 0.00],
                ]
            ),
            id="Line, boundary only",
        ),
    ],
)
def test_nearest_point_to(
    region: BaseRegionOfInterest,
    points_of_interest: xr.DataArray,
    other_fn_args: dict[str, Any],
    expected_output: xr.DataArray,
    request,
) -> None:
    if isinstance(region, str):
        region = request.getfixturevalue(region)
    if isinstance(points_of_interest, str):
        points_of_interest = request.getfixturevalue(points_of_interest)
    if isinstance(expected_output, str):
        expected_output = request.get(expected_output)
    elif isinstance(expected_output, np.ndarray):
        expected_output = xr.DataArray(
            expected_output,
            dims=["time", "nearest point"],
        )

    points_of_interest = points_of_interest
    nearest_points = region.nearest_point_to(
        points_of_interest, **other_fn_args
    )

    xr.testing.assert_allclose(nearest_points, expected_output)


@pytest.mark.parametrize(
    ["region", "position", "fn_kwargs", "possible_nearest_points"],
    [
        pytest.param(
            "unit_square",
            [0.5, 0.5],
            {"boundary": True},
            [
                np.array([0.0, 0.5]),
                np.array([0.5, 0.0]),
                np.array([1.0, 0.5]),
                np.array([0.5, 1.0]),
            ],
            id="Centre of the unit square",
        ),
        pytest.param(
            "unit_line_in_x",
            [0.5, 0.0],
            {"boundary": True},
            [
                np.array([0.0, 0.0]),
                np.array([1.0, 0.0]),
            ],
            id="Boundary of a line",
        ),
    ],
)
def test_tie_breaks(
    region: BaseRegionOfInterest,
    position: np.ndarray,
    fn_kwargs: dict[str, Any],
    possible_nearest_points: list[np.ndarray],
    request,
) -> None:
    """Check behaviour when points are tied for nearest.

    This can only occur when we have a Polygonal region, or a multi-line 1D
    region. In this case, there may be multiple points in the region of
    interest that are tied for closest. ``shapely`` does not actually document
    how it breaks ties here, but we can at least check that it identifies one
    of the possible correct points.
    """
    if isinstance(region, str):
        region = request.getfixturevalue(region)
    if not isinstance(position, np.ndarray | xr.DataArray):
        position = np.array(position)

    nearest_point_found = region.nearest_point_to(position, **fn_kwargs)

    sq_dist_to_nearest_pt = np.sum((nearest_point_found - position) ** 2)

    n_matches = 0
    for possibility in possible_nearest_points:
        # All possibilities should be approximately the same distance away
        # from the position
        assert np.isclose(
            np.sum((possibility - position) ** 2), sq_dist_to_nearest_pt
        )
        # We should match at least one possibility,
        # track to see if we do.
        if np.isclose(nearest_point_found, possibility).all():
            n_matches += 1
    assert n_matches == 1