diff --git a/pyproject.toml b/pyproject.toml index e355c42..ad97f8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,8 @@ dependencies = [ "xarray", "scipy", "pandas", + "dask", + "netCDF4", ] [dependency-groups] 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/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