diff --git a/pyproject.toml b/pyproject.toml index c43e7fa..9d4a8d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,9 +39,7 @@ warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true -[tool.mypy.overrides] -module = "tests.*" -disallow_untyped_defs = false +exclude = ["^tests/"] [tool.ruff] line-length = 110 diff --git a/src/floatmatcher/results.py b/src/floatmatcher/results.py new file mode 100644 index 0000000..a9cfa4d --- /dev/null +++ b/src/floatmatcher/results.py @@ -0,0 +1,36 @@ +# results.py: output object of matchup + +import numpy as np +import xarray as xr + +from dataclasses import dataclass +from numpy.typing import NDArray +from .pointset import PointSet + + +@dataclass +class MatchupResult: + values: dict[str, NDArray[np.float64]] + distance_km: NDArray[np.float64] + time_delta: NDArray[np.float64] + valid: NDArray[np.bool_] + points: PointSet + + def to_dataset(self) -> xr.Dataset: + """ + Reinject the colocalized values into the dataset + The source Dataset travels inside the PointSet + """ + ds = self.points.origin_ds + dim = self.points.origin_dim + if ds is None or dim is None: + raise ValueError( + "Cannot reinject: these points have no origin dataset " + "(they came from raw arrays). Use the MatchupResult directly." + ) + + out = ds.copy() + for k, v in self.values.items(): + out[f"{k}_coloc"] = (dim, v) + + return out diff --git a/tests/test_results.py b/tests/test_results.py new file mode 100644 index 0000000..992ba86 --- /dev/null +++ b/tests/test_results.py @@ -0,0 +1,68 @@ +# tests/test_results.py +import numpy as np +import pytest +import xarray as xr + +from floatmatcher.results import MatchupResult +from floatmatcher.pointset import PointSet + + +@pytest.fixture +def result(): + """ MatchupResult object - 3 profiles """ + ds = xr.Dataset( + {"LONGITUDE": (("N_PROF",), np.array([-40.0, -30.0, -20.0])), + "LATITUDE": (("N_PROF",), np.array([35.0, 36.0, 37.0])), + "TIME": (("N_PROF",), np.array(["2018-01-03", "2018-01-01", "2018-01-02"], + dtype="datetime64[ns]"))}, + coords={"N_PROF": np.array([10, 11, 12])}, + ) + points = PointSet( + lon=ds["LONGITUDE"].values, + lat=ds["LATITUDE"].values, + time=ds["TIME"].values, + origin_dim="N_PROF", + origin_ds=ds, + ) + return MatchupResult( + values={"dummy_variable": np.array([280.0, 281.0, np.nan])}, + distance_km=np.array([1.0, 2.0, np.nan]), + time_delta=np.array([0.1, 0.2, np.nan]), + valid=np.array([True, True, False]), + points=points, + ) + +# @pytest.fixture +# def argo_like_ds(): +# return xr.Dataset( +# {"LONGITUDE": (("N_PROF",), np.array([-40.0, -30.0, -20.0])), +# "LATITUDE": (("N_PROF",), np.array([35.0, 36.0, 37.0])), +# "TIME": (("N_PROF",), np.array(["2018-01-03", "2018-01-01", "2018-01-02"], +# dtype="datetime64[ns]"))}, +# coords={"N_PROF": np.array([10, 11, 12])}, +# ) + + +# ---------- to_dataset ---------- + +def test_to_dataset_adds_coloc_variable(result): + out = result.to_dataset() + assert out["dummy_variable_coloc"].dims == ("N_PROF",) + assert out["dummy_variable_coloc"].sel(N_PROF=10) == 280.0 + + +def test_to_dataset_does_not_mutate_the_source(result): + result.to_dataset() + assert "dummy_variable_coloc" not in result.points.origin_ds.data_vars + + +def test_to_dataset_without_provenance_raises(): + """Raw-array points carry no dataset, so they cannot be reinjected.""" + points = PointSet([-45.0, -44.0, -43.0], [32.0, 33.0, 34.0], + np.array(["2015-01-01", "2015-01-02", "2015-01-03"], + dtype="datetime64[ns]")) + res = MatchupResult(values={"dummy_variable": np.zeros(3)}, + distance_km=np.zeros(3), time_delta=np.zeros(3), + valid=np.ones(3, bool), points=points) + with pytest.raises(ValueError): + res.to_dataset()