diff --git a/pyproject.toml b/pyproject.toml index 514cec2..e355c42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "numpy<2.5", # 2.5+ stubs need Python 3.12 syntax; "xarray", "scipy", + "pandas", ] [dependency-groups] @@ -42,7 +43,7 @@ disallow_untyped_defs = true exclude = ["^tests/"] [[tool.mypy.overrides]] -module = ["scipy.*"] +module = ["scipy.*", "pandas.*"] ignore_missing_imports = true [tool.ruff] 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/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"]