Enrich

enrich

Time-major enrichment of many trajectory points at once.

Remote files are organised by time (a global field per hour for ARCO-ERA5, one message per variable/level/step for ARPEGE GRIB2), so the cost of a time slice is paid once for every flight that crosses it. Points of all flights are grouped by the interval [t0, t0 + step) of the source's native time step; the source says which two cubes bracket each interval (and, for forecasts, from which run). The points using one time slice are split along the source's own tiling (:mod:baro.met.tiling): each group of touching occupied tiles becomes one cube, fetched once over the region and pressure range of its points, with bounded look-ahead, and released as soon as the last interval using it is done. With a cache_dir, every byte read is kept (:mod:baro.met.rangecache), so the cost of a campaign does not depend on how its flights are split into calls.

EnrichResult dataclass

Interpolated fields (NaN where flags != 0) and run statistics.

Source code in packages/barograph-met/src/baro/met/enrich.py

@dataclass
class EnrichResult:
    """Interpolated fields (NaN where ``flags != 0``) and run statistics."""

    fields: dict[str, NDArray[np.float64]]
    flags: NDArray[np.uint8]
    stats: dict[str, Any] = field(default_factory=dict)

    def __getitem__(self, name: str) -> NDArray[np.float64]:
        return self.fields[name]

    @property
    def valid(self) -> NDArray[np.bool_]:
        """Mask of the points with a value (no flag set)."""
        return np.asarray(self.flags == Flag.OK)

valid property

Mask of the points with a value (no flag set).

Enricher

Interpolate a gridded source along trajectories.

Parameters:

Name	Type	Description	Default
source	Source	Gridded source, e.g. :func:baro.met.arco_model_levels or :class:baro.met.arpege.ArpegeSource.	required
cache_dir	str | Path | None	Optional on-disk cache of every byte read (immutable data only), shared between runs, calls and processes: a later call downloads only what no earlier call has read. It is attached to source (which then keeps using it, also through other enrichers).	None
mode	Mode	"flight" (default) reads only the region and levels the points need. "global" reads whole fields on every level (FastMeteo-like), so that one cache serves any later flight.	'flight'
domain	Domain | None	Fixed domain for every cube. By default the points of each time slice are split along the source's tiling, and each group gets the box of its own points.	None
margin_deg	float	Margin around the points when the domain is derived from them.	0.5
level_margin	int	Extra levels kept above and below the bracketing levels.	1
io_workers	int	Concurrent range requests, and cubes fetched ahead of interpolation.	16
vertical	VerticalMethod	"log" or "linear" in pressure; "actual" surface pressure or "standard" (101325 Pa, reproduces a fixed altitude->level table).	'log'
on_uncovered	Literal['warn', 'raise', 'flag']	What to do with points outside the period the source covers (e.g. before the start of an archive): "warn" (default) flags them OUT_OF_COVERAGE and emits a warning with the covered period, "raise" raises :class:DataUnavailableError, "flag" only flags.	'warn'

Source code in packages/barograph-met/src/baro/met/enrich.py

class Enricher:
    """Interpolate a gridded source along trajectories.

    Parameters
    ----------
    source:
        Gridded source, e.g. :func:`baro.met.arco_model_levels` or
        :class:`baro.met.arpege.ArpegeSource`.
    cache_dir:
        Optional on-disk cache of every byte read (immutable data only), shared
        between runs, calls and processes: a later call downloads only what no
        earlier call has read. It is attached to ``source`` (which then keeps
        using it, also through other enrichers).
    mode:
        ``"flight"`` (default) reads only the region and levels the points
        need. ``"global"`` reads whole fields on every level (FastMeteo-like),
        so that one cache serves any later flight.
    domain:
        Fixed domain for every cube. By default the points of each time slice
        are split along the source's tiling, and each group gets the box of
        its own points.
    margin_deg:
        Margin around the points when the domain is derived from them.
    level_margin:
        Extra levels kept above and below the bracketing levels.
    io_workers / prefetch:
        Concurrent range requests, and cubes fetched ahead of interpolation.
    vertical / surface_pressure:
        ``"log"`` or ``"linear"`` in pressure; ``"actual"`` surface pressure or
        ``"standard"`` (101325 Pa, reproduces a fixed altitude->level table).
    on_uncovered:
        What to do with points outside the period the source covers (e.g.
        before the start of an archive): ``"warn"`` (default) flags them
        ``OUT_OF_COVERAGE`` and emits a warning with the covered period,
        ``"raise"`` raises :class:`DataUnavailableError`, ``"flag"`` only flags.
    """

    def __init__(  # noqa: PLR0913 - keyword-only configuration of the public enricher
        self,
        source: Source,
        *,
        cache_dir: str | Path | None = None,
        mode: Mode = "flight",
        domain: Domain | None = None,
        margin_deg: float = 0.5,
        level_margin: int = 1,
        io_workers: int = 16,
        prefetch: int = 2,
        vertical: VerticalMethod = "log",
        surface_pressure: SurfacePressureMode = "actual",
        on_uncovered: Literal["warn", "raise", "flag"] = "warn",
    ) -> None:
        if on_uncovered not in ("warn", "raise", "flag"):
            raise ValueError(f"unknown on_uncovered {on_uncovered!r}")
        self.on_uncovered = on_uncovered
        if mode not in ("flight", "global"):
            raise ValueError(f"unknown mode {mode!r}")
        self.source = source
        self.cache = RangeCache(cache_dir) if cache_dir is not None else None
        if self.cache is not None:
            source.use_cache(self.cache)
        self.mode = mode
        self.domain = domain if domain is not None or mode == "flight" else Domain.whole_globe()
        self.margin_deg = margin_deg
        self.level_margin = _ALL_LEVELS if mode == "global" else level_margin
        self.io_workers = io_workers
        self.prefetch = max(prefetch, 1)
        self.vertical = vertical
        self.surface_pressure = surface_pressure
        self._lock = threading.Lock()

    # -- public API ----------------------------------------------------------------------
    def enrich(
        self,
        time_utc: ArrayLike,
        lat: ArrayLike,
        lon: ArrayLike,
        *,
        altitude_ft: ArrayLike | None = None,
        pressure_pa: ArrayLike | None = None,
    ) -> EnrichResult:
        """Fields at each point. Give barometric (pressure) altitude in ft or pressure in Pa."""
        if (altitude_ft is None) == (pressure_pa is None):
            raise ValueError("give exactly one of altitude_ft or pressure_pa")
        t = np.asarray(time_utc).astype("datetime64[ns]")
        la = np.asarray(lat, dtype=np.float64)
        lo = np.asarray(lon, dtype=np.float64)
        if pressure_pa is not None:
            p = np.asarray(pressure_pa, dtype=np.float64)
        else:
            p = pressure_from_altitude_ft(np.asarray(altitude_ft, dtype=np.float64))
        if not (t.shape == la.shape == lo.shape == p.shape) or t.ndim != 1:
            raise ValueError("inputs must be 1-D arrays of equal length")
        return self._run(t, la, lo, p)

    # -- internals -----------------------------------------------------------------------
    def _run(
        self,
        t: NDArray[np.datetime64],
        lat: NDArray[np.float64],
        lon: NDArray[np.float64],
        p: NDArray[np.float64],
    ) -> EnrichResult:
        started = time.perf_counter()
        before = [s.snapshot() for s in self.source.stats]
        n = t.size
        values = {name: np.full(n, np.nan) for name in self.source.fields}
        flags = np.full(n, Flag.INVALID_INPUT, dtype=np.uint8)
        idx = self._usable(t, lat, lon, p, flags)
        step = self.source.time_step
        t_ok = t[idx]
        t0 = t_ok - (t_ok - _EPOCH) % step
        w = (t_ok - t0) / step
        la, lo, pp = lat[idx], lon[idx], p[idx]
        points = _Points(w, la, lo, pp)
        intervals, requests = self._plan(t0, w, la, lo, pp)
        counters: dict[str, Any] = {
            "cubes": len(requests),
            "fetched": 0,
            "unavailable": 0,
            "cube_bytes": 0,
        }
        cached_before = self._cache_hits()
        with (
            ThreadPoolExecutor(self.io_workers, thread_name_prefix="traj-io") as io,
            ThreadPoolExecutor(self.prefetch, thread_name_prefix="traj-cube") as cubes,
        ):
            for iv, got in self._stream(intervals, requests, io, cubes, counters):
                for sel, r0, r1 in _pairs(iv):
                    cube1 = None if r1 < 0 else got[r1]
                    v, f = self._interpolate(got[r0], cube1, points.take(sel))
                    for name in values:
                        values[name][idx[sel]] = v[name]
                    flags[idx[sel]] = f
        after = [s.snapshot() for s in self.source.stats]
        for name in ("bytes", "requests", "throttled"):
            counters[name] = int(sum(a[name] - b[name] for a, b in zip(after, before, strict=True)))
        counters["cache_bytes"] = self._cache_hits() - cached_before
        counters["seconds"] = time.perf_counter() - started
        return EnrichResult(values, flags, counters)

    def _usable(
        self,
        t: NDArray[np.datetime64],
        lat: NDArray[np.float64],
        lon: NDArray[np.float64],
        p: NDArray[np.float64],
        flags: NDArray[np.uint8],
    ) -> NDArray[np.int64]:
        """Indices of the valid, covered points; flags the uncovered ones."""
        ok = ~np.isnat(t) & np.isfinite(lat) & np.isfinite(lon) & np.isfinite(p) & (p > 0)
        uncovered = self._check_coverage(t, ok)
        flags[uncovered] = Flag.OUT_OF_COVERAGE
        return np.flatnonzero(ok & ~uncovered)

    def _cache_hits(self) -> int:
        """Bytes served by the cache so far."""
        return 0 if self.cache is None else self.cache.stats.snapshot()["hit_bytes"]

    def _check_coverage(
        self, t: NDArray[np.datetime64], ok: NDArray[np.bool_]
    ) -> NDArray[np.bool_]:
        """Points whose bracketing interval falls outside the source's coverage."""
        coverage = self.source.coverage()
        uncovered = np.zeros(t.size, dtype=bool)
        if coverage is None or not ok.any():
            return uncovered
        start = np.datetime64(coverage[0], "ns")
        end = np.datetime64(coverage[1], "ns")
        uncovered[ok] = (t[ok] < start) | (t[ok] > end)
        n = int(uncovered.sum())
        if n == 0 or self.on_uncovered == "flag":
            return uncovered
        before = int((uncovered & (t < start)).sum())
        after = n - before
        span = f"{t[uncovered].min()} .. {t[uncovered].max()}"
        message = (
            f"{n} point(s) outside the period covered by {self.source.source_id} "
            f"({coverage[0]} .. {coverage[1]} UTC): {before} before, {after} after "
            f"(their times span {span}); they are flagged OUT_OF_COVERAGE"
        )
        if self.on_uncovered == "raise":
            raise DataUnavailableError(message)
        warnings.warn(message, UncoveredPointsWarning, skip_file_prefixes=(_PACKAGE_DIR,))
        return uncovered

    def _plan(
        self,
        t0: NDArray[np.datetime64],
        w: NDArray[np.float64],
        lat: NDArray[np.float64],
        lon: NDArray[np.float64],
        p: NDArray[np.float64],
    ) -> tuple[list[_Interval], list[_CubeRequest]]:
        """Intervals in time order, and the cube requests (one per occupied-tile group)."""
        groups = _group(t0)
        keys = {start: self.source.cube_keys(start) for start in groups}
        uses: dict[CubeKey, list[NDArray[np.int64]]] = {}
        for start, members in groups.items():
            k0, k1 = keys[start]
            uses.setdefault(k0, []).append(members)
            later = members[w[members] > 0]
            if later.size:
                uses.setdefault(k1, []).append(later)
        requests: list[_CubeRequest] = []
        lookup: dict[CubeKey, tuple[NDArray[np.int64], NDArray[np.int64]]] = {}
        for key, parts in uses.items():
            sel = np.sort(np.concatenate(parts))
            tiling = None if self.domain is not None else self.source.tiling(key)
            labels = partition(lat[sel], lon[sel], self.margin_deg, tiling)
            ids = np.empty(sel.size, dtype=np.int64)
            for label in range(int(labels.max()) + 1 if sel.size else 0):
                part = sel[labels == label]
                ids[labels == label] = len(requests)
                requests.append(self._request(key, lat[part], lon[part], p[part]))
            lookup[key] = (sel, ids)
        intervals = []
        for start, members in groups.items():
            k0, k1 = keys[start]
            req1 = np.full(members.size, -1, dtype=np.int64)
            needs_next = w[members] > 0
            if needs_next.any():
                req1[needs_next] = _request_ids(lookup[k1], members[needs_next])
            intervals.append(_Interval(start, members, _request_ids(lookup[k0], members), req1))
        return intervals, requests

    def _request(
        self,
        key: CubeKey,
        lat: NDArray[np.float64],
        lon: NDArray[np.float64],
        p: NDArray[np.float64],
    ) -> _CubeRequest:
        """Request for the cube ``key`` over the region and pressures of some points."""
        domain = self.domain or Domain.from_points(lat, lon, self.margin_deg)
        return _CubeRequest(key, domain, (float(p.min()), float(p.max())))

    def _stream(
        self,
        intervals: list[_Interval],
        requests: list[_CubeRequest],
        io: ThreadPoolExecutor,
        pool: ThreadPoolExecutor,
        counters: dict[str, Any],
    ) -> Iterator[tuple[_Interval, dict[int, Cube | None]]]:
        """Intervals with their cubes; cubes are fetched ahead and freed after last use."""

        def fetch(req: int) -> Future[Cube | None]:
            """Submit the download of one cube."""
            return pool.submit(self._download, requests[req], io, counters)

        ahead = _Prefetcher(intervals, fetch, self.prefetch)
        for iv in intervals:
            yield iv, ahead.cubes(iv)
            ahead.release(iv)

    def _download(
        self, req: _CubeRequest, io: ThreadPoolExecutor, counters: dict[str, Any]
    ) -> Cube | None:
        try:
            cube = self.source.fetch(
                req.key, req.domain, req.p_range, level_margin=self.level_margin, executor=io
            )
        except DataUnavailableError:
            self._count(counters, "unavailable")
            return None
        self._count(counters, "fetched")
        self._count(counters, "cube_bytes", cube.nbytes)
        return cube

    def _count(self, counters: dict[str, Any], name: str, amount: int = 1) -> None:
        with self._lock:
            counters[name] += amount

    def _interpolate(
        self,
        cube0: Cube | None,
        cube1: Cube | None,
        pts: _Points,
    ) -> tuple[dict[str, NDArray[np.float64]], NDArray[np.uint8]]:
        n = pts.w.size
        names = list(self.source.fields)
        if cube0 is None:
            return {k: np.full(n, np.nan) for k in names}, np.full(n, Flag.MISSING_DATA, np.uint8)
        opts = SampleOptions(self.vertical, self.surface_pressure)
        v0, flags = sample(cube0, pts.lat, pts.lon, pts.p, opts)
        later = pts.w > 0
        if not later.any():
            return v0, flags
        if cube1 is None:
            v1 = {k: np.full(int(later.sum()), np.nan) for k in names}
            f1 = np.full(int(later.sum()), Flag.MISSING_DATA, np.uint8)
        else:
            nxt = pts.take(later)
            v1, f1 = sample(cube1, nxt.lat, nxt.lon, nxt.p, opts)
        flags = flags.copy()
        flags[later] |= np.where(flags[later] == 0, f1, 0).astype(np.uint8)
        out = {}
        wl = pts.w[later]
        for name in names:
            blended = v0[name].copy()
            blended[later] = (1.0 - wl) * v0[name][later] + wl * v1[name]
            blended[flags != 0] = np.nan
            out[name] = blended
        return out, flags

enrich(time_utc, lat, lon, *, altitude_ft=None, pressure_pa=None)

Fields at each point. Give barometric (pressure) altitude in ft or pressure in Pa.

Source code in packages/barograph-met/src/baro/met/enrich.py

def enrich(
    self,
    time_utc: ArrayLike,
    lat: ArrayLike,
    lon: ArrayLike,
    *,
    altitude_ft: ArrayLike | None = None,
    pressure_pa: ArrayLike | None = None,
) -> EnrichResult:
    """Fields at each point. Give barometric (pressure) altitude in ft or pressure in Pa."""
    if (altitude_ft is None) == (pressure_pa is None):
        raise ValueError("give exactly one of altitude_ft or pressure_pa")
    t = np.asarray(time_utc).astype("datetime64[ns]")
    la = np.asarray(lat, dtype=np.float64)
    lo = np.asarray(lon, dtype=np.float64)
    if pressure_pa is not None:
        p = np.asarray(pressure_pa, dtype=np.float64)
    else:
        p = pressure_from_altitude_ft(np.asarray(altitude_ft, dtype=np.float64))
    if not (t.shape == la.shape == lo.shape == p.shape) or t.ndim != 1:
        raise ValueError("inputs must be 1-D arrays of equal length")
    return self._run(t, la, lo, p)

UncoveredPointsWarning

Bases: UserWarning

Some points fall outside the period the source covers.

Source code in packages/barograph-met/src/baro/met/enrich.py

class UncoveredPointsWarning(UserWarning):
    """Some points fall outside the period the source covers."""