From 82b6bf11290981d395367397a2b13112e62371ad Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 17:34:24 +0000 Subject: [PATCH 1/3] add matchup and neighbors features and tests --- src/floatmatcher/matchup.py | 22 ++++++++++++ src/floatmatcher/neighbors.py | 42 +++++++++++++++++++++++ tests/test_matchup.py | 15 +++++++++ tests/test_neighbors.py | 63 +++++++++++++++++++++++++++++++++++ 4 files changed, 142 insertions(+) create mode 100644 src/floatmatcher/matchup.py create mode 100644 src/floatmatcher/neighbors.py create mode 100644 tests/test_matchup.py create mode 100644 tests/test_neighbors.py diff --git a/src/floatmatcher/matchup.py b/src/floatmatcher/matchup.py new file mode 100644 index 0000000..987b613 --- /dev/null +++ b/src/floatmatcher/matchup.py @@ -0,0 +1,22 @@ +# nearest.py: nearest-neighbor matchup via separated spatial/temporal KDTrees. +# +# The spatial half is prepared ONCE (prepare) and reused on every temporal +# packet (match_packet), because the grid geometry is identical across packets. +# The packet loop itself lives in the orchestrator; this module provides the +# two halves and a single-pass `match` for the non-batched case. + +import numpy as np + +class NearestNeighbor(): + """Nearest-neighbor matchup method""" + + def __init__(self, max_dist_km: int = 25, + max_time: np.timedelta64 = np.timedelta64(1, "D"), + k_nearest : int = 1) -> None : + self.max_dist_km = max_dist_km + self.max_time = max_time + self.k_nearest = k_nearest + + @property + def max_time_seconds(self) -> float: + return float(self.max_time / np.timedelta64(1, "s")) diff --git a/src/floatmatcher/neighbors.py b/src/floatmatcher/neighbors.py new file mode 100644 index 0000000..df0d810 --- /dev/null +++ b/src/floatmatcher/neighbors.py @@ -0,0 +1,42 @@ +# neighbors.py: spatial and temporal KDTree lookups over a reference grid. +# +# Spatial and temporal are SEPARATED because they now have different lifetimes: +# - SpatialIndex is built ONCE (grid geometry is identical on every packet); +# - TemporalIndex is built PER packet (only the time axis changes). +# On a regular grid the spatial positions repeat at every time step, so +# "closest in space" and "closest in time" are independent questions. + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial import cKDTree + +from .pointset import PointSet +from .constants import REF_TIME, TIME_UNIT + + +def _to_seconds(times: NDArray[np.datetime64]) -> NDArray[np.float64]: + """Convert datetime64 to floating-point seconds since a fixed epoch. + + Working in a common float unit lets the 1D KDTree measure time distance, + and returning *seconds* makes the max_time_seconds constraint directly comparable. + """ + delta = np.asarray(times, dtype=f"datetime64[{TIME_UNIT}]") - REF_TIME + seconds: NDArray[np.float64] = delta / np.timedelta64(1, "s") + return seconds + + +def spatial_nearest(grid_xyz: NDArray[np.float64], points: PointSet, + k: int = 1) -> tuple[NDArray[np.float64], NDArray[np.int64]]: + dist: NDArray[np.float64] + idx: NDArray[np.int64] + dist, idx = cKDTree(grid_xyz).query(points.xyz, k=k) + return dist, idx + +def temporal_nearest(grid_times: NDArray[np.datetime64], points: PointSet, + k: int = 1) -> tuple[NDArray[np.float64], NDArray[np.int64]]: + time_delta: NDArray[np.float64] + idx: NDArray[np.int64] + grid_tree = cKDTree(_to_seconds(grid_times)[:, None]) + time_delta, idx = grid_tree.query(_to_seconds(points.time)[:, None], k=k) + return time_delta, idx + diff --git a/tests/test_matchup.py b/tests/test_matchup.py new file mode 100644 index 0000000..8911bb3 --- /dev/null +++ b/tests/test_matchup.py @@ -0,0 +1,15 @@ +# tests/test_matchup.py + +import numpy as np + +from floatmatcher.matchup import NearestNeighbor + + +def test_default_tolerance_is_one_day_in_seconds(): + assert NearestNeighbor().max_time_seconds == 86400.0 + + +def test_max_time_accepts_any_unit(): + assert NearestNeighbor(max_time=np.timedelta64(3, "h")).max_time_seconds == 10800.0 + assert NearestNeighbor(max_time=np.timedelta64(90, "m")).max_time_seconds == 5400.0 + assert NearestNeighbor(max_time=np.timedelta64(2, "D")).max_time_seconds == 172800.0 diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py new file mode 100644 index 0000000..b48a4d7 --- /dev/null +++ b/tests/test_neighbors.py @@ -0,0 +1,63 @@ +# tests/test_neighbors.py + +import numpy as np +import numpy.testing as npt +import pytest + +from floatmatcher.geo import lonlat_to_xyz +from floatmatcher.neighbors import spatial_nearest, temporal_nearest +from floatmatcher.pointset import PointSet + +grid_lon = [0.0, 10.0, 0.0, 10.0] +grid_lat = [0.0, 0.0, 10.0, 10.0] +grid_xyz = lonlat_to_xyz(grid_lon, grid_lat) +grid_times = np.array(["2015-01-01", "2015-01-02", "2015-01-03", + "2015-01-04", "2015-01-05"], dtype="datetime64[ns]") + + +@pytest.fixture +def points_on_nodes(): + """Three points sitting exactly on grid nodes 2, 0 and 3.""" + return PointSet( + lon=[0.0, 0.0, 10.0], + lat=[10.0, 0.0, 10.0], + time=np.array(["2015-01-03T03:00", "2015-01-01T00:00", "2015-01-04T23:00"], + dtype="datetime64[ns]"), + ) + + +def test_spatial_nearest_finds_the_node_under_each_point(points_on_nodes): + dist, idx = spatial_nearest(grid_xyz, points_on_nodes) + npt.assert_array_equal(idx, [2, 0, 3]) + npt.assert_allclose(dist, 0.0, atol=1e-9) + + +def test_spatial_nearest_distance_is_in_kilometres(): + points = PointSet(lon=[1.0], lat=[0.0], + time=np.array(["2015-01-01"], dtype="datetime64[ns]")) + dist, idx = spatial_nearest(grid_xyz, points) + expected = np.linalg.norm(points.xyz[0] - grid_xyz[0]) + npt.assert_array_equal(idx, [0]) + npt.assert_allclose(dist, [expected]) + + +def test_spatial_nearest_k_returns_sorted_neighbours(points_on_nodes): + dist, idx = spatial_nearest(grid_xyz, points_on_nodes, k=3) + assert dist.shape == (3, 3) + assert idx.shape == (3, 3) + npt.assert_array_equal(idx[:, 0], [2, 0, 3]) + assert (np.diff(dist, axis=1) >= 0).all() + + +def test_temporal_nearest_delta_is_in_seconds(points_on_nodes): + time_delta, idx = temporal_nearest(grid_times, points_on_nodes) + npt.assert_array_equal(idx, [2, 0, 4]) + npt.assert_allclose(time_delta, [10800.0, 0.0, 3600.0]) + + +def test_temporal_nearest_k_returns_sorted_neighbours(points_on_nodes): + time_delta, idx = temporal_nearest(grid_times, points_on_nodes, k=2) + assert idx.shape == (3, 2) + npt.assert_array_equal(idx[:, 0], [2, 0, 4]) + npt.assert_allclose(time_delta[:, 0], [10800.0, 0.0, 3600.0]) + assert (np.diff(time_delta, axis=1) >= 0).all() From 09cbef2e2e2a7d26211563b9fd54e2c85e9ee119 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 17:35:13 +0000 Subject: [PATCH 2/3] edit to add scipy in overrides + dependencies --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 9d4a8d2..31bd535 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ authors = [ dependencies = [ "numpy<2.5", # 2.5+ stubs need Python 3.12 syntax; "xarray", + "scipy", ] [dependency-groups] @@ -39,6 +40,10 @@ warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true +[[tool.mypy.overrides]] +module = ["scipy.*"] +ignore_missing_imports = true + exclude = ["^tests/"] [tool.ruff] From d54a0771bb24b0a2110c275b213148bb5866d828 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 17:47:11 +0000 Subject: [PATCH 3/3] mis typed mypy fix --- pyproject.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 31bd535..514cec2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,13 +39,12 @@ python_version = "3.11" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true +exclude = ["^tests/"] [[tool.mypy.overrides]] module = ["scipy.*"] ignore_missing_imports = true -exclude = ["^tests/"] - [tool.ruff] line-length = 110 target-version = "py311"