diff --git a/pyproject.toml b/pyproject.toml index 5a57127..ad97f8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,12 @@ authors = [ {name = "Clement WEBER", email = "clement.weber@pokapok.org"}, ] dependencies = [ - "numpy", + "numpy<2.5", # 2.5+ stubs need Python 3.12 syntax; "xarray", + "scipy", + "pandas", + "dask", + "netCDF4", ] [dependency-groups] @@ -38,8 +42,15 @@ python_version = "3.11" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true +exclude = ["^tests/"] + +[[tool.mypy.overrides]] +module = ["scipy.*", "pandas.*"] +ignore_missing_imports = true [tool.ruff] -line-length = 88 +line-length = 110 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).""" + 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/src/floatmatcher/geo.py b/src/floatmatcher/geo.py new file mode 100644 index 0000000..879f585 --- /dev/null +++ b/src/floatmatcher/geo.py @@ -0,0 +1,20 @@ +# 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]) + + 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/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/src/floatmatcher/orchestrator.py b/src/floatmatcher/orchestrator.py new file mode 100644 index 0000000..0c82665 --- /dev/null +++ b/src/floatmatcher/orchestrator.py @@ -0,0 +1,106 @@ +# orchestrator.py: the top-level coordinator (public API of the library). + + +from collections.abc import Sequence + +import numpy as np + +from .gridset import GridSet +from .matchup import NearestNeighbor +# from .interpolation import Interpolation +from .pointset import PointSet +from .results import MatchupResult +from .products import Product +from .flatgrid import FlatGrid +from .utils import _select_variables +from .neighbors import spatial_nearest, temporal_nearest + +class Orchestrator: + """Orchestrator gather all Points Pointset(), variables needed, setup Product(). + It prepares all elements to call the method in .match() function. + All constraints and parametrization lives in .match() + """ + + def __init__(self, points: PointSet, + variables: str | list[str] | None, product: Product) -> None: + self.points = points + self.variables = variables + self.product = product + self._files: list[str] | None = None + + @property + def files(self) -> list[str]: + if self._files is None: + # always give points just not used in case of ExplicitFiles resolver + self._files = self.product.files_for(self.points) + return self._files + + def match(self, method: NearestNeighbor) -> MatchupResult: + if isinstance(method, NearestNeighbor): + return self._match_nearest(method) + # elif isinstance(method, Interpolation): + # return self._match_interp() + else: + raise AttributeError("unknown method") + + def _match_nearest(self, method: NearestNeighbor) -> MatchupResult: + print("matchin nearest method : no batching method yet") + files_to_process = self.files # property is only trigger when called the first time + grid_full = self._open_lazy_grid(files_to_process) # GridSet object + + # starting by lonlat2xy on spatial grid + flat_grid = FlatGrid.from_grid(grid_full.dataset) # return FlatGrid object + grid_stacked = flat_grid.xyz + + # starting Nearest method : apply kdtree on spatial + dist_km, spatial_idx = spatial_nearest(grid_stacked, self.points, k=method.k_nearest) + valid_spatial = dist_km <= method.max_dist_km + + idx_count = len(self.points.lon) + if grid_full.regime == "3D": + assert flat_grid.time is not None + time_delta, temporal_idx = temporal_nearest(flat_grid.time, + self.points, + k=method.k_nearest + ) + valid = valid_spatial & (time_delta <= method.max_time_seconds) + else: + time_delta = np.full(idx_count, np.nan) + temporal_idx = None + valid = valid_spatial + + # read ONLY at valid points: no wasted read for out-of-window points + # select indexes of spatial and time + idx = np.where(valid)[0] + node_idx = spatial_idx[idx] # grid node index of retained points + tsel_idx = None # set temporal case + if temporal_idx is not None: + tsel_idx = temporal_idx[idx] # apply on every node. filtering is made after (costless) + + # retreive data only at good positions : select in dataset stacked of FlatGrid object + picked = flat_grid.read_values(node_idx, tsel_idx) + + # scatter each variable's valid values back to full PointSet-length. + # picked is from _stacked which is flatten, not PointSet lenght :) + values = {} + for var, vals in picked.items(): + full = np.full(idx_count, np.nan) + full[idx] = vals + values[var] = full + + # invalid points carry no meaningful distance/time either + dist_out = np.full(idx_count, np.nan) + dist_out[idx] = dist_km[idx] + time_delta_out = np.full(idx_count, np.nan) + time_delta_out[idx] = time_delta[idx] + + return MatchupResult(values=values, distance_km=dist_out, + time_delta=time_delta_out, valid=valid, points=self.points) + + + def _open_lazy_grid(self, paths: Sequence[str]) -> GridSet: + """Open set of files -> normalized, variable selected and validated GridSet.""" + ds_raw = self.product.open_paths(paths) + ds = _select_variables(self.product.normalize(ds_raw), self.variables) + return GridSet(ds) + 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/src/floatmatcher/products.py b/src/floatmatcher/products.py new file mode 100644 index 0000000..4e61caf --- /dev/null +++ b/src/floatmatcher/products.py @@ -0,0 +1,112 @@ +import xarray as xr +from abc import ABC, abstractmethod +from pathlib import Path +from collections.abc import Sequence +from typing import Self +from dask.utils import SerializableLock + + +from .pointset import PointSet +from .resolver import FileResolver, ExplicitFiles, PathTemplate + + +_NETCDF_LOCK = SerializableLock() + +# ───────────────────────────────────────────────────────────── + +# ───────────────────────────────────────────────────────────── + +def to_standard(ds: xr.Dataset, mapping: dict[str, str]) -> xr.Dataset: + """ + Rename keys found in mapping. If key not found, skip and let it as it is. + set variables lon/lat/time as coords if not already + """ + + renamable = set(ds.variables) | set(ds.dims) + valid={} + for src, dst in mapping.items() : + if src in renamable: + valid[src] = dst + + ds = ds.rename(valid) + + # transform data var into coordinate if not already + to_promote = [] + for name in ("lon", "lat", "time"): + if name in ds.data_vars: + to_promote.append(name) + + ds = ds.set_coords(to_promote) + return ds + + + +class Product(ABC): + COORD_MAP: dict[str, str] = {} + + def __init__(self, resolver: FileResolver | None = None) -> None: + self.resolver = resolver + + @classmethod + def from_local(cls, + path: str | Path | Sequence[str] | None = None, + pattern: str | None = None) -> Self: + resolver: FileResolver + if pattern and not path: + raise ValueError("pattern given without a root path") + if not path: + raise ValueError("no path or pattern provided for opening data") + + if pattern: + # a pattern is substituted under ONE root, not a list of paths + if not isinstance(path, (str, Path)): + raise ValueError("a pattern needs a single root path, not a list") + resolver = PathTemplate(path, pattern) + else: + resolver = ExplicitFiles(path) + + return cls(resolver) + + def files_for(self, points: PointSet | None = None) -> list[str]: + if self.resolver is None: + raise ValueError( + "this product carries no source; build it with " + f"{type(self).__name__}.from_local(...) before resolving files" + ) + return self.resolver.files_for(points) + + def open_paths(self, paths: Sequence[str]) -> xr.Dataset: + """ open several files as a xr mf_dataset with options + // lock option fix made by AI to avoid segfault : open_mfdataset + returns dask arrays read by multiple threads, but netCDF4/HDF5 + isn't thread-safe here: concurrent access to xarray's file cache + corrupted memory and killed the process. A single shared lock serialises + those reads // + """ + if not paths: + raise ValueError("empty file list") + return xr.open_mfdataset( + paths, + combine="by_coords", + lock=_NETCDF_LOCK, + engine="netcdf4", + compat="override", + coords="minimal", + ) + + @abstractmethod + def normalize(self, ds_raw: xr.Dataset) -> xr.Dataset: + ... + + +class ERA5Product(Product): + COORD_MAP = {"longitude": "lon", "latitude": "lat", "valid_time": "time"} + def normalize(self, raw: xr.Dataset) -> xr.Dataset: + ds = to_standard(raw, self.COORD_MAP) + return ds + +class LUTProduct(Product): + COORD_MAP = {"lon": "lon", "lat": "lat"} + def normalize(self, raw: xr.Dataset) -> xr.Dataset: + ds = to_standard(raw, self.COORD_MAP) + return ds \ No newline at end of file diff --git a/src/floatmatcher/profile_loader.py b/src/floatmatcher/profile_loader.py new file mode 100644 index 0000000..b125f7f --- /dev/null +++ b/src/floatmatcher/profile_loader.py @@ -0,0 +1,105 @@ +from typing import Any + +import pandas as pd +import xarray as xr +from numpy.typing import NDArray + +from .pointset import PointSet +from .exceptions import ProfileFormatError + +# ───────────────────────────────────────────────────────────── +# Private xarray extraction helpers +# (the name-lookup logic lives HERE, never in PointSet) +# ───────────────────────────────────────────────────────────── + +def _find_key(ds: xr.Dataset, name: str) -> str: + if name in ds.coords or name in ds.variables: + return name + raise ProfileFormatError(f"No dataset name matches: {name}") + + +def _get(ds: xr.Dataset, name: str) -> xr.DataArray: + """Return the DataArray object""" + return ds[_find_key(ds, name)] + + +def _extract(ds: xr.Dataset, name: str) -> NDArray[Any]: + """Return the .values (raw ndarray) — the common case.""" + return _get(ds, name).values + + +# ───────────────────────────────────────────────────────────── +# ProfileLoader: converges any point source into a PointSet. +# Never validates itself: it extracts and delegates to PointSet. +# ───────────────────────────────────────────────────────────── + +class ProfileLoader: + + # --- 1. Raw arrays: simplest case --- + @staticmethod + def from_arrays(lon: Any, lat: Any, time: Any) -> PointSet: + # Nothing to extract, no origin_* (bare arrays). + return PointSet(lon, lat, time) + + # --- 2. pandas DataFrame --- + @staticmethod + def from_dataframe(df: pd.DataFrame, lon: str = "longitude", + lat: str = "latitude", time: str = "date") -> PointSet: + # Stable, named columns -> NO multi-name lookup here. + return PointSet(df[lon], df[lat], df[time], + origin_dim=df.index.name or "index") + + # --- 3. xarray Dataset --- + @staticmethod + def from_xrdataset(ds: xr.Dataset, lon: str = "LONGITUDE", + lat: str = "LATITUDE", time: str = "TIME") -> PointSet: + # The point-dimension name is DISCOVERED (da_lon.dims[0]), never assumed. + da_lon = _get(ds, lon) + point_dim = da_lon.dims[0] # DISCOVERED, never assumed + + # 2. extract the three arrays as ndarrays + lon_arr = da_lon.values + lat_arr = _extract(ds, lat) + time_arr = None + candidates = [time, "JULD"] + for name in candidates: + if name in ds.coords or name in ds.variables: + time_arr = _extract(ds, name) + break + if time_arr is None: + raise ProfileFormatError( + f"No dataset name matches time candidates: {candidates}") + + + return PointSet(lon=lon_arr, + lat=lat_arr, + time=time_arr, + origin_dim=str(point_dim), + origin_ds=ds + ) + + # # --- 4. argopy float: an xarray Dataset from DataFetcher().load().data --- + # @staticmethod + # def from_argopy_float(ds: xr.Dataset) -> PointSet: + # """argopy float Dataset -> PointSet, ONE position per profile. + + # argopy returns a measurement-level layout (dim ``N_POINTS``: lon/lat/time + # repeat across each profile's depth levels). Colocalization wants one + # position per profile, so collapse to the profile layout (dim ``N_PROF``) + # when a measurement layout is given. + # """ + # if "N_POINTS" in ds.dims: + # ds = _to_profiles(ds) # N_POINTS -> (N_PROF, N_LEVELS) + # return ProfileLoader.from_xrdataset(ds, lon="LONGITUDE", lat="LATITUDE", + # time="TIME") + + # # --- 5. argopy index: a pandas DataFrame, one row per profile --- + # @staticmethod + # def from_argopy_index(index: pd.DataFrame) -> PointSet: + # """argopy index DataFrame -> PointSet. + + # IndexFetcher().load().index is a DataFrame with longitude/latitude/date, + # one row per profile: exactly what from_dataframe consumes. + # """ + # return ProfileLoader.from_dataframe(index, lon="longitude", + # lat="latitude", time="date") diff --git a/src/floatmatcher/resolver.py b/src/floatmatcher/resolver.py new file mode 100644 index 0000000..61ede74 --- /dev/null +++ b/src/floatmatcher/resolver.py @@ -0,0 +1,148 @@ +# resolver.py: +# +# The resolver return a list of files to open, it doesn't open anything. +# The resolver is split into a dummy class ExplicitFiles returning all files listed +# or all files under a given path and a smarter one returning all files sorted along +# a given pattern given by the user. + +import warnings +from abc import ABC, abstractmethod +from collections.abc import Iterable +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd + +from .pointset import PointSet + +# ───────────────────────────────────────────────────────────── +# Path resolution (pure helper, no disk access) +# ───────────────────────────────────────────────────────────── + +def resolve_path(root: str | Path, pattern: str, date: Any) -> str: + """Substitute a date's components into `pattern` (under `root`). available fields are + year, month, day, hour, minute (pick what the tree needs, with padding ex: + "{year}/{month:02d}/era5_{year}{month:02d}{day:02d}.nc"). + """ + ts = pd.Timestamp(date) + rel = pattern.format( + year=ts.year, month=ts.month, day=ts.day, hour=ts.hour, minute=ts.minute + ) + return str(Path(root) / rel) + + +# ───────────────────────────────────────────────────────────── +# Resolvers: which files to return ? +# ───────────────────────────────────────────────────────────── + +class FileResolver(ABC): + """Answers: which files must be opened for these points?""" + + @abstractmethod + def files_for(self, points: PointSet | None = None) -> list[str]: + ... + + +class ExplicitFiles(FileResolver): + """Simplest resolver: the user lists the files explicitly + or a parent directory containing all files to process. + """ + + def __init__(self, paths: str | Path | Iterable[str | Path], + pattern: str = "*.nc") -> None: + if isinstance(paths, (str, Path)): + self._paths: list[str | Path] = [paths] + else: + self._paths = list(paths) + self.pattern = pattern + + def files_for(self, points: PointSet | None = None) -> list[str]: + resolved: list[str] = [] + for entry in self._paths: + path = Path(entry) + if path.is_dir(): + found = [] + for child in path.rglob(self.pattern): + found.append(str(child)) + found.sort() + resolved.extend(found) + else: + resolved.append(str(path)) + return resolved + + +class PathTemplate(FileResolver): + """Resolve files from a declared directory pattern (ex: Datarmor). + + The user declares the shape of their tree once; for each date PRESENT in the + points, one path is built by substituting the date into the pattern. Only + the dates actually needed are resolved (no directory scan). The field type + (single-levels_inst, pressure-levels, ...) is a literal part of the pattern + chosen by the user. + + Missing files are filtered out (with a warning): the user's time tolerance + in the matchup, not the resolver, decides what gap is acceptable. But if + NOTHING resolves to an existing file, that is a configuration error -> raise. + """ + + def __init__(self, root: str | Path, pattern: str, + granularity: Literal["Y", "M", "D", "h", "m", "s"] = "D", + pad: int = 1) -> None: + self.root = root + self.pattern = pattern + self.granularity = granularity + self.pad = pad + + def files_for(self, points: PointSet | None = None) -> list[str]: + if points is None: + raise ValueError( + "PathTemplate resolves one file per date present in the points; " + "pass a PointSet, or use ExplicitFiles to list a directory." + ) + + dates = np.unique( + np.asarray(points.time).astype(f"datetime64[{self.granularity}]") + ) + + # A point at 23h or 1h AM may well find its nearest step in the NEXT or PREVISOU file, so + # the neighbours of every date are requested too. + step = np.timedelta64(1, self.granularity) + wanted = dates + for shift in range(1, self.pad + 1): + wanted = np.concatenate([wanted, dates - shift * step, dates + shift * step]) + wanted = np.unique(wanted) + + resolved = [] + for d in wanted: + resolved.append(resolve_path(self.root, self.pattern, d)) + + # Only the dates of the requested points are worth warning about: + # padded neighbours are absent at the edges of any archive and does + # not raise warnings. + asked_for = set() + for d in dates: + asked_for.add(resolve_path(self.root, self.pattern, d)) + + existing = [] + missing = [] + for p in resolved: + if Path(p).exists(): + existing.append(p) + elif p in asked_for: + missing.append(p) + + if missing: + warnings.warn( + f"PathTemplate: {len(missing)} expected file(s) not found and " + f"ignored (e.g. {missing[0]}). The matchup's time tolerance decides " + f"whether the remaining dates are close enough." + ) + if not existing: + example = resolved[0] if resolved else "" + raise FileNotFoundError( + "PathTemplate: no file found for any of the requested dates. " + f"Check `root` and `pattern` (a path resolved e.g. to: {example})." + ) + return existing + 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/src/floatmatcher/utils.py b/src/floatmatcher/utils.py new file mode 100644 index 0000000..343f783 --- /dev/null +++ b/src/floatmatcher/utils.py @@ -0,0 +1,22 @@ +# utils.py + +# All general methods are grouped here to avoid "@staticmethod" +# in classes and lightning the writings. Furthermore, reusable +# in other files without whole class import + +import xarray as xr + + +def _select_variables(ds: xr.Dataset, variables: str | list[str] | None = None) -> xr.Dataset: + """Keep only the requested variables. None -> keep all""" + if variables is None: + return ds + + wanted = [variables] if isinstance(variables, str) else list(variables) + missing = [v for v in wanted if v not in ds.data_vars] + if missing: + raise ValueError( + f"variables not found in the dataset: {missing}. " + f"Available: {list(ds.data_vars)}." + ) + return ds[wanted] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..21d9f8f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,82 @@ +# 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 + + +# 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 + + + +# ---------- 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 + + +# ---------- fake era5 ---------- + +@pytest.fixture +def era5_files(tmp_path): + """Five daily ERA5-like files (2015/06/01 to 2015/06/05). + first timestep : + lon=9 lon=10 lon=11 lon=350 + lat=40 1.03 11.03 21.03 31.03 + lat=30 101.03 111.03 121.03 131.03 + lat=20 201.03 211.03 221.03 231.03 + """ + lat = [40.0, 30.0, 20.0] + lon = [9.0, 10.0, 11.0, 350.0] + sst = np.array([[ 0., 10., 20., 30.], + [100., 110., 120., 130.], + [200., 210., 220., 230.]]) + + paths = [] + for day in range(1, 6): + times = np.array([f"2015-06-{day:02d}T03", f"2015-06-{day:02d}T23"], + dtype="datetime64[ns]") + shape_lst = list(np.shape(sst)) + shape_lst.append(len(times)) + shape = tuple(shape_lst) + field = np.zeros(shape) + + for t in range(len(times)) : + h = times[t].astype("datetime64[h]").astype(int) % 24 /100 + field[:,:,t] = sst+day+h + dims = ("latitude", "longitude", "valid_time") + ds = xr.Dataset( + {"sst": (dims, field), "t2m": (dims, field + 1000.0)}, + coords={"latitude": lat, "longitude": lon, "valid_time": times}, + ) + path = tmp_path / f"era5_201506{day:02d}.nc" + ds.to_netcdf(path) + paths.append(str(path)) + + return paths diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..0fdf7e6 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,39 @@ +# 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.""" + 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)) + +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_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 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) + 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" 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() diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py new file mode 100644 index 0000000..c87dba3 --- /dev/null +++ b/tests/test_orchestrator.py @@ -0,0 +1,71 @@ +# tests/test_orchestrator.py +# + +import numpy as np +import numpy.testing as npt + +from floatmatcher.matchup import NearestNeighbor +from floatmatcher.orchestrator import Orchestrator +from floatmatcher.products import ERA5Product +from floatmatcher.profile_loader import ProfileLoader + + +def test_match_nearest_get_good_node(era5_files): + points = ProfileLoader.from_arrays( + lon=[11.1, 9.05, 349.9], + lat=[30.05, 39.95, 20.1], + time=np.array(["2015-06-01T02", "2015-06-02T00", "2015-06-02T22"], dtype="datetime64[ns]"), + ) + product = ERA5Product.from_local(path=era5_files) + + res = Orchestrator(points, "sst", product).match(NearestNeighbor()) + + assert res.valid.all() + assert set(res.values) == {"sst"} # only sst has been taken into account + npt.assert_allclose(res.values["sst"], [121.03, 1.23, 232.23], atol=1e-4) + assert (res.distance_km < 25).all() + assert (res.time_delta < 5 * 3600).all() # seconds - 5h + assert len(np.unique(res.distance_km)) == 3 + + +def test_match_nearest_rejects_out_of_range_points(era5_files): + points = ProfileLoader.from_arrays( + lon=[11.1, 230, 11.1], + lat=[30.05, -40.0, 30.05], + time=np.array(["2015-06-01T02", # p0 in range + "2015-06-01T02", # p1 far away in space + "2015-09-01T02"], # p2 far away in time + dtype="datetime64[ns]"), + ) + product = ERA5Product.from_local(path=era5_files) + orch = Orchestrator(points, "sst", product) + + res = orch.match(NearestNeighbor()) + + assert res.valid.tolist() == [True, False, False] + npt.assert_allclose(res.values["sst"][:1], [121.03], atol=1e-4) + assert np.isnan(res.values["sst"][1:]).all() # no value for the rejected + assert np.isnan(res.distance_km[1:]).all() # nor distance + assert np.isnan(res.time_delta[1:]).all() # nor time gap + + # p1 was out on DISTANCE only + loose_dist = orch.match(NearestNeighbor(max_dist_km=100000)) + assert loose_dist.valid.tolist() == [True, True, False] + + # p2 was out on TIME only + loose_time = orch.match(NearestNeighbor(max_time=np.timedelta64(300, "D"))) + assert loose_time.valid.tolist() == [True, False, True] + + +def test_match_nearest_over_two_variable(era5_files): + points = ProfileLoader.from_arrays( + lon=[11.1], lat=[30.05], + time=np.array(["2015-06-01T02"], dtype="datetime64[ns]"), + ) + product = ERA5Product.from_local(path=era5_files) + + res = Orchestrator(points, ["sst", "t2m"], product).match(NearestNeighbor()) + + assert set(res.values) == {"sst", "t2m"} + npt.assert_allclose(res.values["sst"], [121.03], atol=1e-4) + npt.assert_allclose(res.values["t2m"], [1121.03], atol=1e-4) 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 + + diff --git a/tests/test_products.py b/tests/test_products.py new file mode 100644 index 0000000..81c836d --- /dev/null +++ b/tests/test_products.py @@ -0,0 +1,114 @@ +# tests/test_products.py + +import numpy as np +import numpy.testing as npt +import xarray as xr +import pytest + +from floatmatcher.resolver import PathTemplate, ExplicitFiles +from floatmatcher.products import ( + to_standard, + ERA5Product, + LUTProduct, +) + + +def _era5_raw(): + """Raw ERA5-like dataset with source names longitude/latitude/time.""" + return xr.Dataset( + {"t2m": (("valid_time", "latitude", "longitude"), np.zeros((2, 3, 3)))}, + coords={ + "valid_time": np.array(["2015-01-01", "2015-01-02"], dtype="datetime64[ns]"), + "latitude": [0.0, 1.0, 2.0], + "longitude": [10.0, 11.0, 12.0], + }, + ) + + +# ───────────── to_standard ───────────── + +def test_rename_maps_present_keys(): + out = to_standard(_era5_raw(), {"longitude": "lon", "latitude": "lat", "valid_time": "time"}) + assert "lon" in out.coords and "lat" in out.coords and "time" in out.coords + + +def test_rename_ignores_absent_keys(): + out = to_standard(_era5_raw(), {"longitude": "lon", "foo": "bar"}) + assert "lon" in out.coords + assert "bar" not in out.variables # 'foo' does not exist -> skipped, not an error + + +def test_rename_mapping_already_well_named(): + ds = xr.Dataset({"chl": (("lat", "lon"), np.zeros((2, 2)))}, + coords={"lat": [0.0, 1.0], "lon": [10.0, 11.0]}) + out = to_standard(ds, {"lon": "lon", "lat": "lat"}) + assert "lon" in out.coords and "lat" in out.coords + + +def test_rename_bare_dimension(): + # a dimension without a coordinate can still be renamed + ds = xr.Dataset({"v": (("obs",), [1.0, 2.0, 3.0])}) + out = to_standard(ds, {"obs": "points"}) + assert "points" in out.dims and "obs" not in out.dims + + +def test_data_variable_is_promoted_to_coord(): + """lon/lat carried as data variables (anonymous dims) become coordinates: + GridSet reads them from ds.coords, never from ds.data_vars.""" + ds = xr.Dataset( + { + "sst": (("y", "x"), np.zeros((2, 2))), + "latitude": (("y",), np.array([0.0, 1.0])), + "longitude": (("x",), np.array([10.0, 11.0])), + }, + ) + assert "latitude" in ds.data_vars # not a coord at this point + + out = to_standard(ds, {"longitude": "lon", "latitude": "lat"}) + + assert "lon" in out.coords and "lat" in out.coords + assert "lon" not in out.data_vars and "lat" not in out.data_vars + npt.assert_allclose(out["lon"].values, [10.0, 11.0]) + + +def test_rename_does_not_affect_input(): + ds = _era5_raw() + _ = to_standard(ds, {"longitude": "lon"}) + assert "longitude" in ds.coords # original left untouched + + + +# ───────────── Products ───────────── + +def test_era5_normalize(): + out = ERA5Product().normalize(_era5_raw()) + assert "lon" in out.coords and "lat" in out.coords and "time" in out.coords + + +def test_lut_normalize(): + ds = xr.Dataset({"chl": (("lat", "lon"), np.zeros((2, 2)))}, + coords={"lat": [0.0, 1.0], "lon": [-10.0, -9.0]}) + out = LUTProduct().normalize(ds) + assert "lon" in out.coords and "lat" in out.coords + + +def test_from_local_with_pattern_builds_a_pathtemplate(): + product = ERA5Product.from_local(path="/data", pattern="{year}/x.nc") + assert isinstance(product.resolver, PathTemplate) + + +def test_from_local_without_pattern_builds_explicitfiles(): + product = ERA5Product.from_local(path="/data/2018") + assert isinstance(product.resolver, ExplicitFiles) + + +def test_from_local_returns_the_good_product(): + assert isinstance(ERA5Product.from_local(path="/data"), ERA5Product) + assert isinstance(LUTProduct.from_local(path="/data"), LUTProduct) + + +def test_from_local_without_path_raise_error(): + with pytest.raises(ValueError): + ERA5Product.from_local(pattern="{year}/x.nc") # pattern sans path + with pytest.raises(ValueError): + ERA5Product.from_local() \ No newline at end of file diff --git a/tests/test_profile_loader.py b/tests/test_profile_loader.py new file mode 100644 index 0000000..2f40cae --- /dev/null +++ b/tests/test_profile_loader.py @@ -0,0 +1,186 @@ +# Tests for ProfileLoader and helpers. + + +import numpy as np +import numpy.testing as npt +import pandas as pd +import pytest +import xarray as xr + +from floatmatcher.profile_loader import ( + ProfileLoader, + _find_key, + _get, + _extract, +) +from floatmatcher.exceptions import ProfileFormatError + + +@pytest.fixture +def argopy_like_ds(): + """xarray Dataset in the argopy N_PROF layout, with a real coordinate + on the point dimension -> exercises label provenance.""" + return xr.Dataset( + { + "LONGITUDE": (("N_PROF",), np.array([-45.0, -44.0, -43.0])), + "LATITUDE": (("N_PROF",), np.array([32.0, 33.0, 34.0])), + "TIME": (("N_PROF",), np.array( + ["2015-01-01", "2015-01-02", "2015-01-03"], dtype="datetime64[ns]")), + }, + coords={"N_PROF": np.arange(3)}, + ) + + +@pytest.fixture +def juld_ds(): + """xarray Dataset whose time variable is named JULD, not TIME + -> exercises the TIME/JULD tolerance.""" + return xr.Dataset( + { + "LONGITUDE": (("obs",), np.array([-45.0, -44.0, -43.0])), + "LATITUDE": (("obs",), np.array([32.0, 33.0, 34.0])), + "JULD": (("obs",), np.array( + ["2015-01-01", "2015-01-02", "2015-01-03"], dtype="datetime64[ns]")), + }, + ) + + +@pytest.fixture +def df_datetime_index(): + """DataFrame whose index is a DatetimeIndex (non-integer)""" + idx = pd.to_datetime(["2015-01-01", "2015-01-02", "2015-01-03"]) + idx.name = "profile_date" + return pd.DataFrame( + { + "longitude": [-45.0, -44.0, -43.0], + "latitude": [32.0, 33.0, 34.0], + "date": idx, + }, + index=idx, + ) + + + +# ───────────────────────────────────────────────────────────── +# Helpers: _find_key / _get / _extract +# ───────────────────────────────────────────────────────────── + +def test_find_key_matches_variable_or_coord(argopy_like_ds): + # variable match + assert _find_key(argopy_like_ds, "LONGITUDE") == "LONGITUDE" + # coordinate match (N_PROF is a coord here) + assert _find_key(argopy_like_ds, "N_PROF") == "N_PROF" + + +def test_find_key_raises_when_no_candidate_matches(argopy_like_ds): + with pytest.raises(ProfileFormatError): + _find_key(argopy_like_ds, "DUMMY") + + +def test_get_returns_dataarray_extract_returns_ndarray(argopy_like_ds): + da = _get(argopy_like_ds, "LONGITUDE") + assert isinstance(da, xr.DataArray) # object, has .dims + arr = _extract(argopy_like_ds, "LONGITUDE") + assert isinstance(arr, np.ndarray) # raw values + + +# ───────────────────────────────────────────────────────────── +# from_arrays +# ───────────────────────────────────────────────────────────── + +def test_from_arrays_builds_without_provenance(): + ps = ProfileLoader.from_arrays( + lon=[-45.0, -44.0], + lat=[32.0, 33.0], + time=np.array(["2015-01-01", "2015-01-02"], dtype="datetime64[ns]"), + ) + npt.assert_allclose(ps.lon, [-45.0, -44.0]) + assert ps.origin_dim is None + assert ps.origin_ds is None + + +def test_from_arrays_length_mismatch_raises(): + # PointSet enforces equal lengths; the loader should surface that. + with pytest.raises(ValueError): + ProfileLoader.from_arrays( + lon=[-45.0, -44.0, -43.0], + lat=[32.0, 33.0], # shorter on purpose + time=np.array(["2015-01-01", "2015-01-02"], dtype="datetime64[ns]"), + ) + + +# ───────────────────────────────────────────────────────────── +# from_dataframe +# ───────────────────────────────────────────────────────────── + +def test_from_dataframe_uses_the_index_name(df_datetime_index): + """origin_dim comes from df.index.name""" + ps = ProfileLoader.from_dataframe(df_datetime_index) + + assert ps.origin_dim == "profile_date" + npt.assert_allclose(ps.lon, [-45.0, -44.0, -43.0]) + npt.assert_allclose(ps.lat, [32.0, 33.0, 34.0]) + npt.assert_array_equal(ps.time, + np.array(["2015-01-01", "2015-01-02", "2015-01-03"], + dtype="datetime64[ns]")) + + +def test_from_dataframe_unnamed_index_falls_back_to_index(): + df = pd.DataFrame({ + "longitude": [-45.0, -44.0], + "latitude": [32.0, 33.0], + "date": pd.to_datetime(["2015-01-01", "2015-01-02"]), + }) # default RangeIndex, name is None + ps = ProfileLoader.from_dataframe(df) + assert ps.origin_dim == "index" + + +# ───────────────────────────────────────────────────────────── +# from_xrdataset +# ───────────────────────────────────────────────────────────── + +def test_from_xrdataset_time_juld_tolerance(juld_ds): + # time defaults to "TIME"; JULD must still be found as a fallback. + ps = ProfileLoader.from_xrdataset(juld_ds) + npt.assert_array_equal(ps.time, + np.array(["2015-01-01", "2015-01-02", "2015-01-03"], + dtype="datetime64[ns]")) + + +def test_from_xrdataset_user_can_override_names(): + ds = xr.Dataset({ + "my_lon": (("obs",), np.array([-45.0, -44.0])), + "my_lat": (("obs",), np.array([32.0, 33.0])), + "my_time": (("obs",), np.array( + ["2015-01-01", "2015-01-02"], dtype="datetime64[ns]")), + }) + ps = ProfileLoader.from_xrdataset(ds, lon="my_lon", lat="my_lat", time="my_time") + npt.assert_allclose(ps.lon, [-45.0, -44.0]) + + +def test_from_xrdataset_missing_coordinate_raises(): + ds = xr.Dataset({ + "LONGITUDE": (("obs",), np.array([-45.0, -44.0])), + # LATITUDE missing on purpose + "TIME": (("obs",), np.array( + ["2015-01-01", "2015-01-02"], dtype="datetime64[ns]")), + }) + with pytest.raises(ProfileFormatError): + ProfileLoader.from_xrdataset(ds) + + +# ───────────── origin_ds: the provenance the reinjection relies on ───────────── + +def test_from_xrdataset_carries_the_original_dataset(argopy_like_ds): + ps = ProfileLoader.from_xrdataset(argopy_like_ds) + + assert ps.origin_ds is argopy_like_ds + assert ps.origin_dim == "N_PROF" + + +def test_from_dataframe_has_no_original_dataset(df_datetime_index): + """A DataFrame is not a Dataset: origin_dim is known, origin_ds is not.""" + ps = ProfileLoader.from_dataframe(df_datetime_index) + + assert ps.origin_ds is None + assert ps.origin_dim == "profile_date" diff --git a/tests/test_resolver.py b/tests/test_resolver.py new file mode 100644 index 0000000..254be09 --- /dev/null +++ b/tests/test_resolver.py @@ -0,0 +1,140 @@ +# tests/test_resolver.py + + + +import warnings +from pathlib import Path + +import numpy as np +import pytest + +from floatmatcher.resolver import resolve_path, ExplicitFiles, PathTemplate +from floatmatcher.pointset import PointSet + + +PATTERN = "{year}/{month:02d}/era5_{year}{month:02d}{day:02d}.nc" + +@pytest.fixture +def points_object(): + """A 3-point PointSet carrying provenance (as if extracted from a dataset).""" + return PointSet( + lon=[-45.0, -44.0, -43.0], + lat=[32.0, 33.0, 34.0], + time=np.array(["2015-01-01", "2015-01-02", "2015-01-03"], dtype="datetime64[ns]"), + origin_dim="N_POINTS", + ) + + +def _make_file(root: Path, y, m, d): + """Create an empty file at the templated path.""" + p = root / f"{y}" / f"{m:02d}" / f"era5_{y}{m:02d}{d:02d}.nc" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text("") + return p + + +def _pts(*timestamps): + """A PointSet at those instants; positions are irrelevant here.""" + n = len(timestamps) + return PointSet(lon=[0.0] * n, lat=[0.0] * n, + time=np.array(timestamps, dtype="datetime64[ns]")) + + +# ───────────── resolve_path (pure, no disk) ───────────── + +def test_resolve_path(): + p = resolve_path("/data", PATTERN, np.datetime64("2015-01-01")) + assert Path(p) == Path("/data/2015/01/era5_20150101.nc") + + +# ───────────── ExplicitFiles ───────────── + +def test_explicit_single_path_becomes_list(): + assert ExplicitFiles("a.nc").files_for() == ["a.nc"] + + +def test_explicit_list_is_returned_directly(): + assert ExplicitFiles(["a.nc", "b.nc"]).files_for() == ["a.nc", "b.nc"] + + +def test_explicitfiles_expands_a_directory(tmp_path): + for day in (1, 2, 3): # creates 3 files .nc with good pattern + _make_file(tmp_path, 2015, 1, day) + (tmp_path / "notes.txt").write_text("ignored") # add a .txt in files list + out = ExplicitFiles(str(tmp_path)).files_for() + assert len(out) == 3 # only 3 files matching .nc found + assert out == sorted(out) # list already sorted + assert [Path(f).name for f in out] == ["era5_20150101.nc", + "era5_20150102.nc", + "era5_20150103.nc"] + + +def test_explicitfiles_mixes_files_and_directories(tmp_path): + """A named file is kept as is, a directory is expanded, and both land in + the same list -- in the order the caller gave them.""" + loose = _make_file(tmp_path / "loose", 2016, 5, 9) + tree = tmp_path / "tree" + inner = _make_file(tree, 2015, 1, 1) + + out = ExplicitFiles([str(loose), str(tree)]).files_for() + + assert out == [str(loose), str(inner)] + + +# ───────────── PathTemplate (fake tree via tmp_path) ───────────── +# points_object has dates 2015-01-01 / 02 / 03. + +def test_pathtemplate_resolves_present_dates(tmp_path, points_object): + for d in (1, 2, 3): + _make_file(tmp_path, 2015, 1, d) + out = PathTemplate(str(tmp_path), PATTERN).files_for(points_object) + assert len(out) == 3 + for p in out : + assert Path(p).exists() + + +def test_pathtemplate_filters_missing_with_warning(tmp_path, points_object): + _make_file(tmp_path, 2015, 1, 1) # day 2 missing on purpose + _make_file(tmp_path, 2015, 1, 3) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + PathTemplate(str(tmp_path), PATTERN).files_for(points_object) + assert len(w) == 1 # only 1 warning "not found" + assert "era5_20150102.nc" in str(w[0].message) # explicit what is missig aka the 02/01/2015 + + +def test_pathtemplate_all_missing_raises(tmp_path, points_object): + with pytest.raises(FileNotFoundError): + PathTemplate(str(tmp_path), PATTERN).files_for(points_object) + + +def test_same_day_points_give_one_file_and_no_warning(tmp_path): + """ several points the SAME day -> a single file""" + _make_file(tmp_path, 2015, 1, 1) # neighbors doesnt exists + pts = _pts("2015-01-01T00", "2015-01-01T06", "2015-01-01T18") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + out = PathTemplate(str(tmp_path), PATTERN, pad=1).files_for(pts) + assert len(w) == 0 # no warning found even if neighbors are missing + assert [Path(f).name for f in out] == ["era5_20150101.nc"] # good file has beed found + + +# ───────────── padding ───────────── + +def test_pad_pulls_the_neighbouring_files(tmp_path): + """A point in the day may match in the next or previous file --> returned both""" + for day in (1, 2, 3): + _make_file(tmp_path, 2015, 1, day) + + out = PathTemplate(str(tmp_path), PATTERN, pad=1).files_for(_pts("2015-01-02")) + + assert [Path(f).name for f in out] == ["era5_20150101.nc", + "era5_20150102.nc", + "era5_20150103.nc"] + + +def test_pad_zero_keeps_only_the_point_dates(tmp_path): + for day in (1, 2, 3): + _make_file(tmp_path, 2015, 1, day) + out = PathTemplate(str(tmp_path), PATTERN, pad=0).files_for(_pts("2015-01-02T23")) + assert [Path(f).name for f in out] == ["era5_20150102.nc"] 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() diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..90d5d85 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,47 @@ +# tests/test_utils.py + +import numpy as np +import pytest +import xarray as xr + +from floatmatcher.utils import _select_variables + + +def _grid_4vars(): + return xr.Dataset( + {v: (("lat", "lon"), np.zeros((2, 2))) for v in ["u10", "v10", "t2m", "sst"]}, + coords={"lat": [0.0, 1.0], "lon": [10.0, 11.0]}, + ) + +# ───────────────────────────────────────────────────────────── +# _select_variables +# ───────────────────────────────────────────────────────────── + +def test_select_none_keeps_all(): + out = _select_variables(_grid_4vars(), None) + assert set(out.data_vars) == {"u10", "v10", "t2m", "sst"} + + +def test_select_single_string(): + out = _select_variables(_grid_4vars(), "sst") + assert list(out.data_vars) == ["sst"] + + +def test_select_list_of_variables(): + out = _select_variables(_grid_4vars(), ["sst", "t2m"]) + assert list(out.data_vars) == ["sst", "t2m"] + + +def test_select_unknown_raises_and_lists_available(): + with pytest.raises(ValueError) as e: + _select_variables(_grid_4vars(), "xxx") + msg = str(e.value) + assert "xxx" in msg # names the bad variable + assert "sst" in msg # and lists what IS available + + +def test_select_partial_unknown_raises(): + # one valid + one invalid -> still raises (all-or-nothing) + with pytest.raises(ValueError): + _select_variables(_grid_4vars(), ["sst", "nope"]) +