firecube 0.1.4__py3-none-any.whl

firecube/core/zarr/__init__.py

@@ -0,0 +1,69 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Zarr utilities (validation, scrub, multires, state).
+
+Public surface for CLI/API and plugins. Keep modules small and layered:
+- `validation`: read-only Zarr inspection.
+- `scrub`: mutating maintenance built on validation + ChunkManager.
+- `multires`: multiresolution builders.
+- `state`: shared Firecube state-array helpers.
+"""
+
+import warnings
+
+# Zarr v3 warns about non-Zarr files (like our manifest) in the store directory.
+try:
+    from zarr.errors import ZarrUserWarning
+
+    warnings.filterwarnings("ignore", category=ZarrUserWarning)
+except ImportError:
+    pass
+
+from firecube.core.zarr.io import ZarrIO
+from firecube.core.zarr.multires import MultiresConfig, ZarrMultiresBuilder
+from firecube.core.zarr.region_writer import RegionZarrWriter, RegionZarrWriterProtocol
+from firecube.core.zarr.scrub import ScrubResult, run_scrub
+from firecube.core.zarr.state import (
+    DEFAULT_STATE_MEANING,
+    ensure_timestamp_state_array,
+    expand_time_index_ranges_to_chunk_boundaries,
+    update_timestamp_state,
+)
+from firecube.core.zarr.validation import (
+    ZarrValidationReport,
+    find_extra_chunks,
+    group_exists,
+    read_chunk_grid,
+    validate_group_with_fs,
+)
+
+__all__ = [
+    "DEFAULT_STATE_MEANING",
+    "MultiresConfig",
+    "RegionZarrWriter",
+    "RegionZarrWriterProtocol",
+    "ScrubResult",
+    "ZarrIO",
+    "ZarrMultiresBuilder",
+    "ZarrValidationReport",
+    "ensure_timestamp_state_array",
+    "expand_time_index_ranges_to_chunk_boundaries",
+    "find_extra_chunks",
+    "group_exists",
+    "read_chunk_grid",
+    "run_scrub",
+    "update_timestamp_state",
+    "validate_group_with_fs",
+]

firecube/core/zarr/_reserved_attrs.py

@@ -0,0 +1,47 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Registry of array attribute keys that firecube manages internally.
+
+Plugins must not set these keys in ZarrArraySpec.attrs; doing so would
+conflict with firecube's own internal bookkeeping.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+__all__ = ["RESERVED_ARRAY_ATTRS", "assert_attrs_safe"]
+
+RESERVED_ARRAY_ATTRS: frozenset[str] = frozenset(
+    {
+        "_ARRAY_DIMENSIONS",
+        "_FillValue",
+        "firecube_run_id",
+        "firecube_span_id",
+        "firecube_internal",
+        "firecube_static_written",
+    }
+)
+
+
+def assert_attrs_safe(attrs: Mapping[str, Any]) -> None:
+    """Raise ValueError if *attrs* contains any reserved key."""
+    for key in attrs:
+        if key in RESERVED_ARRAY_ATTRS:
+            raise ValueError(
+                f"Reserved attr {key!r} is managed by firecube, not the plugin. "
+                "Remove it from ZarrArraySpec.attrs."
+            )

firecube/core/zarr/_reserved_root_attrs.py

@@ -0,0 +1,74 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Guard reserved root attrs supplied by external writers.
+
+This guard is for code paths that accept root attrs from EXTERNAL sources — user CLI options,
+plugin output, or any other non-firecube author. Use this guard wherever such attrs are about
+to be persisted to a Zarr root.
+
+Firecube-internal writers MUST NOT call this guard. The slot-index service
+(ChunkManager.ensure_slot_index_model._mirror_attrs) is the authoritative writer for the
+reserved root attrs and writes them DIRECTLY. Drift detection on subsequent runs is handled by
+the precedence matrix in _apply_slot_model_precedence, not by this guard.
+
+In short: "assert no reserved root attrs appear in user/plugin-supplied attrs" — that is the
+entire contract.
+
+This module is DISTINCT from _reserved_attrs.py which guards array attributes. Root attributes
+and array attributes are separate protection surfaces.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from firecube.core.slot_index import (
+    SLOT_INDEX_MODEL_ATTR,
+    SLOT_INDEX_MODEL_IDENTITY_HASH_ATTR,
+)
+
+__all__ = ["RESERVED_ROOT_ATTRS", "assert_root_attrs_safe"]
+
+RESERVED_ROOT_ATTRS: frozenset[str] = frozenset(
+    {
+        SLOT_INDEX_MODEL_ATTR,
+        SLOT_INDEX_MODEL_IDENTITY_HASH_ATTR,
+    }
+)
+
+
+def assert_root_attrs_safe(attrs: Mapping[str, Any]) -> None:
+    """Raise ValueError if any key in attrs is a reserved firecube root attribute name.
+
+    The check is NAME-ONLY: the value content of the attr is NOT inspected.
+    Even a "correct-looking" value from user code is refused because this guard
+    exists to prevent user/plugin code paths from writing into the reserved namespace,
+    not to validate the value.
+
+    Args:
+        attrs: Mapping of root attribute names to values (typically from user/plugin
+               code before persisting to a Zarr root).
+
+    Raises:
+        ValueError: If any key in attrs overlaps with RESERVED_ROOT_ATTRS.
+    """
+    offending = RESERVED_ROOT_ATTRS.intersection(attrs)
+    if offending:
+        raise ValueError(
+            f"attrs contain reserved firecube root attribute(s): {sorted(offending)!r}. "
+            f"Reserved root attributes are managed internally by firecube "
+            f"(see RESERVED_ROOT_ATTRS = {sorted(RESERVED_ROOT_ATTRS)!r})."
+        )

firecube/core/zarr/io.py

@@ -0,0 +1,88 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+
+import xarray as xr
+import zarr
+
+if TYPE_CHECKING:
+    from firecube.core.filesystem.store_factory import ZarrStoreHandle
+    from firecube.core.storage.session import StorageSession
+    from firecube.core.storage.uri import StorageUri
+
+
+class ZarrIO:
+    def __init__(self, session: StorageSession) -> None:
+        self._session = session
+
+    def create_store(self, uri: StorageUri, mode: str = "w") -> ZarrStoreHandle:
+        """Create a driver-aware Zarr store handle for a typed storage URI.
+
+        The returned handle can be splatted into zarr/xarray APIs via
+        ``handle.zarr_kwargs()`` and is constructed using the parent storage
+        session's configured driver.
+        """
+        from firecube.core.filesystem import create_session_zarr_store
+        from firecube.core.storage.session import storage_config_from_binding
+
+        return create_session_zarr_store(
+            uri=uri,
+            storage_config=storage_config_from_binding(self._session._binding),
+            mode=mode,
+        )
+
+    def open_group(self, uri: StorageUri, mode: str = "r") -> Any:
+        kwargs: dict[str, Any] = {**self.create_store(uri, mode).zarr_kwargs(), "mode": mode}
+        return zarr.open_group(**kwargs)
+
+    def open_dataset(self, uri: StorageUri, group: str = "", **xr_kwargs: Any) -> xr.Dataset:
+        """Open an xarray dataset from a Zarr group using the session store.
+
+        ``uri`` must be a canonical ``StorageUri``; raw URI strings are not
+        accepted by this elevated session API. Additional keyword arguments are
+        forwarded to ``xr.open_zarr`` (e.g. ``decode_times=False`` for raw
+        attribute inspection by structural validators).
+        """
+        handle = self.create_store(uri, mode="r")
+        return xr.open_zarr(**handle.zarr_kwargs(), group=group, **xr_kwargs)
+
+    def write_dataset(
+        self,
+        ds: xr.Dataset,
+        uri: StorageUri,
+        group: str,
+        mode: str = "w",
+        **kwargs: Any,
+    ) -> None:
+        """Write an xarray dataset to a Zarr group using the session store.
+
+        Store construction is centralized through ``create_store`` so local,
+        remote, fsspec, and obstore targets follow the same session-bound driver
+        configuration. ``zarr_format`` may be supplied for xarray-style callers;
+        only format 3 is supported by the underlying Firecube write helper.
+        """
+        handle = self.create_store(uri, mode)
+        zarr_format = kwargs.pop("zarr_format", 3)
+        if zarr_format != 3:
+            raise ValueError("session.zarr.write_dataset only supports zarr_format=3")
+        ds.to_zarr(
+            **handle.zarr_kwargs(),
+            group=group,
+            mode=cast(Any, mode),
+            zarr_format=zarr_format,
+            **kwargs,
+        )

firecube/core/zarr/layers.py

@@ -0,0 +1,184 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Core logic for building multi-resolution Zarr layers."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable, MutableMapping
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+import xarray as xr
+
+if TYPE_CHECKING:
+    from firecube.core.storage.session import StorageSession
+
+from firecube.core.storage.uri import StorageUri
+
+DEFAULT_MULTIRES_RESOLUTIONS: tuple[float, ...] = (1.0, 0.5)
+
+
+def build_multires_layers(
+    zarr_store: str | Path | MutableMapping[str, bytes],
+    *,
+    session: StorageSession,
+    resolutions: Iterable[float] = DEFAULT_MULTIRES_RESOLUTIONS,
+    time_dim_name: str = "timestamp",
+    silent: bool = True,
+    logger: logging.Logger | None = None,
+    group: str | None = None,
+    strict: bool = False,
+    ds_input: xr.Dataset | None = None,
+) -> list[str]:
+    """Build coarsened multi-resolution layers inside an existing Zarr store.
+
+    Parameters
+    ----------
+    zarr_store:
+        Store path or mapping. When a group is provided, zarr_store is assumed
+        to point at the *root* of the product and `group` selects the dataset.
+    group:
+        Optional group path (e.g. "F048/FWI"). When omitted, the function
+        attempts to open the root dataset directly. Pass ``group`` explicitly
+        if the store uses a grouped layout.
+    strict:
+        When True, failures to open the base dataset are raised instead of
+        being converted into an empty result. This is useful for CLI/API
+        operations where a failure should abort the request.
+    ds_input:
+        Optional xarray Dataset to coarsen. When provided, the function coarsens
+        only this data and appends it to multiresolution layers (instead of
+        rebuilding everything from the store).
+    """
+
+    groups_written: list[str] = []
+
+    # Normalize resolutions
+    if resolutions is True:
+        resolutions = list(DEFAULT_MULTIRES_RESOLUTIONS)
+    elif isinstance(resolutions, str):
+        try:
+            # Handle "1.0,0.5" or similar
+            parsed = [float(r.strip()) for r in resolutions.split(",") if r.strip()]
+        except (ValueError, TypeError) as exc:
+            raise ValueError(f"Invalid resolutions: {resolutions!r}") from exc
+        if not parsed:
+            raise ValueError(f"Invalid resolutions: {resolutions!r}")
+        resolutions = parsed
+
+    if not resolutions or not isinstance(resolutions, Iterable):
+        return groups_written
+
+    def _get_base_ds(
+        z_store: str,
+        grp: str | None,
+    ) -> xr.Dataset:
+        z_uri = StorageUri.parse(z_store)
+        open_kwargs = {"chunks": "auto", "consolidated": False}
+        if grp:
+            return session.zarr.open_dataset(z_uri, **open_kwargs, group=grp)
+        try:
+            return session.zarr.open_dataset(z_uri, **open_kwargs)
+        except Exception as exc:
+            raise ValueError(
+                f"Cannot open root dataset at {z_store!r}. "
+                "If the store uses a grouped layout, pass the group explicitly."
+            ) from exc
+
+    ds_base: xr.Dataset
+    is_incremental = False
+
+    try:
+        if ds_input is not None:
+            ds_base = ds_input
+            is_incremental = True
+        else:
+            ds_base = _get_base_ds(str(zarr_store), group)
+    except Exception as exc:
+        if logger and not silent:
+            logger.warning("Failed to open base dataset for multires: %s", exc)
+        if strict:
+            raise
+        return groups_written
+
+    ds_base = ds_base.sortby(["lat", "lon"])
+
+    dlat = (
+        float(abs(ds_base.lat.diff("lat").mean().values))
+        if ds_base.sizes.get("lat", 0) > 1
+        else 1.0
+    )
+    dlon = (
+        float(abs(ds_base.lon.diff("lon").mean().values))
+        if ds_base.sizes.get("lon", 0) > 1
+        else 1.0
+    )
+
+    for res in resolutions:
+        lat_factor = max(1, round(res / dlat))
+        lon_factor = max(1, round(res / dlon))
+
+        # When a group is provided, write multiresolution layers under that
+        # group (e.g. F048/multires/1.0deg). Otherwise, default to a top-level
+        # multires/ path.
+        prefix = f"{group.rstrip('/')}/" if group else ""
+        group_name = f"{prefix}multires/{res}deg"
+
+        # Decide mode: if incremental and group already exists, append.
+        mode: Literal["w", "a"] = "w"
+        if is_incremental:
+            try:
+                base_s = str(zarr_store).rstrip("/")
+                target_uri = f"{base_s}/{group_name}"
+                if session.exists(StorageUri.parse(target_uri)):
+                    mode = "a"
+            except Exception:
+                pass
+
+        if logger and not silent:
+            msg = "Building multi-resolution layer"
+            if mode == "a":
+                msg = "Appending to multi-resolution layer"
+            logger.info(
+                msg,
+                extra={
+                    "group": group_name,
+                    "lat_factor": lat_factor,
+                    "lon_factor": lon_factor,
+                    "mode": mode,
+                },
+            )
+
+        ds_coarse = cast(
+            Any, ds_base.coarsen(lat=lat_factor, lon=lon_factor, boundary="trim")
+        ).mean(skipna=True)
+
+        session.zarr.write_dataset(
+            ds_coarse,
+            StorageUri.parse(str(zarr_store)),
+            group=group_name,
+            mode=mode,
+            append_dim=time_dim_name if mode == "a" else None,
+            zarr_format=3,
+            safe_chunks=False,
+            align_chunks=True,
+        )
+
+        if logger and not silent:
+            logger.info("Multi-resolution layer written", extra={"layer": group_name})
+        groups_written.append(group_name)
+
+    return groups_written

firecube/core/zarr/multires.py

@@ -0,0 +1,178 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Core helpers for building multi-resolution Zarr layers.
+
+This module provides a generic, storage-aware multiresolution builder
+that can operate on any existing Zarr product using the shared
+StorageConfig and ChunkManager infrastructure.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, cast
+
+from firecube.core.config import StorageConfig
+from firecube.core.controlplane import ChunkManager
+from firecube.core.product.identity import ProductIdentity
+from firecube.core.product.target import ProductTarget
+from firecube.core.storage.binding import StorageBinding
+from firecube.core.storage.driver_config import StorageDriverConfig
+from firecube.core.storage.session import StorageSession
+from firecube.core.storage.uri import StorageUri
+from firecube.core.zarr.layers import build_multires_layers
+
+log = logging.getLogger("firecube.core.zarr.multires")
+
+
+@dataclass
+class MultiresConfig:
+    """Configuration for building multiresolution Zarr layers."""
+
+    product: str
+    group: str
+    resolutions: list[float]
+    time_dim_name: str = "timestamp"
+
+
+class ZarrMultiresBuilder:
+    """Generic multiresolution builder for Zarr products.
+
+    This operates on an existing Zarr product by:
+      - resolving the store URI from StorageConfig
+      - opening the target group
+      - building multiresolution views using shared zarr_utils helpers
+
+    It is intentionally generic; product-specific helpers (e.g. ExampleProductMultires)
+    should build MultiresConfig instances and call .build().
+    """
+
+    def __init__(
+        self,
+        storage_config: StorageConfig,
+        *,
+        product_name: str,
+        product_uri: StorageUri | None = None,
+        chunk_manager: ChunkManager | None = None,
+        session: StorageSession | None = None,
+    ) -> None:
+        if not product_name:
+            raise ValueError("ZarrMultiresBuilder requires a non-empty product_name.")
+        self.storage_config = storage_config
+        self.product_name = product_name
+        self.session = session
+        resolved_product_uri = product_uri or (
+            session.product.product_uri if session is not None else None
+        )
+        if resolved_product_uri is None:
+            raise ValueError("ZarrMultiresBuilder requires product_uri or session.")
+        self.product_uri: StorageUri = resolved_product_uri
+        driver_config = StorageDriverConfig.from_storage_config(storage_config)
+        # This binding intentionally uses the sentinel identity
+        # "__firecube_controlplane__" so multires can address the cube root as a
+        # product-agnostic control-plane scope instead of a real product name.
+        # The sentinel bypasses the product-name short-circuit in
+        # _ControlRootResolver. ChunkManager is currently instantiated but not
+        # used in build(); it is reserved here for future control-plane hooks.
+        self.chunk_manager = chunk_manager or ChunkManager(
+            binding=StorageBinding(
+                identity=ProductIdentity.from_uri(
+                    self.product_uri.parent().join("__firecube_controlplane__"),
+                    "zarr",
+                    product_name="__firecube_controlplane__",
+                ),
+                driver=driver_config,
+            )
+        )
+
+    def _store_uri_for(self, product: str) -> str:
+        """Build the underlying store URI for a given product."""
+        product_uri = (
+            self.product_uri
+            if product.strip("/") == self.product_name
+            else self.product_uri.parent().join(product)
+        )
+        return product_uri.to_str()
+
+    def _session_for(self, cfg: MultiresConfig) -> StorageSession:
+        if self.session is not None:
+            return self.session
+
+        driver_config = StorageDriverConfig.from_storage_config(self.storage_config)
+        resolved_product = ProductTarget.resolve(
+            cfg.product,
+            driver_config,
+            product_name=cfg.product,
+            plugin_default_format="zarr",
+            default_base_uri=self.product_uri.parent(),
+        )
+        return StorageSession(
+            StorageBinding(
+                identity=ProductIdentity(
+                    product_name=resolved_product.product_name,
+                    product_uri=resolved_product.product_uri,
+                    control_root_uri=resolved_product.control_root_uri,
+                    format=resolved_product.format,
+                ),
+                driver=driver_config,
+            )
+        )
+
+    def build(self, cfg: MultiresConfig) -> dict:
+        """Build multiresolution layers for the given configuration."""
+        if not cfg.resolutions:
+            raise ValueError("At least one resolution must be provided for multires.")
+
+        store_uri = self._store_uri_for(cfg.product)
+        log.info(
+            "Building Zarr multiresolution layers",
+            extra={
+                "product": cfg.product,
+                "group": cfg.group,
+                "store_uri": store_uri,
+                "resolutions": cfg.resolutions,
+            },
+        )
+
+        # Delegate to shared helper
+        build_fn = cast(Any, build_multires_layers)
+        layers = build_fn(
+            store_uri,
+            resolutions=cfg.resolutions,
+            silent=False,
+            logger=log,
+            group=cfg.group,
+            strict=True,
+            session=self._session_for(cfg),
+            time_dim_name=cfg.time_dim_name,
+        )
+
+        log.info(
+            "Zarr multiresolution build complete",
+            extra={
+                "product": cfg.product,
+                "group": cfg.group,
+                "store_uri": store_uri,
+                "layers": layers,
+            },
+        )
+
+        return {
+            "product": cfg.product,
+            "group": cfg.group,
+            "store_uri": store_uri,
+            "layers": layers,
+        }