Sources

sources

Gridded sources: which arrays, which vertical coordinate, how to build a cube.

GridSource

An hourly, regular lat/lon zarr store with a level dimension.

vertical lists all levels of the store's level axis, in the store's order (top to bottom). Cubes carry only the levels needed to bracket the requested pressures, read with block-level range requests. With a cache (:meth:use_cache), the field chunks of hours older than final_after are kept on disk; metadata and recent (ERA5T) hours are always re-read.

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

class GridSource:
    """An hourly, regular lat/lon zarr store with a level dimension.

    ``vertical`` lists all levels of the store's level axis, in the store's
    order (top to bottom). Cubes carry only the levels needed to bracket the
    requested pressures, read with block-level range requests. With a cache
    (:meth:`use_cache`), the field chunks of hours older than ``final_after``
    are kept on disk; metadata and recent (ERA5T) hours are always re-read.
    """

    def __init__(  # noqa: PLR0913 - keyword-only configuration of a public source
        self,
        source_id: str,
        store: ZarrStore,
        vertical: VerticalCoordinate,
        *,
        fields: dict[str, str] | None = None,
        level_array: str,
        surface_pressure: SurfacePressure | None = None,
        final_after: np.timedelta64 = FINAL_AFTER,
    ) -> None:
        self.source_id = source_id
        self.final_after = final_after
        self.store = store
        self.vertical = vertical
        self.fields = dict(fields or DEFAULT_FIELDS)
        self._cache_fields = set(self.fields.values())
        self.surface_pressure = surface_pressure
        if vertical.needs_surface_pressure and surface_pressure is None:
            raise ValueError("hybrid levels need a surface pressure array")
        store_levels = np.asarray(store.read_1d(level_array), dtype=np.int64)
        missing = np.setdiff1d(vertical.ids, store_levels)
        if missing.size:
            raise ValueError(f"{source_id}: levels {missing.tolist()} not in the store")
        # position on the store's level axis of each level of ``vertical``
        self._store_index = np.searchsorted(store_levels, vertical.ids)
        if not np.array_equal(store_levels[self._store_index], vertical.ids):
            raise ValueError(f"{source_id}: store level axis must be sorted top to bottom")
        self.lat = store.read_1d("latitude").astype(np.float64)
        self.lon = store.read_1d("longitude").astype(np.float64)
        self._epoch = _hour_epoch(store.array("time").attrs.get("units", ""))
        self._t0 = int(store.read_1d("time", 0, 1)[0])
        self._checked_time_chunks: set[int] = set()
        self._lock = threading.Lock()
        self._tiling: Tiling | None = None
        self._tiling_known = False

    def select_fields(self, names: Sequence[str]) -> GridSource:
        """Independent field view sharing stores and byte caches with this source.

        Only listed output fields are fetched. Cache eligibility retains all
        original fields, so attaching a cache through a view cannot narrow the
        parent's cache policy. Configure shared stores before concurrent use.
        """
        if not names or len(set(names)) != len(names):
            raise ValueError("fields must be a nonempty sequence of unique names")
        unknown = set(names) - self.fields.keys()
        if unknown:
            raise ValueError(f"unknown fields: {sorted(unknown)}")
        view = copy(self)
        view.fields = {name: self.fields[name] for name in names}
        view._tiling = None
        view._tiling_known = False
        return view

    @property
    def stats(self) -> list[TransferStats]:
        """Transfer counters of every underlying byte store."""
        stores = [self.store.store]
        if self.surface_pressure is not None and self.surface_pressure.store is not self.store:
            stores.append(self.surface_pressure.store.store)
        return [s.stats for s in stores]

    def time_index(self, hour: np.datetime64) -> int:
        """Index of ``hour`` on the store's time axis, verified against the time array.

        The axis is assumed hourly and contiguous; each time chunk is checked
        once (a few KiB) before its indices are trusted.
        """
        hours = int((np.datetime64(hour, "h") - self._epoch) / np.timedelta64(1, "h"))
        idx = hours - self._t0
        meta = self.store.array("time")
        if not 0 <= idx < meta.shape[0]:
            raise KeyError(f"{hour} outside the store's time axis")
        size = meta.chunks[0]
        chunk = idx // size
        with self._lock:
            known = chunk in self._checked_time_chunks
        if not known:
            start, stop = chunk * size, min((chunk + 1) * size, meta.shape[0])
            values = self.store.read_1d("time", start, stop)
            if not np.array_equal(values, np.arange(start, stop) + self._t0):
                raise RuntimeError(f"{self.source_id}: time axis is not hourly-contiguous")
            with self._lock:
                self._checked_time_chunks.add(chunk)
        return idx

    time_step = np.timedelta64(1, "h")

    def coverage(self) -> tuple[np.datetime64, np.datetime64] | None:
        """Published period, from the store's ``valid_time_start``/``valid_time_stop``.

        For ARCO-ERA5 the model-level store lags ERA5T (about three months).
        """
        attrs = self.store.attrs
        start, stop = attrs.get("valid_time_start"), attrs.get("valid_time_stop")
        if not start or not stop:
            return None
        begin = np.datetime64(start, "h")
        end = (
            np.datetime64(stop, "D") + np.timedelta64(23, "h")
            if len(stop) <= len("YYYY-MM-DD")
            else np.datetime64(stop, "h")
        )
        return begin, np.datetime64(end, "h")

    def tiling(self, key: CubeKey) -> Tiling | None:
        """Latitude bands of about one Blosc block (every longitude), the same for every hour.

        Read once from the header of one fixed chunk (the first published hour),
        so every call and process probes the same, cacheable, few KiB.
        """
        if self._tiling_known:
            return self._tiling
        coverage = self.coverage()
        try:
            t = 0 if coverage is None else self.time_index(coverage[0])
        except KeyError:
            t = 0
        name = next(iter(self.fields.values()))
        rows = self.store.block_rows(name, [t, 0][: self._lead_dims(name)])
        if rows is not None and len(self.lat) > 1:
            step = rows * abs(float(self.lat[0] - self.lat[1]))
            bands = int(np.ceil(len(self.lat) / rows))
            self._tiling = Tiling(lat0=float(self.lat[0]) - bands * step, lat_step=step)
        self._tiling_known = True
        return self._tiling

    def _lead_dims(self, name: str) -> int:
        """Number of leading (time, level) dimensions of array ``name``."""
        return len(self.store.array(name).shape) - 2

    def configure_decoded_cache(self, max_bytes: int) -> None:
        """Bound decoded Blosc reuse to final field chunks, using ``final_after``.

        Field views share this reader and its budget. Recent ERA5T chunks never
        qualify. Surface pressure in a separate reader has its own budget.
        """
        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 0:
            raise ValueError("max_bytes must be a nonnegative integer")
        if not max_bytes:
            self.store.configure_decoded_cache(0, immutable=lambda key: False)
            return
        arrays = set(self._cache_fields)
        if self.surface_pressure is not None and self.surface_pressure.store is self.store:
            arrays.add(self.surface_pressure.array)
        self.store.configure_decoded_cache(
            max_bytes, immutable=_final_keys(self.store, arrays, self.final_after)
        )

    def use_cache(self, cache: RangeCache) -> None:
        """Keep the final hours' field chunks in ``cache`` (see the class docstring)."""
        arrays = {*self._cache_fields, "time"}  # past time chunks never change either
        stores = [(self.store, arrays)]
        if self.surface_pressure is not None:
            sp = self.surface_pressure
            if sp.store is self.store:
                arrays.add(sp.array)
            else:
                stores.append((sp.store, {sp.array}))
        for zarr, names in stores:
            inner = zarr.store.inner if isinstance(zarr.store, CachedStore) else zarr.store
            final = _final_keys(zarr, names, self.final_after)
            zarr.store = CachedStore(inner, cache, namespace_for(inner.location), final)

    def cube_keys(self, t0: np.datetime64) -> tuple[CubeKey, CubeKey]:
        """Analyses at ``t0`` and one hour later (no run: reanalysis)."""
        h = np.datetime64(t0, "h")
        return CubeKey(h), CubeKey(h + self.time_step)

    def fetch(
        self,
        key: CubeKey,
        domain: Domain,
        p_range: tuple[float, float],
        *,
        level_margin: int = 1,
        executor: Executor | None = None,
    ) -> Cube:
        """Read one hour over ``domain`` with the levels bracketing ``p_range`` (Pa)."""
        hour = key.valid
        try:
            t = self.time_index(hour)
        except KeyError as exc:
            raise DataUnavailableError(f"{self.source_id}: {hour} not in the store") from exc
        grid = GridIndex.locate(domain, self.lat, self.lon)
        ps: NDArray[np.float32] | None = None
        ps_lo = ps_hi = STANDARD_SURFACE_PRESSURE
        if self.surface_pressure is not None:
            sp = self.surface_pressure
            ps = sp.store.read_rows(sp.array, [(t,)], grid.rows, executor=executor)[0][:, grid.cols]
            if np.isfinite(ps).any():
                ps_lo, ps_hi = float(np.nanmin(ps)), float(np.nanmax(ps))
        positions = self.vertical.select_bracketing(
            p_range[0], p_range[1], ps_lo, ps_hi, level_margin
        )
        lead = [(t, int(k)) for k in self._store_index[positions]]
        fields = {
            name: self.store.read_rows(array, lead, grid.rows, executor=executor)[:, :, grid.cols]
            for name, array in self.fields.items()
        }
        return Cube(
            source=self.source_id,
            hour=np.datetime64(hour, "h"),
            lat=grid.lat,
            lon=grid.lon,
            vertical=self.vertical.subset(positions),
            fields=fields,
            ps=ps,
        )

stats property

Transfer counters of every underlying byte store.

configure_decoded_cache(max_bytes)

Bound decoded Blosc reuse to final field chunks, using final_after.

Field views share this reader and its budget. Recent ERA5T chunks never qualify. Surface pressure in a separate reader has its own budget.

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

def configure_decoded_cache(self, max_bytes: int) -> None:
    """Bound decoded Blosc reuse to final field chunks, using ``final_after``.

    Field views share this reader and its budget. Recent ERA5T chunks never
    qualify. Surface pressure in a separate reader has its own budget.
    """
    if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 0:
        raise ValueError("max_bytes must be a nonnegative integer")
    if not max_bytes:
        self.store.configure_decoded_cache(0, immutable=lambda key: False)
        return
    arrays = set(self._cache_fields)
    if self.surface_pressure is not None and self.surface_pressure.store is self.store:
        arrays.add(self.surface_pressure.array)
    self.store.configure_decoded_cache(
        max_bytes, immutable=_final_keys(self.store, arrays, self.final_after)
    )

coverage()

Published period, from the store's valid_time_start/valid_time_stop.

For ARCO-ERA5 the model-level store lags ERA5T (about three months).

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

def coverage(self) -> tuple[np.datetime64, np.datetime64] | None:
    """Published period, from the store's ``valid_time_start``/``valid_time_stop``.

    For ARCO-ERA5 the model-level store lags ERA5T (about three months).
    """
    attrs = self.store.attrs
    start, stop = attrs.get("valid_time_start"), attrs.get("valid_time_stop")
    if not start or not stop:
        return None
    begin = np.datetime64(start, "h")
    end = (
        np.datetime64(stop, "D") + np.timedelta64(23, "h")
        if len(stop) <= len("YYYY-MM-DD")
        else np.datetime64(stop, "h")
    )
    return begin, np.datetime64(end, "h")

cube_keys(t0)

Analyses at t0 and one hour later (no run: reanalysis).

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

def cube_keys(self, t0: np.datetime64) -> tuple[CubeKey, CubeKey]:
    """Analyses at ``t0`` and one hour later (no run: reanalysis)."""
    h = np.datetime64(t0, "h")
    return CubeKey(h), CubeKey(h + self.time_step)

fetch(key, domain, p_range, *, level_margin=1, executor=None)

Read one hour over domain with the levels bracketing p_range (Pa).

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

def fetch(
    self,
    key: CubeKey,
    domain: Domain,
    p_range: tuple[float, float],
    *,
    level_margin: int = 1,
    executor: Executor | None = None,
) -> Cube:
    """Read one hour over ``domain`` with the levels bracketing ``p_range`` (Pa)."""
    hour = key.valid
    try:
        t = self.time_index(hour)
    except KeyError as exc:
        raise DataUnavailableError(f"{self.source_id}: {hour} not in the store") from exc
    grid = GridIndex.locate(domain, self.lat, self.lon)
    ps: NDArray[np.float32] | None = None
    ps_lo = ps_hi = STANDARD_SURFACE_PRESSURE
    if self.surface_pressure is not None:
        sp = self.surface_pressure
        ps = sp.store.read_rows(sp.array, [(t,)], grid.rows, executor=executor)[0][:, grid.cols]
        if np.isfinite(ps).any():
            ps_lo, ps_hi = float(np.nanmin(ps)), float(np.nanmax(ps))
    positions = self.vertical.select_bracketing(
        p_range[0], p_range[1], ps_lo, ps_hi, level_margin
    )
    lead = [(t, int(k)) for k in self._store_index[positions]]
    fields = {
        name: self.store.read_rows(array, lead, grid.rows, executor=executor)[:, :, grid.cols]
        for name, array in self.fields.items()
    }
    return Cube(
        source=self.source_id,
        hour=np.datetime64(hour, "h"),
        lat=grid.lat,
        lon=grid.lon,
        vertical=self.vertical.subset(positions),
        fields=fields,
        ps=ps,
    )

select_fields(names)

Independent field view sharing stores and byte caches with this source.

Only listed output fields are fetched. Cache eligibility retains all original fields, so attaching a cache through a view cannot narrow the parent's cache policy. Configure shared stores before concurrent use.

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

def select_fields(self, names: Sequence[str]) -> GridSource:
    """Independent field view sharing stores and byte caches with this source.

    Only listed output fields are fetched. Cache eligibility retains all
    original fields, so attaching a cache through a view cannot narrow the
    parent's cache policy. Configure shared stores before concurrent use.
    """
    if not names or len(set(names)) != len(names):
        raise ValueError("fields must be a nonempty sequence of unique names")
    unknown = set(names) - self.fields.keys()
    if unknown:
        raise ValueError(f"unknown fields: {sorted(unknown)}")
    view = copy(self)
    view.fields = {name: self.fields[name] for name in names}
    view._tiling = None
    view._tiling_known = False
    return view

tiling(key)

Latitude bands of about one Blosc block (every longitude), the same for every hour.

Read once from the header of one fixed chunk (the first published hour), so every call and process probes the same, cacheable, few KiB.

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

def tiling(self, key: CubeKey) -> Tiling | None:
    """Latitude bands of about one Blosc block (every longitude), the same for every hour.

    Read once from the header of one fixed chunk (the first published hour),
    so every call and process probes the same, cacheable, few KiB.
    """
    if self._tiling_known:
        return self._tiling
    coverage = self.coverage()
    try:
        t = 0 if coverage is None else self.time_index(coverage[0])
    except KeyError:
        t = 0
    name = next(iter(self.fields.values()))
    rows = self.store.block_rows(name, [t, 0][: self._lead_dims(name)])
    if rows is not None and len(self.lat) > 1:
        step = rows * abs(float(self.lat[0] - self.lat[1]))
        bands = int(np.ceil(len(self.lat) / rows))
        self._tiling = Tiling(lat0=float(self.lat[0]) - bands * step, lat_step=step)
    self._tiling_known = True
    return self._tiling

time_index(hour)

Index of hour on the store's time axis, verified against the time array.

The axis is assumed hourly and contiguous; each time chunk is checked once (a few KiB) before its indices are trusted.

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

def time_index(self, hour: np.datetime64) -> int:
    """Index of ``hour`` on the store's time axis, verified against the time array.

    The axis is assumed hourly and contiguous; each time chunk is checked
    once (a few KiB) before its indices are trusted.
    """
    hours = int((np.datetime64(hour, "h") - self._epoch) / np.timedelta64(1, "h"))
    idx = hours - self._t0
    meta = self.store.array("time")
    if not 0 <= idx < meta.shape[0]:
        raise KeyError(f"{hour} outside the store's time axis")
    size = meta.chunks[0]
    chunk = idx // size
    with self._lock:
        known = chunk in self._checked_time_chunks
    if not known:
        start, stop = chunk * size, min((chunk + 1) * size, meta.shape[0])
        values = self.store.read_1d("time", start, stop)
        if not np.array_equal(values, np.arange(start, stop) + self._t0):
            raise RuntimeError(f"{self.source_id}: time axis is not hourly-contiguous")
        with self._lock:
            self._checked_time_chunks.add(chunk)
    return idx

use_cache(cache)

Keep the final hours' field chunks in cache (see the class docstring).

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

def use_cache(self, cache: RangeCache) -> None:
    """Keep the final hours' field chunks in ``cache`` (see the class docstring)."""
    arrays = {*self._cache_fields, "time"}  # past time chunks never change either
    stores = [(self.store, arrays)]
    if self.surface_pressure is not None:
        sp = self.surface_pressure
        if sp.store is self.store:
            arrays.add(sp.array)
        else:
            stores.append((sp.store, {sp.array}))
    for zarr, names in stores:
        inner = zarr.store.inner if isinstance(zarr.store, CachedStore) else zarr.store
        final = _final_keys(zarr, names, self.final_after)
        zarr.store = CachedStore(inner, cache, namespace_for(inner.location), final)

SurfacePressure dataclass

Where a hybrid-level source reads its surface pressure.

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

@dataclass(frozen=True)
class SurfacePressure:
    """Where a hybrid-level source reads its surface pressure."""

    store: ZarrStore
    array: str = "surface_pressure"

arco_model_levels(*, levels=None, base_url=ARCO_BASE_URL, fields=None)

ARCO-ERA5 L137 model levels (0.25 deg, hourly) + surface pressure from the PL store.

levels restricts the hybrid levels used (default: all 137).

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

def arco_model_levels(
    *,
    levels: ArrayLike | None = None,
    base_url: str = ARCO_BASE_URL,
    fields: dict[str, str] | None = None,
) -> GridSource:
    """ARCO-ERA5 L137 model levels (0.25 deg, hourly) + surface pressure from the PL store.

    ``levels`` restricts the hybrid levels used (default: all 137).
    """
    ml = _open(f"{base_url}/{ARCO_MODEL_LEVEL}")
    pl = _open(f"{base_url}/{ARCO_PRESSURE_LEVEL}")
    vertical = model_levels(levels)
    return GridSource(
        _source_id("arco-era5-ml", model_levels().ids, vertical.ids),
        ml,
        vertical,
        fields=fields,
        level_array="hybrid",
        surface_pressure=SurfacePressure(pl),
    )

arco_pressure_levels(*, levels=None, base_url=ARCO_BASE_URL, fields=None)

ARCO-ERA5 pressure levels (0.25 deg, hourly); levels in hPa (default: all 37).

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

def arco_pressure_levels(
    *,
    levels: ArrayLike | None = None,
    base_url: str = ARCO_BASE_URL,
    fields: dict[str, str] | None = None,
) -> GridSource:
    """ARCO-ERA5 pressure levels (0.25 deg, hourly); ``levels`` in hPa (default: all 37)."""
    pl = _open(f"{base_url}/{ARCO_PRESSURE_LEVEL}")
    all_ids = np.asarray(pl.read_1d("level"), dtype=np.int64)
    vertical = pressure_levels(all_ids if levels is None else levels)
    return GridSource(
        _source_id("arco-era5-pl", all_ids, vertical.ids),
        pl,
        vertical,
        fields=fields,
        level_array="level",
    )