climate-ref 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

climate_ref/datasets/base.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Protocol, cast
+from typing import Any, Protocol, cast
 
 import pandas as pd
 from loguru import logger
@@ -35,6 +35,31 @@ def _log_duplicate_metadata(
         )
 
 
+class DatasetParsingFunction(Protocol):
+    """
+    Protocol for a function that parses metadata from a file or directory
+    """
+
+    def __call__(self, file: str, **kwargs: Any) -> dict[str, Any]:
+        """
+        Parse a file or directory and return metadata for the dataset
+
+        Parameters
+        ----------
+        file
+            File or directory to parse
+
+        kwargs
+            Additional keyword arguments to pass to the parsing function.
+
+        Returns
+        -------
+        :
+            Data catalog containing the metadata for the dataset
+        """
+        ...
+
+
 class DatasetAdapter(Protocol):
     """
     An adapter to provide a common interface for different dataset types
@@ -173,7 +198,7 @@ class DatasetAdapter(Protocol):
         slug = unique_slugs[0]
 
         dataset_metadata = data_catalog_dataset[list(self.dataset_specific_metadata)].iloc[0].to_dict()
-        dataset, created = db.get_or_create(DatasetModel, slug=slug, **dataset_metadata)
+        dataset, created = db.get_or_create(DatasetModel, defaults=dataset_metadata, slug=slug)
         if not created:
             logger.warning(f"{dataset} already exists in the database. Skipping")
             return None
@@ -212,6 +237,7 @@ class DatasetAdapter(Protocol):
                 {
                     **{k: getattr(file, k) for k in self.file_specific_metadata},
                     **{k: getattr(file.dataset, k) for k in self.dataset_specific_metadata},
+                    "finalised": file.dataset.finalised,
                 }
                 for file in result
             ],

climate_ref/datasets/cmip6.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
-import traceback
 import warnings
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 import pandas as pd
-import xarray as xr
 from ecgtools import Builder
-from ecgtools.parsers.utilities import extract_attr_with_regex  # type: ignore
 from loguru import logger
 
-from climate_ref.datasets.base import DatasetAdapter
+from climate_ref.config import Config
+from climate_ref.datasets.base import DatasetAdapter, DatasetParsingFunction
+from climate_ref.datasets.cmip6_parsers import parse_cmip6_complete, parse_cmip6_drs
 from climate_ref.models.dataset import CMIP6Dataset
 
 
@@ -22,16 +21,19 @@ def _parse_datetime(dt_str: pd.Series[str]) -> pd.Series[datetime | Any]:
     """
 
     def _inner(date_string: str | None) -> datetime | None:
-        if not date_string:
+        if not date_string or pd.isnull(date_string):
             return None
 
         # Try to parse the date string with and without milliseconds
-        try:
-            dt = datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S")
-        except ValueError:
-            dt = datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S.%f")
+        for fmt in ("%Y-%m-%d", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M:%S.%f"):
+            try:
+                return datetime.strptime(date_string, fmt)
+            except ValueError:
+                continue
 
-        return dt
+        # If all parsing attempts fail, log an error and return None
+        logger.error(f"Failed to parse date string: {date_string}")
+        return None
 
     return pd.Series(
         [_inner(dt) for dt in dt_str],
@@ -44,15 +46,16 @@ def _apply_fixes(data_catalog: pd.DataFrame) -> pd.DataFrame:
     def _fix_parent_variant_label(group: pd.DataFrame) -> pd.DataFrame:
         if group["parent_variant_label"].nunique() == 1:
             return group
-        group["parent_variant_label"] = group["variant_label"].iloc[0]
+        group["parent_variant_label"] = group["parent_variant_label"].iloc[0]
 
         return group
 
-    data_catalog = (
-        data_catalog.groupby("instance_id")
-        .apply(_fix_parent_variant_label, include_groups=False)
-        .reset_index(level="instance_id")
-    )
+    if "parent_variant_label" in data_catalog:
+        data_catalog = (
+            data_catalog.groupby("instance_id")
+            .apply(_fix_parent_variant_label, include_groups=False)
+            .reset_index(level="instance_id")
+        )
 
     if "branch_time_in_child" in data_catalog:
         data_catalog["branch_time_in_child"] = _clean_branch_time(data_catalog["branch_time_in_child"])
@@ -68,88 +71,6 @@ def _clean_branch_time(branch_time: pd.Series[str]) -> pd.Series[float]:
     return pd.to_numeric(branch_time.astype(str).str.replace("D", ""), errors="coerce")
 
 
-def parse_cmip6(file: str) -> dict[str, Any]:
-    """
-    Parser for CMIP6
-
-    This function parses the CMIP6 dataset and returns a dictionary with the metadata.
-    This was copied from the ecgtools package, but we want to log the exception when it fails.
-    """
-    keys = sorted(
-        {
-            "activity_id",
-            "branch_method",
-            "branch_time_in_child",
-            "branch_time_in_parent",
-            "experiment",
-            "experiment_id",
-            "frequency",
-            "grid",
-            "grid_label",
-            "institution_id",
-            "nominal_resolution",
-            "parent_activity_id",
-            "parent_experiment_id",
-            "parent_source_id",
-            "parent_time_units",
-            "parent_variant_label",
-            "realm",
-            "product",
-            "source_id",
-            "source_type",
-            "sub_experiment",
-            "sub_experiment_id",
-            "table_id",
-            "variable_id",
-            "variant_label",
-        }
-    )
-
-    try:
-        with xr.open_dataset(file, chunks={}, use_cftime=True) as ds:
-            info = {key: ds.attrs.get(key) for key in keys}
-            info["member_id"] = info["variant_label"]
-
-            variable_id = info["variable_id"]
-            if variable_id:  # pragma: no branch
-                attrs = ds[variable_id].attrs
-                for attr in ["standard_name", "long_name", "units"]:
-                    info[attr] = attrs.get(attr)
-
-            # Set the default of # of vertical levels to 1
-            vertical_levels = 1
-            start_time, end_time = None, None
-            init_year = None
-            try:
-                vertical_levels = ds[ds.cf["vertical"].name].size
-            except (KeyError, AttributeError, ValueError):
-                ...
-
-            try:
-                start_time, end_time = str(ds.cf["T"][0].data), str(ds.cf["T"][-1].data)
-            except (KeyError, AttributeError, ValueError):
-                ...
-            if info.get("sub_experiment_id"):  # pragma: no branch
-                init_year = extract_attr_with_regex(info["sub_experiment_id"], r"\d{4}")
-                if init_year:  # pragma: no cover
-                    init_year = int(init_year)
-            info["vertical_levels"] = vertical_levels
-            info["init_year"] = init_year
-            info["start_time"] = start_time
-            info["end_time"] = end_time
-            if not (start_time and end_time):
-                info["time_range"] = None
-            else:
-                info["time_range"] = f"{start_time}-{end_time}"
-        info["path"] = str(file)
-        info["version"] = extract_attr_with_regex(str(file), regex=r"v\d{4}\d{2}\d{2}|v\d{1}") or "v0"
-        return info
-
-    except Exception:
-        logger.exception(f"Failed to parse {file}")
-        return {"INVALID_ASSET": file, "TRACEBACK": traceback.format_exc()}
-
-
 class CMIP6DatasetAdapter(DatasetAdapter):
     """
     Adapter for CMIP6 datasets
@@ -191,6 +112,7 @@ class CMIP6DatasetAdapter(DatasetAdapter):
         "standard_name",
         "long_name",
         "units",
+        "finalised",
         slug_column,
     )
 
@@ -208,8 +130,30 @@ class CMIP6DatasetAdapter(DatasetAdapter):
         "grid_label",
     )
 
-    def __init__(self, n_jobs: int = 1):
+    def __init__(self, n_jobs: int = 1, config: Config | None = None):
         self.n_jobs = n_jobs
+        self.config = config or Config.default()
+
+    def get_parsing_function(self) -> DatasetParsingFunction:
+        """
+        Get the parsing function for CMIP6 datasets based on configuration
+
+        The parsing function used is determined by the `cmip6_parser` configuration value:
+        - "drs": Use the DRS parser (default)
+        - "complete": Use the complete parser that extracts all available metadata
+
+        Returns
+        -------
+        :
+            The appropriate parsing function based on configuration
+        """
+        parser_type = self.config.cmip6_parser
+        if parser_type == "complete":
+            logger.info("Using complete CMIP6 parser")
+            return parse_cmip6_complete
+        else:
+            logger.info(f"Using DRS CMIP6 parser (config value: {parser_type})")
+            return parse_cmip6_drs
 
     def find_local_datasets(self, file_or_directory: Path) -> pd.DataFrame:
         """
@@ -228,6 +172,8 @@ class CMIP6DatasetAdapter(DatasetAdapter):
         :
             Data catalog containing the metadata for the dataset
         """
+        parsing_function = self.get_parsing_function()
+
         with warnings.catch_warnings():
             # Ignore the DeprecationWarning from xarray
             warnings.simplefilter("ignore", DeprecationWarning)
@@ -237,7 +183,7 @@ class CMIP6DatasetAdapter(DatasetAdapter):
                 depth=10,
                 include_patterns=["*.nc"],
                 joblib_parallel_kwargs={"n_jobs": self.n_jobs},
-            ).build(parsing_func=parse_cmip6)  # type: ignore
+            ).build(parsing_func=parsing_function)
 
         datasets: pd.DataFrame = builder.df.drop(["init_year"], axis=1)
 
@@ -254,6 +200,14 @@ class CMIP6DatasetAdapter(DatasetAdapter):
             lambda row: "CMIP6." + ".".join([row[item] for item in drs_items]), axis=1
         )
 
+        # Add in any missing metadata columns
+        missing_columns = set(self.dataset_specific_metadata + self.file_specific_metadata) - set(
+            datasets.columns
+        )
+        if missing_columns:
+            for column in missing_columns:
+                datasets[column] = pd.NA
+
         # Temporary fix for some datasets
         # TODO: Replace with a standalone package that contains metadata fixes for CMIP6 datasets
         datasets = _apply_fixes(datasets)

climate_ref/datasets/cmip6_parsers.py

@@ -0,0 +1,189 @@
+"""
+CMIP6 parser functions for extracting metadata from netCDF files
+
+Additional non-official DRS's may be added in the future.
+"""
+
+import traceback
+from typing import Any
+
+import xarray as xr
+from ecgtools.parsers.cmip import parse_cmip6_using_directories  # type: ignore
+from ecgtools.parsers.utilities import extract_attr_with_regex  # type: ignore
+from loguru import logger
+
+
+def _parse_daterange(date_range: str) -> tuple[str | None, str | None]:
+    """
+    Parse a date range string into start and end dates
+
+    The output from this is an estimated date range until the file is completely parsed.
+
+    Parameters
+    ----------
+    date_range
+        Date range string in the format "YYYYMM-YYYYMM"
+
+    Returns
+    -------
+    :
+        Tuple containing start and end dates as strings in the format "YYYY-MM-DD"
+    """
+    try:
+        start, end = date_range.split("-")
+        if len(start) != 6 or len(end) != 6:  # noqa: PLR2004
+            raise ValueError("Date range must be in the format 'YYYYMM-YYYYMM'")
+
+        start = f"{start[:4]}-{start[4:6]}-01"
+        # Up to the 30th of the month, assuming a 30-day month
+        # These values will be corrected later when the file is parsed
+        end = f"{end[:4]}-{end[4:6]}-30"
+
+        return start, end
+    except ValueError:
+        logger.error(f"Invalid date range format: {date_range}")
+        return None, None
+
+
+def parse_cmip6_complete(file: str, **kwargs: Any) -> dict[str, Any]:
+    """
+    Complete parser for CMIP6 files
+
+    This parser loads each file and extracts all available metadata.
+
+    For some filesystems this may be slow, as it involves a lot of I/O operations.
+
+    Parameters
+    ----------
+    file
+        File to parse
+    kwargs
+        Additional keyword arguments (not used, but required for compatibility)
+
+    Returns
+    -------
+    :
+        Dictionary with extracted metadata
+    """
+    keys = sorted(
+        {
+            "activity_id",
+            "branch_method",
+            "branch_time_in_child",
+            "branch_time_in_parent",
+            "experiment",
+            "experiment_id",
+            "frequency",
+            "grid",
+            "grid_label",
+            "institution_id",
+            "nominal_resolution",
+            "parent_activity_id",
+            "parent_experiment_id",
+            "parent_source_id",
+            "parent_time_units",
+            "parent_variant_label",
+            "realm",
+            "product",
+            "source_id",
+            "source_type",
+            "sub_experiment",
+            "sub_experiment_id",
+            "table_id",
+            "variable_id",
+            "variant_label",
+        }
+    )
+
+    try:
+        with xr.open_dataset(file, chunks={}, use_cftime=True) as ds:
+            info = {key: ds.attrs.get(key) for key in keys}
+            info["member_id"] = info["variant_label"]
+
+            variable_id = info["variable_id"]
+            if variable_id:  # pragma: no branch
+                attrs = ds[variable_id].attrs
+                for attr in ["standard_name", "long_name", "units"]:
+                    info[attr] = attrs.get(attr)
+
+            # Set the default of # of vertical levels to 1
+            vertical_levels = 1
+            start_time, end_time = None, None
+            init_year = None
+            try:
+                vertical_levels = ds[ds.cf["vertical"].name].size
+            except (KeyError, AttributeError, ValueError):
+                ...
+
+            try:
+                start_time, end_time = str(ds.cf["T"][0].data), str(ds.cf["T"][-1].data)
+            except (KeyError, AttributeError, ValueError):
+                ...
+            if info.get("sub_experiment_id"):  # pragma: no branch
+                init_year = extract_attr_with_regex(info["sub_experiment_id"], r"\d{4}")
+                if init_year:  # pragma: no cover
+                    init_year = int(init_year)
+            info["vertical_levels"] = vertical_levels
+            info["init_year"] = init_year
+            info["start_time"] = start_time
+            info["end_time"] = end_time
+            if not (start_time and end_time):
+                info["time_range"] = None
+            else:
+                info["time_range"] = f"{start_time}-{end_time}"
+        info["path"] = str(file)
+        info["version"] = extract_attr_with_regex(str(file), regex=r"v\d{4}\d{2}\d{2}|v\d{1}") or "v0"
+
+        # Mark the dataset as finalised
+        # This is used to indicate that the dataset has been fully parsed and is ready for use
+        info["finalised"] = True
+
+        return info
+
+    except Exception:
+        logger.exception(f"Failed to parse {file}")
+        return {"INVALID_ASSET": file, "TRACEBACK": traceback.format_exc()}
+
+
+def parse_cmip6_drs(file: str, **kwargs: Any) -> dict[str, Any]:
+    """
+    DRS parser for CMIP6 files
+
+    This parser extracts metadata according to the CMIP6 Data Reference Syntax (DRS).
+    This includes the essential metadata required to identify the dataset and is included in the filename.
+
+    Parameters
+    ----------
+    file
+        File to parse
+    kwargs
+        Additional keyword arguments (not used, but required for compatibility)
+
+    Returns
+    -------
+    :
+        Dictionary with extracted metadata
+    """
+    info: dict[str, Any] = parse_cmip6_using_directories(file)
+
+    if "INVALID_ASSET" in info:
+        logger.warning(f"Failed to parse {file}: {info['INVALID_ASSET']}")
+        return info
+
+    # The member_id is technically incorrect
+    # but for simplicity we are going to ignore sub-experiments for the DRS parser
+    info["variant_label"] = info["member_id"]
+
+    # Rename the `dcpp_init_year` key to `init_year` if it exists
+    if "dcpp_init_year" in info:
+        info["init_year"] = info.pop("dcpp_init_year")
+
+    if info.get("time_range"):
+        # Parse the time range if it exists
+        start_time, end_time = _parse_daterange(info["time_range"])
+        info["start_time"] = start_time
+        info["end_time"] = end_time
+
+    info["finalised"] = False
+
+    return info

climate_ref/datasets/obs4mips.py

@@ -15,8 +15,17 @@ from climate_ref.datasets.cmip6 import _parse_datetime
 from climate_ref.models.dataset import Dataset, Obs4MIPsDataset
 
 
-def parse_obs4mips(file: str) -> dict[str, Any | None]:
-    """Parser for obs4mips"""
+def parse_obs4mips(file: str, **kwargs: Any) -> dict[str, Any]:
+    """
+    Parser for obs4mips
+
+    Parameters
+    ----------
+    file
+        File to parse
+    kwargs
+        Additional keyword arguments (not used, but required for protocol compatibility)
+    """
     keys = sorted(
         list(
             {
@@ -106,6 +115,7 @@ class Obs4MIPsDatasetAdapter(DatasetAdapter):
 
     dataset_specific_metadata = (
         "activity_id",
+        "finalised",
         "frequency",
         "grid",
         "grid_label",
@@ -159,7 +169,7 @@ class Obs4MIPsDatasetAdapter(DatasetAdapter):
             depth=10,
             include_patterns=["*.nc"],
             joblib_parallel_kwargs={"n_jobs": self.n_jobs},
-        ).build(parsing_func=parse_obs4mips)  # type: ignore[arg-type]
+        ).build(parsing_func=parse_obs4mips)
 
         datasets = builder.df
         if datasets.empty:
@@ -178,4 +188,5 @@ class Obs4MIPsDatasetAdapter(DatasetAdapter):
         datasets["instance_id"] = datasets.apply(
             lambda row: "obs4MIPs." + ".".join([row[item] for item in drs_items]), axis=1
         )
+        datasets["finalised"] = True
         return datasets

climate_ref/executor/__init__.py

@@ -9,7 +9,14 @@ The simplest executor is the `LocalExecutor`, which runs the diagnostic in the s
 This is useful for local testing and debugging.
 """
 
-from .hpc import HPCExecutor
+from climate_ref_core.exceptions import InvalidExecutorException
+
+try:
+    from .hpc import HPCExecutor
+except InvalidExecutorException as exc:
+    # This exception is reraised when importing the executor as `climate_ref.executors.HPCExecutor`
+    HPCExecutor = exc  # type: ignore
+
 from .local import LocalExecutor
 from .result_handling import handle_execution_result
 from .synchronous import SynchronousExecutor

climate_ref/executor/hpc.py

@@ -5,12 +5,18 @@ If you want to
 - run REF under the HPC workflows
 - run REF in multiple nodes
 
+The `HPCExecutor` requires the optional `parsl` dependency.
+This dependency (and therefore this executor) is not available on Windows.
 """
 
 try:
     import parsl
 except ImportError:  # pragma: no cover
-    raise ImportError("The HPCExecutor requires the `parsl` package")
+    from climate_ref_core.exceptions import InvalidExecutorException
+
+    raise InvalidExecutorException(
+        "climate_ref_core.executor.hpc.HPCExecutor", "The HPCExecutor requires the `parsl` package"
+    )
 
 import os
 import time
@@ -21,7 +27,7 @@ from loguru import logger
 from parsl import python_app
 from parsl.config import Config as ParslConfig
 from parsl.executors import HighThroughputExecutor
-from parsl.launchers import SrunLauncher
+from parsl.launchers import SimpleLauncher, SrunLauncher
 from parsl.providers import SlurmProvider
 from tqdm import tqdm
 
@@ -34,6 +40,7 @@ from climate_ref_core.exceptions import DiagnosticError, ExecutionError
 from climate_ref_core.executor import execute_locally
 
 from .local import ExecutionFuture, process_result
+from .pbs_scheduler import SmartPBSProvider
 
 
 @python_app
@@ -96,8 +103,9 @@ class HPCExecutor:
         self.account = str(executor_config.get("account", os.environ.get("USER")))
         self.username = executor_config.get("username", os.environ.get("USER"))
         self.partition = str(executor_config.get("partition")) if executor_config.get("partition") else None
+        self.queue = str(executor_config.get("queue")) if executor_config.get("queue") else None
         self.qos = str(executor_config.get("qos")) if executor_config.get("qos") else None
-        self.req_nodes = int(executor_config.get("req_nodes", 1))
+        self.req_nodes = int(executor_config.get("req_nodes", 1)) if self.scheduler == "slurm" else 1
         self.walltime = str(executor_config.get("walltime", "00:10:00"))
         self.log_dir = str(executor_config.get("log_dir", "runinfo"))
 
@@ -181,21 +189,47 @@ class HPCExecutor:
     def _initialize_parsl(self) -> None:
         executor_config = self.config.executor.config
 
-        provider = SlurmProvider(
-            account=self.account,
-            partition=self.partition,
-            qos=self.qos,
-            nodes_per_block=self.req_nodes,
-            max_blocks=int(executor_config.get("max_blocks", 1)),
-            scheduler_options=executor_config.get("scheduler_options", "#SBATCH -C cpu"),
-            worker_init=executor_config.get("worker_init", "source .venv/bin/activate"),
-            launcher=SrunLauncher(
-                debug=True,
-                overrides=executor_config.get("overrides", ""),
-            ),
-            walltime=self.walltime,
-            cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
-        )
+        provider: SlurmProvider | SmartPBSProvider
+        if self.scheduler == "slurm":
+            provider = SlurmProvider(
+                account=self.account,
+                partition=self.partition,
+                qos=self.qos,
+                nodes_per_block=self.req_nodes,
+                max_blocks=int(executor_config.get("max_blocks", 1)),
+                scheduler_options=executor_config.get("scheduler_options", "#SBATCH -C cpu"),
+                worker_init=executor_config.get("worker_init", "source .venv/bin/activate"),
+                launcher=SrunLauncher(
+                    debug=True,
+                    overrides=executor_config.get("overrides", ""),
+                ),
+                walltime=self.walltime,
+                cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
+            )
+
+        elif self.scheduler == "pbs":
+            provider = SmartPBSProvider(
+                account=self.account,
+                queue=self.queue,
+                worker_init=executor_config.get("worker_init", "source .venv/bin/activate"),
+                nodes_per_block=_to_int(executor_config.get("nodes_per_block", 1)),
+                cpus_per_node=_to_int(executor_config.get("cpus_per_node", None)),
+                ncpus=_to_int(executor_config.get("ncpus", None)),
+                mem=executor_config.get("mem", "4GB"),
+                jobfs=executor_config.get("jobfs", "10GB"),
+                storage=executor_config.get("storage", ""),
+                init_blocks=executor_config.get("init_blocks", 1),
+                min_blocks=executor_config.get("min_blocks", 0),
+                max_blocks=executor_config.get("max_blocks", 1),
+                parallelism=executor_config.get("parallelism", 1),
+                scheduler_options=executor_config.get("scheduler_options", ""),
+                launcher=SimpleLauncher(),
+                walltime=self.walltime,
+                cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
+            )
+        else:
+            raise ValueError(f"Unsupported scheduler: {self.scheduler}")
+
         executor = HighThroughputExecutor(
             label="ref_hpc_executor",
             cores_per_worker=self.cores_per_worker if self.cores_per_worker else 1,
@@ -206,8 +240,11 @@ class HPCExecutor:
         )
 
         hpc_config = ParslConfig(
-            run_dir=self.log_dir, executors=[executor], retries=int(executor_config.get("retries", 2))
+            run_dir=self.log_dir,
+            executors=[executor],
+            retries=int(executor_config.get("retries", 2)),
         )
+
         parsl.load(hpc_config)
 
     def run(