diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 10c2e97..e48f3ba 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,3 +32,25 @@ jobs: - name: Test (pytest) run: uv run pytest + + docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v3 + with: + python-version: "3.12" + + - name: Install dependencies (docs group) + run: uv sync --group docs + + - name: Build docs (warnings are errors) + run: uv run sphinx-build -W --keep-going -b html docs docs/_build/html + + - name: Upload built docs + uses: actions/upload-artifact@v4 + with: + name: docs-html + path: docs/_build/html diff --git a/.gitignore b/.gitignore index 206aa22..62e2543 100644 --- a/.gitignore +++ b/.gitignore @@ -145,3 +145,6 @@ uv.lock xxx/ /old* .mongo-data +*.nc +era5_float_6900166/ +*.zip \ No newline at end of file diff --git a/README.md b/README.md index f5f12b5..977616c 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,28 @@ # FloatMatcher -A python library for third-party dataset matchup with Argo floats -test + +Floatmatcher is a a python library for third-party dataset matchup with Argo floats +Third-party datasets can originate from local or distant repositories. + +## Documentation +Asociated documentation is available running : ``` python3 -m http.server 8000 -d docs/_build/html ``` after installation + +## Installation +To install floatmatcher library, use the current git repository : +### uv +``` uv add "floatmatcher @ git+https://github.com/euroargodev/FloatMatcher.git@main" ``` + +### pip +``` pip install "git+https://github.com/euroargodev/FloatMatcher.git@main" ``` + +## Usage +```python +# Import the main builder: +from floatmatcher import floatmatcher +``` + + + +## Additional tutorials + + + diff --git a/examples/tutorial.py b/examples/tutorial.py index 54152ca..b947225 100644 --- a/examples/tutorial.py +++ b/examples/tutorial.py @@ -34,7 +34,6 @@ # 1.3) dataset xr + mapping dictionnary # # 1.4) argopy (no mapping needed) K -# TODO : set_points_from dataframe/dataset # 1.1 example my_fm.set_points(lon, lat, time) @@ -63,9 +62,10 @@ # radius (km) # time -my_fm.set_method(type="nearest", radius=12, +my_fm.set_method(type="interp", radius=12, time_gap=np.timedelta64(3, "h"), - k_nearest=1 + k_nearest=4, n_frames=2, kernel="idw", + power=2 ) diff --git a/src/floatmatcher/flatgrid.py b/src/floatmatcher/flatgrid.py index e8df5b0..827b74b 100644 --- a/src/floatmatcher/flatgrid.py +++ b/src/floatmatcher/flatgrid.py @@ -42,6 +42,11 @@ def xyz(self) -> NDArray[np.float64]: self._xyz = lonlat_to_xyz(self.lon, self.lat) return self._xyz + @property + def variables(self) -> list[str]: + """Names of the variables of the grid : avoid hashable type asked by mypy""" + return [str(v) for v in self._stacked.data_vars] + @classmethod def from_grid(cls, ds: xr.Dataset) -> "FlatGrid": """Flatten a grid, 2D or 3D, into a node cloud @@ -90,6 +95,18 @@ def read_values(self, node_idx: ArrayLike, out[str(var)] = sel.values return out + + def read_variable(self, var: str, node_idx: NDArray, time_idx: NDArray|None) -> NDArray: + """ return variable only on nodes and times selected """ + node = xr.DataArray(np.asarray(node_idx), dims=("pts", "K")) + da_var = self._stacked[var] + if time_idx is None: + return da_var.isel(node=node).transpose("pts", "K").values + tsel = xr.DataArray(np.asarray(time_idx), dims=("pts", "N")) + return da_var.isel(node=node, time=tsel).transpose("pts", "K", "N").values + + + def spatial_nearest(self, points: PointSet, k: int = 1) -> tuple[NDArray[np.float64], NDArray[np.int64]]: """Nearest nodes of the points in space @@ -108,7 +125,7 @@ def spatial_nearest(self, points: PointSet, """ dist: NDArray[np.float64] idx: NDArray[np.int64] - dist, idx = cKDTree(self.xyz).query(points.xyz, k=k) + dist, idx = cKDTree(self.xyz).query(points.xyz, k=k, workers=-1) return dist, idx def temporal_nearest(self, points: PointSet, @@ -138,7 +155,7 @@ def temporal_nearest(self, points: PointSet, if points.time is None: raise ValueError("points have no time: they cannot be matched on a grid with a time axis") grid_tree = cKDTree(_to_seconds(self.time)[:, None]) - time_delta, idx = grid_tree.query(_to_seconds(points.time)[:, None], k=k) + time_delta, idx = grid_tree.query(_to_seconds(points.time)[:, None], k=k, workers=-1) return time_delta, idx diff --git a/src/floatmatcher/methods.py b/src/floatmatcher/methods.py index ff68bf1..c807d91 100644 --- a/src/floatmatcher/methods.py +++ b/src/floatmatcher/methods.py @@ -13,7 +13,7 @@ from .matchup_results import MatchupResult from .pointset import PointSet from .utils import _indented_repr - +from .weights import available_kernels class Method(ABC): """Prototype of the matchup methods @@ -23,9 +23,16 @@ class Method(ABC): points. Subclasses implement ``apply`` and are registered in ``available_methods``. """ + time_gap: np.timedelta64 + def __repr__(self) -> str: return _indented_repr(self) + @property + def time_gap_seconds(self) -> float: + """``time_gap`` in seconds""" + return float(self.time_gap / np.timedelta64(1, "s")) + @abstractmethod def apply(self, dataset: xr.Dataset, points: PointSet) -> MatchupResult: """Match the points with a grid @@ -67,11 +74,6 @@ def __init__(self, radius: int = 25, self.time_gap = time_gap self.k_nearest = k_nearest - @property - def time_gap_seconds(self) -> float: - """``time_gap`` in seconds""" - return float(self.time_gap / np.timedelta64(1, "s")) - def apply(self, dataset: xr.Dataset, points: PointSet) -> MatchupResult: """Match the points with the nearest node of the grid @@ -139,8 +141,124 @@ def apply(self, dataset: xr.Dataset, points: PointSet) -> MatchupResult: time_delta=time_delta_out, valid=valid) + + +class Interpolator(Method): + """Interpolation based on Nearest-neighbor fit method + + Each point takes the values of the combination of closest nodes of the grid, + and of the nearests times step if the grid has a time axis. + Points beyond ``radius`` or ``time_gap`` are invalid. + + Parameters + ---------- + radius: int, default=25 + Maximum distance to the node, in km (straight line on the sphere). + time_gap: numpy.timedelta64, default=1 day + Maximum gap to the time step. Ignored for a grid without time axis. + k_nearest: int, default=1 + Number of neighbors looked up. Only 1 works for now. + """ + + def __init__(self, + radius: int = 25, + time_gap: np.timedelta64 = np.timedelta64(1, "D"), + k_nearest : int = 2, + n_frames: int = 1, + kernel: str = "idw", + power: int = 2 + ) -> None : + self.radius = radius + self.time_gap = time_gap + self.k_nearest = k_nearest + self.n_frames = n_frames + self.kernel = kernel + self.power = power + + + def apply(self, dataset: xr.Dataset, points: PointSet) -> MatchupResult: + """Match the points with the nearest node of the grid""" + + # starting by lonlat2xy on spatial grid + # return FlatGrid object, flatten grid + flat_grid = FlatGrid.from_grid(dataset) + + # apply kdtree on spatial : find k nearest distances and indexes of nearest points + dist_km, spatial_idx = flat_grid.spatial_nearest(points, k=self.k_nearest) + valid_spatial = dist_km <= self.radius + # reshape to fit (n,) --> (n,K) in 1D: + spatial_idx = spatial_idx.reshape(len(points.lon), -1) + dist_km = dist_km.reshape(len(points.lon), -1) + valid_spatial = valid_spatial.reshape(len(points.lon), -1) + + # length of points to process + idx_count = len(points.lon) + + # apply kdtree on time : find k nearest "time distances" + if flat_grid.time is not None: + time_delta, temporal_idx = flat_grid.temporal_nearest(points, k=self.n_frames) + # Reshape also + temporal_idx = temporal_idx.reshape(len(points.lon), -1) + time_delta = time_delta.reshape(len(points.lon), -1) + valid_time = time_delta <= self.time_gap_seconds + else: + temporal_idx = None + time_delta = np.full((idx_count, 1), np.nan) + valid_time = np.ones((idx_count, 1), dtype=bool) + + # load kernel + kernel = available_kernels[self.kernel] + + # compute spatial weights + # return (n, K) matrix + spatial_weights = kernel(dist_km, power=self.power, limit=self.radius) * valid_spatial + + # compute temporal weights + if flat_grid.time is None: + temporal_weights = np.ones((idx_count, 1)) # no time axis + else: + temporal_weights = kernel(time_delta, power=self.power, + limit=self.time_gap_seconds) * valid_time + + # read EVERY points and let the equation select good points + values = {} + for var in flat_grid.variables: + + # load variable into memory, one variable at a time + var_picked = flat_grid.read_variable(var, spatial_idx, temporal_idx) # return (n, K, N) matrix + if temporal_idx is None: + var_picked = var_picked[:, :, None] # return (n, K, None) matrix + + # apply weights for each time step + # (n, N), 1 value wighted per time frame + var_weighted_sum = np.nansum(var_picked * spatial_weights[:, :, None], axis=1) + # don't take nan into account # (n, N), 1 weight per time frame + weights_sum = np.nansum(spatial_weights[:,:,None] * ~np.isnan(var_picked), axis=1) + + var_weighted_frames = np.full(var_weighted_sum.shape, np.nan, dtype=var_picked.dtype) # (n, N) + np.divide(var_weighted_sum, weights_sum, + out=var_weighted_frames, where=weights_sum > 0) # (n, N) + + # contract the frames : (n, N) -> (n,) + var_weighted_sum_time = np.nansum(var_weighted_frames * temporal_weights, axis=1) + weights_sum_time = np.nansum(temporal_weights * ~np.isnan(var_weighted_frames), axis=1) + + var_weighted_single = np.full(idx_count, np.nan, dtype=var_picked.dtype) + + np.divide(var_weighted_sum_time, weights_sum_time, + out=var_weighted_single, where=weights_sum_time > 0) + + values[var] = var_weighted_single + + valid = np.asarray(valid_spatial.any(axis=1) & valid_time.any(axis=1)) # (n,) + + return MatchupResult(values=values, + distance_km=np.where(valid, dist_km.min(axis=1), np.nan), + time_delta=np.where(valid, time_delta.min(axis=1), np.nan), + valid=valid) + # key given to FloatMatcher.set_method(type=...) -> method class available_methods: dict[str, type[Method]] = { - "nearest": NearestNeighbor, + "nearest": NearestNeighbor, "interp": Interpolator } """dict: Method classes by name, used by :meth:`.FloatMatcher.set_method`""" diff --git a/src/floatmatcher/weights.py b/src/floatmatcher/weights.py new file mode 100644 index 0000000..dc01088 --- /dev/null +++ b/src/floatmatcher/weights.py @@ -0,0 +1,49 @@ +# weights.py: weighting kernels for the interpolation methods. +# +# A kernel turns distances into UNNORMALIZED weights, and nothing else: it knows neither the +# radius nor the unit of what it weighs. The same function therefore serves the spatial axis +# (km) and the temporal one (seconds). Masking (radius, time gap) and normalization are the +# method's job, because normalization depends on the NaN mask of the data. + +from collections.abc import Callable + +import numpy as np +from numpy.typing import ArrayLike, NDArray + + +def idw(d: ArrayLike, power: float = 1.0, limit: float | None = None, + eps: float = 1e-3) -> NDArray[np.float64]: + """Inverse distance weighting: ``1 / d**power`` + + A point sitting exactly on a node gives an infinite weight. Distances below ``eps`` are + clamped to it, so that node dominates the sum by orders of magnitude without overflowing. + + Parameters + ---------- + d: array-like of float + Distances, any shape. Not modified. + power: float, default=1.0 + The higher it is, the more the closest neighbours dominate. ``power`` to infinity + falls back on the nearest neighbour. + limit: float, optional + Unused, kept so that every kernel shares one signature. + eps: float, default=1e-3 + Distance below which a node is taken to sit on the point. Far below the resolution + of both axes (1 meter in km, 1 milisecond in seconds). + + Returns + ------- + numpy.ndarray + Unnormalized weights, same shape as ``d``. + """ + dist = np.asarray(d, dtype=float) + dist = np.where( dist < eps, eps, dist) # replace where dist < 1m or 1sec by 1 (avoid division by 0) + weights = 1.0 / dist**power + return weights + + + +available_kernels: dict[str, Callable[..., NDArray[np.float64]]] = { + "idw": idw, +} +"""dict: Weighting functions by name, used by :class:`.Interpolator`""" diff --git a/tests/test_methods.py b/tests/test_methods.py index 8d601bc..4f08c36 100644 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -4,8 +4,9 @@ import numpy.testing as npt import pytest -from floatmatcher.methods import Method, NearestNeighbor, available_methods +from floatmatcher.methods import Interpolator, Method, NearestNeighbor, available_methods from floatmatcher.pointset import PointSet +from helpers import make_grid def test_method_is_abstract(): @@ -91,3 +92,55 @@ def test_apply_on_2d_grid_ignores_time(grid_2d_ds): assert res.valid.tolist() == [True] npt.assert_allclose(res.values["v"], [230.0]) assert np.isnan(res.time_delta[0]) + + +# ---------- Interpolator : IDW ---------- +# +# lon=0 lon=10 +# lat=10 30 40 +# lat=0 10 20 + +radius_wide = 5000 # km: wide enough that nothing is rejected + + +@pytest.fixture +def grid_square_ds(): + ds = make_grid([0.0, 10.0], [0.0, 10.0]) + ds["v"] = 2.0 * ds["lat"] + ds["lon"] + 10.0 + return ds + + +def _interp(grid, lon, lat, k=4, power=1.0): + points = PointSet.from_arrays(lon=[lon], lat=[lat]) + return Interpolator(radius=radius_wide, k_nearest=k, power=power).apply(grid, points) + + +def test_interp_stays_in_the_range_of_the_four_nodes(grid_square_ds): + res = _interp(grid_square_ds, lon=3.0, lat=4.0, k=4) + assert 10.0 <= res.values["v"][0] <= 40.0 + + +def test_interp_power_one_is_linear_halfway_between_two_nodes(grid_square_ds): + # lon=5, lat=0 sits at equal distance from the nodes 10 and 20 + # -> with k=2 and power=1 the weights are equal, which is linear interpolation + res = _interp(grid_square_ds, lon=5.0, lat=0.0, k=2) + npt.assert_allclose(res.values["v"], [(10.0 + 20.0) / 2], rtol=1e-4) + + +def test_interp_k_limits_the_neighbours_even_inside_the_radius(grid_square_ds): + two = _interp(grid_square_ds, lon=5.0, lat=0.0, k=2) + + # the 4 nodes are all inside radius, yet k=2 reads only the two closest + npt.assert_allclose(two.values["v"], [(10.0 + 20.0) / 2], rtol=1e-4) + + + +def test_interp_radius_limits_the_neighbours_below_the_requested_k(grid_square_ds): + points = PointSet.from_arrays(lon=[5.0], lat=[0.0]) + # 800km + res = Interpolator(radius=800, k_nearest=4, power=1.0).apply( + grid_square_ds, points) + + # 4 asked, 2 are reachable + npt.assert_allclose(res.values["v"], [(10.0 + 20.0) / 2], rtol=1e-4) + assert res.valid[0]