climate-ref-core 0.5.0__py3-none-any.whl

climate_ref_core/__init__.py

@@ -0,0 +1,7 @@
+"""
+Rapid evaluating Climate data
+"""
+
+import importlib.metadata
+
+__version__ = importlib.metadata.version("climate-ref-core")

climate_ref_core/constraints.py

@@ -0,0 +1,363 @@
+"""
+Dataset selection constraints
+"""
+
+import sys
+import warnings
+from collections import defaultdict
+from collections.abc import Mapping
+from typing import Protocol, runtime_checkable
+
+if sys.version_info < (3, 11):
+    from typing_extensions import Self
+else:
+    from typing import Self
+
+import numpy as np
+import pandas as pd
+from attrs import frozen
+from loguru import logger
+
+from climate_ref_core.datasets import SourceDatasetType
+from climate_ref_core.exceptions import ConstraintNotSatisfied
+
+
+@runtime_checkable
+class GroupValidator(Protocol):
+    """
+    A constraint that must be satisfied when executing a given diagnostic run.
+
+    All constraints must be satisfied for a given group to be run.
+    """
+
+    def validate(self, group: pd.DataFrame) -> bool:
+        """
+        Validate if the constraint is satisfied by the dataset.
+
+        This is executed after the apply method to determine if the constraint is satisfied.
+        If the constraint is not satisfied, the group will not be executed.
+
+        Parameters
+        ----------
+        group
+            A group of datasets that is being validated.
+
+        Returns
+        -------
+        :
+            Whether the constraint is satisfied
+        """
+        ...
+
+
+@runtime_checkable
+class GroupOperation(Protocol):
+    """
+    An operation to perform on a group of datasets resulting in a new group of datasets.
+
+    !! warning
+
+        Operations should not mutate the input group, but instead return a new group.
+        Mutating the input group may result in unexpected behaviour.
+    """
+
+    def apply(self, group: pd.DataFrame, data_catalog: pd.DataFrame) -> pd.DataFrame:
+        """
+        Perform an operation on the group of datasets.
+
+        A new group of datasets should be returned if modifications are required,
+        and the input group should not be modified. If no modifications are required,
+        return the input group unchanged.
+        If this operation fails, a ConstraintNotSatisfied exception should be raised.
+
+        Parameters
+        ----------
+        group
+            A group of datasets that is being validated.
+        data_catalog
+            The data catalog of datasets
+
+        Raises
+        ------
+        ConstraintNotSatisfied
+            The operation was not successful
+
+        Returns
+        -------
+        :
+            The updated group of datasets
+        """
+        ...
+
+
+GroupConstraint = GroupOperation | GroupValidator
+"""
+A constraint that must be satisfied for a group of datasets to be executed.
+
+This is applied to a group of datasets representing the inputs to a potential diagnostic execution.
+The group must satisfy all constraints to be processed.
+
+This can include operations that are applied to a group of datasets which may modify the group,
+but may also include validators that check if the group satisfies a certain condition.
+"""
+
+
+def apply_constraint(
+    dataframe: pd.DataFrame,
+    constraint: GroupConstraint,
+    data_catalog: pd.DataFrame,
+) -> pd.DataFrame | None:
+    """
+    Apply a constraint to a group of datasets
+
+    Parameters
+    ----------
+    dataframe:
+        The group of datasets to apply the constraint to.
+    constraint
+        The constraint to apply.
+    data_catalog
+        The data catalog of all datasets.
+
+    Returns
+    -------
+    :
+        The updated group of datasets or None if the constraint was not satisfied
+    """
+    try:
+        updated_group = (
+            constraint.apply(dataframe, data_catalog) if isinstance(constraint, GroupOperation) else dataframe
+        )
+
+        valid = constraint.validate(updated_group) if isinstance(constraint, GroupValidator) else True
+        if not valid:
+            logger.debug(f"Constraint {constraint} not satisfied for {dataframe}")
+            raise ConstraintNotSatisfied(f"Constraint {constraint} not satisfied for {dataframe}")
+    except ConstraintNotSatisfied:
+        logger.debug(f"Constraint {constraint} not satisfied for {dataframe}")
+        return None
+
+    return updated_group
+
+
+@frozen
+class RequireFacets:
+    """
+    A constraint that requires a dataset to have certain facets.
+    """
+
+    dimension: str
+    required_facets: tuple[str, ...]
+
+    def validate(self, group: pd.DataFrame) -> bool:
+        """
+        Check that the required facets are present in the group
+        """
+        if self.dimension not in group:
+            logger.warning(f"Dimension {self.dimension} not present in group {group}")
+            return False
+        return all(value in group[self.dimension].values for value in self.required_facets)
+
+
+@frozen
+class AddSupplementaryDataset:
+    """
+    Include e.g. a cell measure or ancillary variable in the selection.
+    """
+
+    supplementary_facets: Mapping[str, str | tuple[str, ...]]
+    """
+    Facets describing the supplementary dataset.
+    """
+
+    matching_facets: tuple[str, ...]
+    """
+    Facets that must match with datasets in the selection.
+    """
+
+    optional_matching_facets: tuple[str, ...]
+    """
+    Select only the best matching datasets based on similarity with these facets.
+    """
+
+    def apply(
+        self,
+        group: pd.DataFrame,
+        data_catalog: pd.DataFrame,
+    ) -> pd.DataFrame:
+        """
+        Add a supplementary dataset to the group.
+        """
+        supplementary_facets: defaultdict[str, tuple[str, ...]] = defaultdict(tuple)
+        for facet, values in self.supplementary_facets.items():
+            supplementary_facets[facet] = values if isinstance(values, tuple) else (values,)
+
+        for facet in self.matching_facets:
+            values = tuple(group[facet].unique())
+            supplementary_facets[facet] += values
+
+        supplementary_group = data_catalog
+        for facet, values in supplementary_facets.items():
+            mask = supplementary_group[facet].isin(values)
+            supplementary_group = supplementary_group[mask]
+
+        if not supplementary_group.empty and self.optional_matching_facets:
+            facets = list(self.matching_facets + self.optional_matching_facets)
+            datasets = group[facets].drop_duplicates()
+            indices = set()
+            for i in range(len(datasets)):
+                scores = (supplementary_group[facets] == datasets.iloc[i]).sum(axis=1)
+                matches = supplementary_group[scores == scores.max()]
+                # Select the latest version if there are multiple matches
+                matches = matches[matches["version"] == matches["version"].max()]
+                indices.add(matches.index[0])
+            supplementary_group = supplementary_group.loc[list(indices)].drop_duplicates()
+
+        return pd.concat([group, supplementary_group])
+
+    @classmethod
+    def from_defaults(
+        cls,
+        variable: str,
+        source_type: SourceDatasetType,
+    ) -> Self:
+        """
+        Include e.g. a cell measure or ancillary variable in the selection.
+
+        The constraint is created using the defaults for the source_type.
+
+        Parameters
+        ----------
+        variable:
+            The name of the variable to add.
+        source_type:
+            The source_type of the variable to add.
+
+        Returns
+        -------
+        :
+            A constraint to include a supplementary variable.
+
+        """
+        kwargs = {
+            SourceDatasetType.CMIP6: {
+                "matching_facets": (
+                    "source_id",
+                    "grid_label",
+                ),
+                "optional_matching_facets": (
+                    "table_id",
+                    "experiment_id",
+                    "member_id",
+                    "version",
+                ),
+            }
+        }
+        variable_facet = {
+            SourceDatasetType.CMIP6: "variable_id",
+        }
+
+        supplementary_facets = {variable_facet[source_type]: variable}
+        return cls(supplementary_facets, **kwargs[source_type])
+
+
+@frozen
+class RequireContiguousTimerange:
+    """
+    A constraint that requires datasets to have a contiguous timerange.
+    """
+
+    group_by: tuple[str, ...]
+    """
+    The fields to group the datasets by. Each group must be contiguous in time
+    to fulfill the constraint.
+    """
+
+    def validate(self, group: pd.DataFrame) -> bool:
+        """
+        Check that all subgroups of the group have a contiguous timerange.
+        """
+        # Maximum allowed time difference between the end of one file and the
+        # start of the next file.
+        max_timedelta = pd.Timedelta(
+            days=31,  # Maximum number of days in a month.
+            hours=1,  # Allow for potential rounding errors.
+        )
+        group = group.dropna(subset=["start_time", "end_time"])
+        if len(group) < 2:  # noqa: PLR2004
+            return True
+
+        for _, subgroup in group.groupby(list(self.group_by)):
+            if len(subgroup) < 2:  # noqa: PLR2004
+                continue
+            sorted_group = subgroup.sort_values("start_time", kind="stable")
+            start_series = sorted_group["start_time"]
+            end_series = sorted_group["end_time"]
+            # Sometimes the elements of start_series.values are of type datetime64[ns]
+            # and sometimes its elements are of type datetime.datetime.
+            # Convert both arrays to datetime.datetime objects to make sure they
+            # can be subtracted.
+            with warnings.catch_warnings():
+                # We have already mitigated the future change in behaviour of DatetimeProperties.to_pydatetime
+                warnings.simplefilter("ignore", FutureWarning)
+
+                if hasattr(start_series, "dt"):
+                    start_array = np.array(start_series.dt.to_pydatetime())
+                else:
+                    start_array = start_series.values  # type: ignore[assignment]
+                if hasattr(end_series, "dt"):
+                    end_array = np.array(end_series.dt.to_pydatetime())
+                else:
+                    end_array = end_series.values  # type: ignore[assignment]
+            diff = start_array[1:] - end_array[:-1]
+            gap_indices = diff > max_timedelta
+            if gap_indices.any():
+                paths = sorted_group["path"]
+                for gap_idx in np.flatnonzero(gap_indices):
+                    logger.debug(
+                        f"Constraint {self.__class__.__name__} not satisfied "
+                        f"because gap larger than {max_timedelta} found between "
+                        f"{paths.iloc[gap_idx]} and {paths.iloc[gap_idx + 1]}"
+                    )
+                return False
+        return True
+
+
+@frozen
+class RequireOverlappingTimerange:
+    """
+    A constraint that requires datasets to have an overlapping timerange.
+    """
+
+    group_by: tuple[str, ...]
+    """
+    The fields to group the datasets by. There must be overlap in time between
+    the groups to fulfill the constraint.
+    """
+
+    def validate(self, group: pd.DataFrame) -> bool:
+        """
+        Check that all subgroups of the group have an overlapping timerange.
+        """
+        group = group.dropna(subset=["start_time", "end_time"])
+        if len(group) < 2:  # noqa: PLR2004
+            return True
+
+        starts = group.groupby(list(self.group_by))["start_time"].min()
+        ends = group.groupby(list(self.group_by))["end_time"].max()
+        return starts.max() < ends.min()  # type: ignore[no-any-return]
+
+
+@frozen
+class SelectParentExperiment:
+    """
+    Include a dataset's parent experiment in the selection
+    """
+
+    def apply(self, group: pd.DataFrame, data_catalog: pd.DataFrame) -> pd.DataFrame:
+        """
+        Include a dataset's parent experiment in the selection
+
+        Not yet implemented
+        """
+        raise NotImplementedError("This is not implemented yet")  # pragma: no cover

climate_ref_core/dataset_registry.py

@@ -0,0 +1,158 @@
+"""
+Data registries for non-published reference data
+
+These data are placeholders until these data have been added to obs4MIPs.
+The CMIP7 Assessment Fas Track REF requires that reference datasets are openly licensed
+before it is included in any published data catalogs.
+"""
+
+import importlib.resources
+import os
+import pathlib
+import shutil
+
+import pooch
+from loguru import logger
+from rich.progress import track
+
+
+def fetch_all_files(
+    registry: pooch.Pooch,
+    name: str,
+    output_dir: pathlib.Path | None,
+    symlink: bool = False,
+) -> None:
+    """
+    Fetch all files associated with a pooch registry and write them to an output directory.
+
+    Pooch fetches, caches and validates the downloaded files.
+    Subsequent calls to this function will not refetch any previously downloaded files.
+
+    Parameters
+    ----------
+    registry
+        Pooch directory containing a set of files that should be fetched.
+    name
+        Name of the registry.
+    output_dir
+        The root directory to write the files to.
+
+        The directory will be created if it doesn't exist,
+        and matching files will be overwritten.
+
+        If no directory is provided, the files will be fetched from the remote server,
+        but not copied anywhere.
+    symlink
+        If True, symlink all files to this directory.
+        Otherwise, perform a copy.
+    """
+    if output_dir:
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    for key in track(registry.registry.keys(), description=f"Fetching {name} data"):
+        fetch_file = registry.fetch(key)
+
+        if output_dir is None:
+            # Just warm the cache and move onto the next file
+            continue
+
+        linked_file = output_dir / key
+        linked_file.parent.mkdir(parents=True, exist_ok=True)
+        if not linked_file.exists():  # pragma: no cover
+            if symlink:
+                logger.info(f"Linking {key} to {linked_file}")
+
+                os.symlink(fetch_file, linked_file)
+            else:
+                logger.info(f"Copying {key} to {linked_file}")
+                shutil.copy(fetch_file, linked_file)
+        else:
+            logger.info(f"File {linked_file} already exists. Skipping.")
+
+
+class DatasetRegistryManager:
+    """
+    A collection of reference datasets registries
+
+    The REF requires additional reference datasets
+    in addition to obs4MIPs data which can be downloaded via ESGF.
+    Each provider may have different sets of reference data that are needed.
+    These are provider-specific datasets are datasets not yet available in obs4MIPs,
+    or are post-processed from obs4MIPs.
+
+    A dataset registry consists of a file that contains a list of files and checksums,
+    in combination with a base URL that is used to fetch the files.
+    [Pooch](https://www.fatiando.org/pooch/latest/) is used within the DataRegistry
+    to manage the caching, downloading and validation of the files.
+
+    All datasets that are registered here are expected to be openly licensed and freely available.
+    """
+
+    def __init__(self) -> None:
+        self._registries: dict[str, pooch.Pooch] = {}
+
+    def __getitem__(self, item: str) -> pooch.Pooch:
+        """
+        Get a registry by name
+        """
+        return self._registries[item]
+
+    def keys(self) -> list[str]:
+        """
+        Get the list of registry names
+        """
+        return list(self._registries.keys())
+
+    def register(  # noqa: PLR0913
+        self,
+        name: str,
+        base_url: str,
+        package: str,
+        resource: str,
+        cache_name: str | None = None,
+        version: str | None = None,
+    ) -> None:
+        """
+        Register a new dataset registry
+
+        This will create a new Pooch registry and add it to the list of registries.
+        This is typically used by a provider to register a new collections of datasets at runtime.
+
+        Parameters
+        ----------
+        name
+            Name of the registry
+
+            This is used to identify the registry
+        base_url
+            Commmon URL prefix for the files
+        package
+            Name of the package containing the registry resource.
+        resource
+            Name of the resource in the package that contains a list of files and checksums.
+
+            This must be formatted in a way that is expected by pooch.
+        version
+            The version of the data.
+
+            Changing the version will invalidate the cache and force a re-download of the data.
+        cache_name
+            Name to use to generate the cache directory.
+
+            This defaults to the value of `name` if not provided.
+        """
+        if cache_name is None:
+            cache_name = "ref"
+
+        registry = pooch.create(
+            path=pooch.os_cache(cache_name),
+            base_url=base_url,
+            version=version,
+            retry_if_failed=5,
+            env="REF_DATASET_CACHE_DIR",
+        )
+        registry.load_registry(str(importlib.resources.files(package) / resource))
+        self._registries[name] = registry
+
+
+dataset_registry_manager = DatasetRegistryManager()

climate_ref_core/datasets.py

@@ -0,0 +1,157 @@
+"""
+Dataset management and filtering
+"""
+
+import enum
+import functools
+import hashlib
+from collections.abc import Iterable
+from typing import Any, Self
+
+import pandas as pd
+from attrs import field, frozen
+
+
+class SourceDatasetType(enum.Enum):
+    """
+    Types of supported source datasets
+    """
+
+    CMIP6 = "cmip6"
+    CMIP7 = "cmip7"
+    obs4MIPs = "obs4mips"
+    PMPClimatology = "pmp-climatology"
+
+    @classmethod
+    @functools.lru_cache(maxsize=1)
+    def ordered(
+        cls,
+    ) -> list[Self]:
+        """
+        Order in alphabetical order according to their value
+
+        Returns
+        -------
+        :
+            Ordered list of dataset types
+        """
+        return sorted(cls, key=lambda x: x.value)
+
+
+def _clean_facets(raw_values: dict[str, str | tuple[str, ...] | list[str]]) -> dict[str, tuple[str, ...]]:
+    """
+    Clean the value of a facet filter to a tuple of strings
+    """
+    result = {}
+
+    for key, value in raw_values.items():
+        if isinstance(value, list):
+            result[key] = tuple(value)
+        elif isinstance(value, str):
+            result[key] = (value,)
+        elif isinstance(value, tuple):
+            result[key] = value
+    return result
+
+
+@frozen
+class FacetFilter:
+    """
+    A filter to apply to a data catalog of datasets.
+    """
+
+    facets: dict[str, tuple[str, ...]] = field(converter=_clean_facets)
+    """
+    Filters to apply to the data catalog.
+
+    The keys are the metadata fields to filter on, and the values are the values to filter on.
+    The result will only contain datasets where for all fields,
+    the value of the field is one of the given values.
+    """
+    keep: bool = True
+    """
+    Whether to keep or remove datasets that match the filter.
+
+    If true (default), datasets that match the filter will be kept else they will be removed.
+    """
+
+
+@frozen
+class DatasetCollection:
+    """
+    Group of datasets required for a given diagnostic execution for a specific source dataset type.
+    """
+
+    datasets: pd.DataFrame
+    slug_column: str
+    """
+    Column in datasets that contains the unique identifier for the dataset
+    """
+    selector: tuple[tuple[str, str], ...] = ()
+    """
+    Unique key, value pairs that were selected during the initial groupby
+    """
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self.datasets, item)
+
+    def __getitem__(self, item: str | list[str]) -> Any:
+        return self.datasets[item]
+
+    def __hash__(self) -> int:
+        # This hashes each item individually and sums them so order doesn't matter
+        return int(pd.util.hash_pandas_object(self.datasets[self.slug_column]).sum())
+
+    def __eq__(self, other: object) -> bool:
+        return self.__hash__() == other.__hash__()
+
+
+class ExecutionDatasetCollection:
+    """
+    The complete set of datasets required for an execution of a diagnostic.
+
+    This may cover multiple source dataset types.
+    """
+
+    def __init__(self, collection: dict[SourceDatasetType | str, DatasetCollection]):
+        self._collection = {SourceDatasetType(k): v for k, v in collection.items()}
+
+    def __contains__(self, key: SourceDatasetType | str) -> bool:
+        if isinstance(key, str):
+            key = SourceDatasetType(key)
+        return key in self._collection
+
+    def __getitem__(self, key: SourceDatasetType | str) -> DatasetCollection:
+        if isinstance(key, str):
+            key = SourceDatasetType(key)
+        return self._collection[key]
+
+    def __hash__(self) -> int:
+        return hash(self.hash)
+
+    def items(self) -> Iterable[tuple[SourceDatasetType, DatasetCollection]]:
+        """
+        Iterate over the datasets in the collection
+        """
+        return self._collection.items()
+
+    @property
+    def hash(self) -> str:
+        """
+        Unique identifier for the collection
+
+        A SHA1 hash is calculated of the combination of the hashes of the individual collections.
+        The value isn't reversible but can be used to uniquely identify the aggregate of the
+        collections.
+
+        Returns
+        -------
+        :
+            SHA1 hash of the collections
+        """
+        # The dataset collection hashes are reproducible,
+        # so we can use them to hash the diagnostic dataset.
+        # This isn't explicitly true for all Python hashes
+        hash_sum = sum(hash(item) for item in self._collection.values())
+        hash_bytes = hash_sum.to_bytes(16, "little", signed=True)
+        return hashlib.sha1(hash_bytes).hexdigest()  # noqa: S324