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"