Feat/neighbors

pyproject.toml

@@ -15,6 +15,7 @@ authors = [
 dependencies = [
     "numpy<2.5",        # 2.5+ stubs need Python 3.12 syntax;
     "xarray",
+    "scipy",
 ]
 
 [dependency-groups]
@@ -38,9 +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
+
 [tool.ruff]
 line-length = 110
 target-version = "py311"

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"))

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
+

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

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()