From 9a786e581fbfe3d7145754e5f6acc9874aa4d4d5 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 1 Sep 2026 11:16:37 +0000 Subject: [PATCH 01/12] add constants and exceptions --- pyproject.toml | 4 +++- src/floatmatcher/constants.py | 16 ++++++++++++++++ src/floatmatcher/exceptions.py | 13 +++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 src/floatmatcher/constants.py create mode 100644 src/floatmatcher/exceptions.py diff --git a/pyproject.toml b/pyproject.toml index 5a57127..df10d22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ authors = [ {name = "Clement WEBER", email = "clement.weber@pokapok.org"}, ] dependencies = [ - "numpy", + "numpy<2.5", # 2.5+ stubs need Python 3.12 syntax; "xarray", ] @@ -42,4 +42,6 @@ disallow_untyped_defs = true [tool.ruff] line-length = 88 target-version = "py311" + +[tool.ruff.lint] select = ["E", "F"] diff --git a/src/floatmatcher/constants.py b/src/floatmatcher/constants.py new file mode 100644 index 0000000..42f3b2b --- /dev/null +++ b/src/floatmatcher/constants.py @@ -0,0 +1,16 @@ +# constants.py: project-wide constants and conventions, gathered in one place. + +import numpy as np + +# --- geography --- +EARTH_RADIUS_KM: float = 6371.0 +"""Mean Earth radius (km), used to project lon/lat onto a sphere in geo.py.""" + +# --- time --- +TIME_UNIT: str = "ns" +"""Canonical datetime64 resolution enforced across the library (points and grids).""" + +REF_TIME: np.datetime64 = np.datetime64("1970-01-01", "ns") +"""Reference epoch: datetime64 values are turned into floating-point days from +here for the 1D temporal KDTree. Any fixed epoch works; what matters is that +points and grid use the SAME one, so their difference is correct.""" diff --git a/src/floatmatcher/exceptions.py b/src/floatmatcher/exceptions.py new file mode 100644 index 0000000..de9deba --- /dev/null +++ b/src/floatmatcher/exceptions.py @@ -0,0 +1,13 @@ +# floatmatcher/exceptions.py +# All library-specific errors live here, under a single base class so callers +# can catch everything from the library with `except FloatMatcherError`. + + +class FloatMatcherError(Exception): + """Base class for all FloatMatcher errors.""" + + +class ProfileFormatError(FloatMatcherError): + """Raised when point data cannot be extracted from the given source + (missing coordinate, unrecognized structure).""" + From f82b7ab41cbd7b6ba1c01a45c9b723409c376981 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 1 Sep 2026 16:22:20 +0000 Subject: [PATCH 02/12] add feature geo and associated tests --- src/floatmatcher/geo.py | 51 ++++++++++++++++++++++++++++++++++++ tests/test_geo.py | 58 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 src/floatmatcher/geo.py create mode 100644 tests/test_geo.py diff --git a/src/floatmatcher/geo.py b/src/floatmatcher/geo.py new file mode 100644 index 0000000..7666a39 --- /dev/null +++ b/src/floatmatcher/geo.py @@ -0,0 +1,51 @@ +# geo.py : all geography transformations + +# imports +import numpy as np +from numpy.typing import ArrayLike, NDArray +from .constants import EARTH_RADIUS_KM + +# functions +#_______ + +def lonlat_to_xyz(lon: ArrayLike, lat: ArrayLike) -> NDArray[np.float64]: + """Convert lon/lat to xyz on the shpere""" + lon_rad = np.radians(lon) + lat_rad = np.radians(lat) + x = EARTH_RADIUS_KM * np.cos(lat_rad) * np.cos(lon_rad) + y = EARTH_RADIUS_KM * np.cos(lat_rad) * np.sin(lon_rad) + z = EARTH_RADIUS_KM * np.sin(lat_rad) + return np.column_stack([x, y, z]) + + +# def convert_lon(lon: ArrayLike, lon_range: str) -> NDArray[np.float64]: +# """Return longitude values re-labelled into the requested convention. +# Both conventions share the 0 meridian (Greenwich) and turn eastward, no sorting. +# Parameters +# ---------- +# lon : array-like +# lon_range : {"-180-180", "0-360"} +# """ +# lon = np.asarray(lon, dtype=float) +# if lon_range == "-180-180": +# return ((lon + 180.0) % 360.0) - 180.0 +# if lon_range == "0-360": +# return lon % 360.0 +# raise ValueError( +# f"unknown lon_range: {lon_range!r} (expected '-180-180' or '0-360')" +# ) + + +# def detect_lon_range(lon: ArrayLike) -> str: +# """Infer the longitude convention of a grid from its values. +# - any value > 180 -> only "0-360" can produce that +# - any value < 0 -> only "-180-180" can produce that +# - all in [0, 180] -> ambiguous (both conventions coincide there), +# default to "-180-180" +# """ +# lon = np.asarray(lon, dtype=float) +# if lon.max() > 180.0: +# return "0-360" +# if lon.min() < 0.0: +# return "-180-180" +# return "-180-180" diff --git a/tests/test_geo.py b/tests/test_geo.py new file mode 100644 index 0000000..e4bb9f7 --- /dev/null +++ b/tests/test_geo.py @@ -0,0 +1,58 @@ +# tests/test_geo.py + +import numpy as np +import numpy.testing as npt +from floatmatcher.constants import EARTH_RADIUS_KM +from floatmatcher.geo import ( + lonlat_to_xyz +) + + +# ───────────── lonlat_to_xyz ───────────── + +def test_reference_points(): + """ set special points to verify the computations """ + lon = [0, 90, 0, 0] + lat = [0, 0, 90, -90] + xyz = lonlat_to_xyz(lon, lat) + + expected = np.array([ + [EARTH_RADIUS_KM, 0, 0], # equator / Greenwich + [0, EARTH_RADIUS_KM, 0], # equator / 90°E + [0, 0, EARTH_RADIUS_KM], # north pole + [0, 0, -EARTH_RADIUS_KM], # south pole + ]) + npt.assert_allclose(xyz, expected, atol=1e-9) + + +def test_output_shape(): + """N points entry -> output (N, 3)""" + xyz = lonlat_to_xyz([0, 10, 20, 30], [0, 10, 20, 30]) + assert xyz.shape == (4, 3) + + +def test_points_lie_on_sphere(): + """every point is at the same distance from the center""" + rng = np.random.default_rng(0) + lon = rng.uniform(-180, 180, size=100) + lat = rng.uniform(-90, 90, size=100) + xyz = lonlat_to_xyz(lon, lat) + + norms = np.linalg.norm(xyz, axis=1) # compute norm of each xyz lon/lat vector + npt.assert_allclose(norms, EARTH_RADIUS_KM, atol=1e-9) + + +def test_antimeridian_points_are_close(): + """2 points near -180 and 180 are close in 3D""" + xyz = lonlat_to_xyz([179, -179], [0, 0]) + dist = np.linalg.norm(xyz[0] - xyz[1]) + # equator: 2° ~222 km + assert dist < 250 + + +def test_xyz_convention_invariant(): + """each convention 0-360 and -180-180 must give the same xyz""" + a = lonlat_to_xyz([350.0], [30.0]) # 350 == -10 : same meridian + b = lonlat_to_xyz([-10.0], [30.0]) + npt.assert_allclose(a, b, atol=1e-9) + From e6806fe62a617be4e936bac13c0aabfb0c9a656a Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 1 Sep 2026 16:22:53 +0000 Subject: [PATCH 03/12] edit pyproject.toml to ease mypy check types on tests --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index df10d22..f1c9fce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,10 @@ warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true +[tool.mypy.overrides] +module = "tests.*" +disallow_untyped_defs = false + [tool.ruff] line-length = 88 target-version = "py311" From fb6a6d9b55cb8869b98c4eb324ff25b782f308e2 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 1 Sep 2026 16:24:53 +0000 Subject: [PATCH 04/12] forgot remaining useless old functions : cleaning --- src/floatmatcher/geo.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/src/floatmatcher/geo.py b/src/floatmatcher/geo.py index 7666a39..879f585 100644 --- a/src/floatmatcher/geo.py +++ b/src/floatmatcher/geo.py @@ -18,34 +18,3 @@ def lonlat_to_xyz(lon: ArrayLike, lat: ArrayLike) -> NDArray[np.float64]: return np.column_stack([x, y, z]) -# def convert_lon(lon: ArrayLike, lon_range: str) -> NDArray[np.float64]: -# """Return longitude values re-labelled into the requested convention. -# Both conventions share the 0 meridian (Greenwich) and turn eastward, no sorting. -# Parameters -# ---------- -# lon : array-like -# lon_range : {"-180-180", "0-360"} -# """ -# lon = np.asarray(lon, dtype=float) -# if lon_range == "-180-180": -# return ((lon + 180.0) % 360.0) - 180.0 -# if lon_range == "0-360": -# return lon % 360.0 -# raise ValueError( -# f"unknown lon_range: {lon_range!r} (expected '-180-180' or '0-360')" -# ) - - -# def detect_lon_range(lon: ArrayLike) -> str: -# """Infer the longitude convention of a grid from its values. -# - any value > 180 -> only "0-360" can produce that -# - any value < 0 -> only "-180-180" can produce that -# - all in [0, 180] -> ambiguous (both conventions coincide there), -# default to "-180-180" -# """ -# lon = np.asarray(lon, dtype=float) -# if lon.max() > 180.0: -# return "0-360" -# if lon.min() < 0.0: -# return "-180-180" -# return "-180-180" From fb424929aea2ec5d520ad365987e6f2103f02d18 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 15:53:53 +0000 Subject: [PATCH 05/12] add PointSet with validation and xyz cache --- src/floatmatcher/pointset.py | 45 +++++++++++++++++++++++ tests/helpers.py | 15 ++++++++ tests/test_pointset.py | 71 ++++++++++++++++++++++++++++++++++++ 3 files changed, 131 insertions(+) create mode 100644 src/floatmatcher/pointset.py create mode 100644 tests/helpers.py create mode 100644 tests/test_pointset.py diff --git a/src/floatmatcher/pointset.py b/src/floatmatcher/pointset.py new file mode 100644 index 0000000..1680494 --- /dev/null +++ b/src/floatmatcher/pointset.py @@ -0,0 +1,45 @@ +# pointset.py: validated container for the input points to colocalize + +from dataclasses import dataclass, field + +import numpy as np +import xarray as xr +from numpy.typing import NDArray +from .geo import lonlat_to_xyz +from .constants import TIME_UNIT + + +@dataclass +class PointSet: + """Validated wrapper around the (lon, lat, time) arrays to colocalize. + + The wrapper carries validation and clear names, but heavy computation + works directly on the underlying NumPy arrays (``points.lon``), never on + the object itself inside loops. + + Dataset travels along origin_ds + """ + + lon: NDArray[np.float64] + lat: NDArray[np.float64] + time: NDArray[np.datetime64] + origin_dim: str | None = None # source dimension name + origin_ds: xr.Dataset | None = None # source dataset + _xyz: NDArray[np.float64] | None = field(default=None, init=False, repr=False) + + + def __post_init__(self) -> None: + """validation for lenghts and dtype of lon/lat/time""" + self.lon = np.asarray(self.lon, dtype=float) + self.lat = np.asarray(self.lat, dtype=float) + self.time = np.asarray(self.time, dtype=f"datetime64[{TIME_UNIT}]") + if not (len(self.lat) == len(self.time) == len(self.lon)): + raise ValueError("lon, lat, time must have the same length") + + @property + def xyz(self) -> NDArray[np.float64]: + """Cartesian 3D coordinates on the sphere, computed once and cached.""" + if self._xyz is None: + self._xyz = lonlat_to_xyz(self.lon, self.lat) + return self._xyz + diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..8336543 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,15 @@ +# tests/helpers.py: helpers shared across the test folder +import numpy as np + + +def set_timestamps(*days): + """Timestamps for the given 2015-01-xx days.""" + stamps = [] + for d in days: + stamps.append(f"2015-01-{d:02d}") + return np.array(stamps, dtype="datetime64[ns]") + + +def daily_timestamps(n): + """n consecutive daily timestamps starting 2015-01-01.""" + return set_timestamps(*range(1, n + 1)) diff --git a/tests/test_pointset.py b/tests/test_pointset.py new file mode 100644 index 0000000..ea11287 --- /dev/null +++ b/tests/test_pointset.py @@ -0,0 +1,71 @@ +# tests/test_pointset.py +import numpy as np +import pytest + +from floatmatcher.pointset import PointSet + +from helpers import daily_timestamps + + +# ───────────── validation ───────────── +def test_lon_dtype(): + """test dtype of time array""" + ps = PointSet([-45.0, -44.0], [32.0, 33.0], daily_timestamps(2)) + assert ps.lon.dtype == float + +def test_lat_dtype(): + """test dtype of time array""" + ps = PointSet([-45.0, -44.0], [32.0, 33.0], daily_timestamps(2)) + assert ps.lat.dtype == float + +def test_time_dtype(): + """test dtype of time array""" + ps = PointSet([-45.0, -44.0], [32.0, 33.0], daily_timestamps(2)) + assert ps.time.dtype == "datetime64[ns]" + + +def test_valid_lengths_ok(): + """Matching lengths build without error.""" + ps = PointSet([-45.0, -44.0], [32.0, 33.0], daily_timestamps(2)) + assert len(ps.lon) == len(ps.lat) == len(ps.time) + + +def test_mismatched_lengths_raise(): + """Different lengths are rejected at construction.""" + with pytest.raises(ValueError): + PointSet([-45.0, -44.0], [32.0], daily_timestamps(2)) # lat too short + + +def test_arrays_are_converted(): + """Python lists become float NumPy arrays after construction.""" + ps = PointSet([1, 2, 3], [4, 5, 6], daily_timestamps(3)) # ints on purpose + assert isinstance(ps.lon, np.ndarray) + assert ps.lon.dtype == float + + +# ───────────── xyz cache ───────────── + +def test_xyz_shape(): + """xyz exposes one (x, y, z) row per point.""" + ps = PointSet([0.0, 90.0, 0.0], [0.0, 0.0, 90.0], daily_timestamps(3)) + assert ps.xyz.shape == (3, 3) + + +def test_xyz_is_cached(): + """xyz is computed once and the same array is returned afterwards. + test @property proper working """ + ps = PointSet([10.0, 20.0], [30.0, 40.0], daily_timestamps(2)) + first = ps.xyz + second = ps.xyz + assert first is second # same object, not just equal → cache hit + + + +# ───────────── provenance ───────────── + +def test_provenance_defaults_to_none(): + """Without a source, provenance fields are None.""" + ps = PointSet([1.0], [2.0], daily_timestamps(1)) + assert ps.origin_dim is None + + From c16dd46654a1f8837ef50fd9a36e083b5099f533 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 16:08:53 +0000 Subject: [PATCH 06/12] add grid set and tests. clean tests and conftest to amtch actual step dev --- src/floatmatcher/gridset.py | 50 ++++++++++++++++++++++++++++++ tests/conftest.py | 29 ++++++++++++++++++ tests/helpers.py | 26 +++++++++++++++- tests/test_gridset.py | 61 +++++++++++++++++++++++++++++++++++++ 4 files changed, 165 insertions(+), 1 deletion(-) create mode 100644 src/floatmatcher/gridset.py create mode 100644 tests/conftest.py create mode 100644 tests/test_gridset.py diff --git a/src/floatmatcher/gridset.py b/src/floatmatcher/gridset.py new file mode 100644 index 0000000..2e11769 --- /dev/null +++ b/src/floatmatcher/gridset.py @@ -0,0 +1,50 @@ +# gridset.py: validated wrapper around a normalized grid dataset + +from dataclasses import dataclass, field +import xarray as xr +import numpy as np + +@dataclass +class GridSet: + """Thin validated wrapper around a normalized grid dataset. + + Wraps a normalized ``xr.Dataset`` (coords ``lat``/``lon``, and ``time`` in + the 3D case), validating it at construction and deriving its regime. The + dataset is kept intact and accessible as ``grid.dataset`` + """ + + dataset: xr.Dataset + # convention the Product promises; None -> skip the check + regime: str = field(init=False) # "2D" or "3D", derived at construction + + def __post_init__(self) -> None: + + if "lon" not in self.dataset.coords or "lat" not in self.dataset.coords: + raise ValueError( + "The dataset given to GridSet object doesn't have lon or lat " + "coordinates" + ) + + if len(self.dataset.data_vars)<1: + raise ValueError("There is no variable in the dataset given to GridSet") + + # test of lat/lon unicity + lon = self.dataset["lon"].values + lat = self.dataset["lat"].values + if len(np.unique(lon)) != len(lon): + raise ValueError("grid: duplicated longitudes in array") + if len(np.unique(lat)) != len(lat): + raise ValueError("grid: duplicated latitudes in array") + + # select regime 3D/2D + if "time" in self.dataset.coords : + self.regime = "3D" + # test time unicity if 3D regime + times = self.dataset["time"].values + if len(np.unique(times)) != len(times): + raise ValueError("grid: duplicate timestamps (overlapping files?)") + else: + self.regime = "2D" + + + \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ab0cbd1 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,29 @@ +# tests/conftest.py: fixtures shared across the test suite + + +import numpy as np +import pytest +import xarray as xr + +from helpers import make_grid, daily_timestamps + + +# ---------- grids ---------- + +lat = [0.0, 1.0, 2.0] +lon = [10.0, 20.0, 30.0, 40.0] + + +@pytest.fixture +def grid_2d_ds(): + ds = make_grid(lat, lon) + ds["v"] = 100.0 * ds["lat"] + ds["lon"] + return ds + + +@pytest.fixture +def grid_3d_ds(): + ds = make_grid(lat, lon, time=daily_timestamps(2)) + ds["v"] = (100.0 * ds["lat"] + ds["lon"] + + xr.DataArray(np.arange(ds.sizes["time"]), dims="time")) + return ds \ No newline at end of file diff --git a/tests/helpers.py b/tests/helpers.py index 8336543..0fdf7e6 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,6 +1,6 @@ # tests/helpers.py: helpers shared across the test folder import numpy as np - +import xarray as xr def set_timestamps(*days): """Timestamps for the given 2015-01-xx days.""" @@ -13,3 +13,27 @@ def set_timestamps(*days): def daily_timestamps(n): """n consecutive daily timestamps starting 2015-01-01.""" return set_timestamps(*range(1, n + 1)) + +def make_grid(lat, lon, time=None, variables=("v",)): + """Build a GridSet-ready Dataset filled with zeros. + """ + lat = np.asarray(lat, dtype=float) + lon = np.asarray(lon, dtype=float) + if isinstance(variables, str): + variables = (variables,) + + if time is None: + dims = ("lat", "lon") + coords = {"lat": lat, "lon": lon} + shape = (lat.size, lon.size) + else: + time = np.asarray(time, dtype="datetime64[ns]") + dims = ("time", "lat", "lon") + coords = {"time": time, "lat": lat, "lon": lon} + shape = (time.size, lat.size, lon.size) + + data_vars = {} + for name in variables: + data_vars[name] = (dims, np.zeros(shape)) + + return xr.Dataset(data_vars, coords=coords) \ No newline at end of file diff --git a/tests/test_gridset.py b/tests/test_gridset.py new file mode 100644 index 0000000..75c2342 --- /dev/null +++ b/tests/test_gridset.py @@ -0,0 +1,61 @@ +# tests/test_gridset.py +import pytest + +from floatmatcher.gridset import GridSet + + +# ───────────── regime derivation ───────────── + +def test_regime_3d(grid_3d_ds): + """A grid with a time coord is 3D.""" + assert GridSet(grid_3d_ds).regime == "3D" + +def test_regime_2d(grid_2d_ds): + """A grid without a time coord is 2D.""" + assert GridSet(grid_2d_ds).regime == "2D" + + +# ───────────── validation ───────────── + +def test_dataset_is_kept(grid_3d_ds): + """GridSet wraps the dataset without transforming it.""" + grid = GridSet(grid_3d_ds) + assert grid.dataset is grid_3d_ds # same object, not a copy + +def test_missing_lat_3d_raises(grid_3d_ds): + """A grid without lat is rejected.""" + with pytest.raises(ValueError): + GridSet(grid_3d_ds.drop_vars("lat")) + +def test_missing_lat_2d_raises(grid_2d_ds): + """A grid without lat is rejected.""" + with pytest.raises(ValueError): + GridSet(grid_2d_ds.drop_vars("lat")) + +def test_no_data_variable_raises(grid_2d_ds): + """A grid with coords but no data variable is rejected.""" + with pytest.raises(ValueError): + GridSet(grid_2d_ds.drop_vars("v")) # removes the only variable + + + + +# ───────────── suffled and duplicated coordinates ───────────── + +def test_shuffled_lon_is_accepted(grid_3d_ds): + """ order does not matter: the KDTree works on a point cloud""" + ds = grid_3d_ds.isel(lon=[2, 0, 1]) # [30,10,20] + assert GridSet(ds).regime == "3D" + + +def test_duplicate_lon_raises(grid_3d_ds): + """two nodes at the same position would make the nearest lookup ambiguous""" + ds = grid_3d_ds.isel(lon=[0, 1, 1]) + with pytest.raises(ValueError): + GridSet(ds) + + +def test_decreasing_time_is_accepted(grid_3d_ds): + """ decreasing time is accepted """ + ds = grid_3d_ds.isel(time=slice(None, None, -1)) + assert GridSet(ds).regime == "3D" From 188327841275d0c623fb3cd4326af6c439aaedd3 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 16:19:45 +0000 Subject: [PATCH 07/12] add flatgrid object and tests --- src/floatmatcher/flatgrid.py | 49 ++++++++++++++++++++++++ tests/test_flatgrid.py | 72 ++++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+) create mode 100644 src/floatmatcher/flatgrid.py create mode 100644 tests/test_flatgrid.py diff --git a/src/floatmatcher/flatgrid.py b/src/floatmatcher/flatgrid.py new file mode 100644 index 0000000..41d2e7d --- /dev/null +++ b/src/floatmatcher/flatgrid.py @@ -0,0 +1,49 @@ +# reference.py: the reference point cloud (flattened grid), internal to nearest + +from dataclasses import dataclass, field + +import numpy as np +import xarray as xr +from numpy.typing import ArrayLike, NDArray + +from .geo import lonlat_to_xyz + + +@dataclass +class FlatGrid: + """Flattened grid as a point cloud, for KDTree lookup.""" + + lon: NDArray[np.float64] + lat: NDArray[np.float64] + time: NDArray[np.datetime64] | None + _stacked: xr.Dataset # lazy xr dataset + _xyz: NDArray[np.float64] | None = field(default=None, init=False, repr=False) + + @property + def xyz(self) -> NDArray[np.float64]: + """Cartesian 3D coordinates of the nodes, computed once and cached.""" + if self._xyz is None: + self._xyz = lonlat_to_xyz(self.lon, self.lat) + return self._xyz + + @classmethod + def from_grid(cls, ds: xr.Dataset) -> "FlatGrid": + """Flatten a grid (2D or 3D) into a node cloud. Values stay lazy.""" + stacked = ds.stack(node=("lat", "lon")) + time = ds["time"].values if "time" in ds.coords else None + return cls(lon=stacked["lon"].values, lat=stacked["lat"].values, + time=time, _stacked=stacked) + + def read_values(self, node_idx: ArrayLike, + tsel_idx: ArrayLike | None = None) -> dict[str, NDArray[np.float64]]: + """Read variable values ONLY at the retained (node[, time]) indices""" + node = xr.DataArray(np.asarray(node_idx), dims="pts") + out = {} + for var in self._stacked.data_vars: # _stacked is lazy ds with only lon/lat/time in memory + da_var = self._stacked[var] + if tsel_idx is None: + sel = da_var.isel(node=node) + else: + sel = da_var.isel(node=node, time=xr.DataArray(np.asarray(tsel_idx), dims="pts")) + out[str(var)] = sel.values + return out diff --git a/tests/test_flatgrid.py b/tests/test_flatgrid.py new file mode 100644 index 0000000..5a6f19b --- /dev/null +++ b/tests/test_flatgrid.py @@ -0,0 +1,72 @@ +# tests/test_flatgrid.py +# +# FlatGrid flattens a grid into a node cloud and stays LAZY: read_values pulls +# only the retained (node[, time]) pairs, never the whole cube. + +import numpy as np +import numpy.testing as npt + +from floatmatcher.flatgrid import FlatGrid + +# Reminder +# 3D (grid_3d_ds): +# n0 n1 n2 n3 n4 n5 n6 n7 n8 n9 n10 n11 +# t=0 10.0 20.0 30.0 40.0 110.0 120.0 130.0 140.0 210.0 220.0 230.0 240.0 +# t=1 11.0 21.0 31.0 41.0 111.0 121.0 131.0 141.0 211.0 221.0 231.0 241.0 + +# 2D (grid_2d_ds) : +# n0 n1 n2 n3 n4 n5 n6 n7 n8 n9 n10 n11 +# t=0 10.0 20.0 30.0 40.0 110.0 120.0 130.0 140.0 210.0 220.0 230.0 240.0 + + +def test_flatten_preserves_alignment(grid_2d_ds): + """Each flattened node keeps the value of its own coordinates""" + flat_grid = FlatGrid.from_grid(grid_2d_ds) + + assert flat_grid.lon.size == 3 * 4 # lat * lon + + values = flat_grid.read_values(np.arange(12))["v"] # 2D: no time index + npt.assert_allclose(values, [10.0, 20.0, 30.0, 40.0, + 110.0, 120.0, 130.0, 140.0, + 210.0, 220.0, 230.0, 240.0]) + + +def test_flatgrid_is_2d_has_no_time(grid_2d_ds): + """A 2D grid produces a FlatGrid with time=None.""" + assert FlatGrid.from_grid(grid_2d_ds).time is None + + +def test_flatgrid_xyz_shape(grid_2d_ds): + """xyz exposes one (x, y, z) row per node.""" + assert FlatGrid.from_grid(grid_2d_ds).xyz.shape == (12, 3) + + +def test_retreive_right_indexes(grid_3d_ds): + """Indices are paired element-wise, not crossed: node[i] with time[i]""" + flat_grid = FlatGrid.from_grid(grid_3d_ds) + + # node 3 read at t=0 -> 40 ; node 4 read at t=1 -> 111 + flat_grid_v = flat_grid.read_values([3, 4], [0, 1])["v"] + expected_values = [40.0, 111.0] + npt.assert_allclose(flat_grid_v, expected_values) + + + +def test_read_values_returns_every_variable(grid_2d_ds): + """read_values loops over data_vars""" + grid_2d_ds["sst"] = grid_2d_ds["v"] + grid_2d_ds["t2m"] = grid_2d_ds["v"] + 1000.0 + + flat_grid = FlatGrid.from_grid(grid_2d_ds) + grid_values = flat_grid.read_values([0, 4]) # read node 0 and 4 on 2d grid --> dict + + assert set(grid_values) == {"v", "sst", "t2m"} + npt.assert_allclose(grid_values["sst"], [10.0, 110.0]) + npt.assert_allclose(grid_values["t2m"], [1010.0, 1110.0]) + + +def test_xyz_is_cached(grid_2d_ds): + """xyz is computed once: the same array object comes back. + test if @property works well""" + flat_grid = FlatGrid.from_grid(grid_2d_ds) + assert flat_grid.xyz is flat_grid.xyz From 2a78f2992bffd22dfd9ca91d564643af22a0a66f Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 16:30:11 +0000 Subject: [PATCH 08/12] add a bit more lenght to lines for ruff 110 instad of 88 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f1c9fce..42456b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ module = "tests.*" disallow_untyped_defs = false [tool.ruff] -line-length = 88 +line-length = 110s target-version = "py311" [tool.ruff.lint] From e63cc1c147de145d3872f7a50b90cefdc64ef90f Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 16:31:11 +0000 Subject: [PATCH 09/12] mistyping corrected --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 42456b8..c43e7fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ module = "tests.*" disallow_untyped_defs = false [tool.ruff] -line-length = 110s +line-length = 110 target-version = "py311" [tool.ruff.lint] From 098296e3a435e09b950cf1ff99f3b41df2ed8175 Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 16:46:47 +0000 Subject: [PATCH 10/12] add results, tests and correct mypy call in pyproject.toml --- pyproject.toml | 4 +-- src/floatmatcher/results.py | 36 ++++++++++++++++++++ tests/test_results.py | 68 +++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 3 deletions(-) create mode 100644 src/floatmatcher/results.py create mode 100644 tests/test_results.py diff --git a/pyproject.toml b/pyproject.toml index c43e7fa..9d4a8d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,9 +39,7 @@ warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true -[tool.mypy.overrides] -module = "tests.*" -disallow_untyped_defs = false +exclude = ["^tests/"] [tool.ruff] line-length = 110 diff --git a/src/floatmatcher/results.py b/src/floatmatcher/results.py new file mode 100644 index 0000000..a9cfa4d --- /dev/null +++ b/src/floatmatcher/results.py @@ -0,0 +1,36 @@ +# results.py: output object of matchup + +import numpy as np +import xarray as xr + +from dataclasses import dataclass +from numpy.typing import NDArray +from .pointset import PointSet + + +@dataclass +class MatchupResult: + values: dict[str, NDArray[np.float64]] + distance_km: NDArray[np.float64] + time_delta: NDArray[np.float64] + valid: NDArray[np.bool_] + points: PointSet + + def to_dataset(self) -> xr.Dataset: + """ + Reinject the colocalized values into the dataset + The source Dataset travels inside the PointSet + """ + ds = self.points.origin_ds + dim = self.points.origin_dim + if ds is None or dim is None: + raise ValueError( + "Cannot reinject: these points have no origin dataset " + "(they came from raw arrays). Use the MatchupResult directly." + ) + + out = ds.copy() + for k, v in self.values.items(): + out[f"{k}_coloc"] = (dim, v) + + return out diff --git a/tests/test_results.py b/tests/test_results.py new file mode 100644 index 0000000..992ba86 --- /dev/null +++ b/tests/test_results.py @@ -0,0 +1,68 @@ +# tests/test_results.py +import numpy as np +import pytest +import xarray as xr + +from floatmatcher.results import MatchupResult +from floatmatcher.pointset import PointSet + + +@pytest.fixture +def result(): + """ MatchupResult object - 3 profiles """ + ds = xr.Dataset( + {"LONGITUDE": (("N_PROF",), np.array([-40.0, -30.0, -20.0])), + "LATITUDE": (("N_PROF",), np.array([35.0, 36.0, 37.0])), + "TIME": (("N_PROF",), np.array(["2018-01-03", "2018-01-01", "2018-01-02"], + dtype="datetime64[ns]"))}, + coords={"N_PROF": np.array([10, 11, 12])}, + ) + points = PointSet( + lon=ds["LONGITUDE"].values, + lat=ds["LATITUDE"].values, + time=ds["TIME"].values, + origin_dim="N_PROF", + origin_ds=ds, + ) + return MatchupResult( + values={"dummy_variable": np.array([280.0, 281.0, np.nan])}, + distance_km=np.array([1.0, 2.0, np.nan]), + time_delta=np.array([0.1, 0.2, np.nan]), + valid=np.array([True, True, False]), + points=points, + ) + +# @pytest.fixture +# def argo_like_ds(): +# return xr.Dataset( +# {"LONGITUDE": (("N_PROF",), np.array([-40.0, -30.0, -20.0])), +# "LATITUDE": (("N_PROF",), np.array([35.0, 36.0, 37.0])), +# "TIME": (("N_PROF",), np.array(["2018-01-03", "2018-01-01", "2018-01-02"], +# dtype="datetime64[ns]"))}, +# coords={"N_PROF": np.array([10, 11, 12])}, +# ) + + +# ---------- to_dataset ---------- + +def test_to_dataset_adds_coloc_variable(result): + out = result.to_dataset() + assert out["dummy_variable_coloc"].dims == ("N_PROF",) + assert out["dummy_variable_coloc"].sel(N_PROF=10) == 280.0 + + +def test_to_dataset_does_not_mutate_the_source(result): + result.to_dataset() + assert "dummy_variable_coloc" not in result.points.origin_ds.data_vars + + +def test_to_dataset_without_provenance_raises(): + """Raw-array points carry no dataset, so they cannot be reinjected.""" + points = PointSet([-45.0, -44.0, -43.0], [32.0, 33.0, 34.0], + np.array(["2015-01-01", "2015-01-02", "2015-01-03"], + dtype="datetime64[ns]")) + res = MatchupResult(values={"dummy_variable": np.zeros(3)}, + distance_km=np.zeros(3), time_delta=np.zeros(3), + valid=np.ones(3, bool), points=points) + with pytest.raises(ValueError): + res.to_dataset() From 82b6bf11290981d395367397a2b13112e62371ad Mon Sep 17 00:00:00 2001 From: Clement WEBER Date: Tue, 8 Sep 2026 17:34:24 +0000 Subject: [PATCH 11/12] 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 12/12] 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]