PMP annual cycle implementation

changelog/221.feature.md

@@ -0,0 +1 @@
+Added PMP's annual cycle metrics

conftest.py

@@ -233,13 +233,16 @@ def _create_definition(
         metric_dataset: MetricDataset | None = None,
         cmip6: DatasetCollection | None = None,
         obs4mips: DatasetCollection | None = None,
+        pmp_climatology: DatasetCollection | None = None,
     ) -> MetricExecutionDefinition:
         if metric_dataset is None:
             datasets = {}
             if cmip6:
                 datasets[SourceDatasetType.CMIP6] = cmip6
             if obs4mips:
                 datasets[SourceDatasetType.obs4MIPs] = obs4mips
+            if pmp_climatology:
+                datasets[SourceDatasetType.PMPClimatology] = pmp_climatology
             metric_dataset = MetricDataset(datasets)
 
         return MetricExecutionDefinition(

packages/ref-metrics-pmp/pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
 ]
 
 [tool.uv.sources]
-pcmdi_metrics = { git = "https://github.com/PCMDI/pcmdi_metrics", rev="v3.9" }
+pcmdi_metrics = { git = "https://github.com/PCMDI/pcmdi_metrics", rev="v3.9.1" }
 
 [project.license]
 text = "Apache-2.0"

packages/ref-metrics-pmp/src/cmip_ref_metrics_pmp/__init__.py

@@ -6,6 +6,7 @@
 
 from cmip_ref_core.dataset_registry import dataset_registry_manager
 from cmip_ref_core.providers import CondaMetricsProvider
+from cmip_ref_metrics_pmp.annual_cycle import AnnualCycle
 from cmip_ref_metrics_pmp.variability_modes import ExtratropicalModesOfVariability
 
 __version__ = importlib.metadata.version("cmip_ref_metrics_pmp")
@@ -21,6 +22,7 @@
 provider.register(ExtratropicalModesOfVariability("PNA"))
 provider.register(ExtratropicalModesOfVariability("NPO"))
 provider.register(ExtratropicalModesOfVariability("SAM"))
+provider.register(AnnualCycle())
 
 
 dataset_registry_manager.register(

packages/ref-metrics-pmp/src/cmip_ref_metrics_pmp/annual_cycle.py

@@ -0,0 +1,239 @@
+import datetime
+import json
+from collections.abc import Iterable
+
+from loguru import logger
+
+from cmip_ref_core.datasets import FacetFilter, SourceDatasetType
+from cmip_ref_core.metrics import (
+    CommandLineMetric,
+    DataRequirement,
+    MetricExecutionDefinition,
+    MetricExecutionResult,
+)
+from cmip_ref_metrics_pmp.pmp_driver import build_glob_pattern, build_pmp_command, process_json_result
+
+
+class AnnualCycle(CommandLineMetric):
+    """
+    Calculate the annual cycle for a dataset
+    """
+
+    def __init__(self) -> None:
+        self.name = "PMP Annual Cycle"
+        self.slug = "pmp-annual-cycle"
+        self.data_requirements = (
+            DataRequirement(
+                source_type=SourceDatasetType.PMPClimatology,
+                # TODO: Add or operators and enable more varaiables
+                filters=(
+                    FacetFilter(
+                        facets={"source_id": ("GPCP-Monthly-3-2", "ERA-5"), "variable_id": ("pr", "ts")}
+                    ),
+                ),
+                group_by=("variable_id", "source_id"),
+            ),
+            DataRequirement(
+                source_type=SourceDatasetType.CMIP6,
+                # TODO: Add or operators and enable more varaiables
+                filters=(
+                    FacetFilter(
+                        facets={
+                            "frequency": "mon",
+                            "experiment_id": ("amip", "historical", "hist-GHG", "piControl"),
+                            "variable_id": ("pr", "ts"),
+                        }
+                    ),
+                ),
+                group_by=("variable_id", "source_id", "experiment_id", "variant_label", "member_id"),
+            ),
+        )
+
+        self.parameter_file_1 = "pmp_param_annualcycle_1-clims.py"
+        self.parameter_file_2 = "pmp_param_annualcycle_2-metrics.py"
+
+    def build_cmd(self, definition: MetricExecutionDefinition) -> Iterable[str]:
+        """
+        Build the command to run the metric
+
+        Parameters
+        ----------
+        definition
+            Definition of the metric execution
+
+        Returns
+        -------
+            Command arguments to execute in the PMP environment
+        """
+        raise NotImplementedError("Function not required")
+
+    def build_cmds(self, definition: MetricExecutionDefinition) -> list[list[str]]:
+        """
+        Build the command to run the metric
+
+        Parameters
+        ----------
+        definition
+            Definition of the metric execution
+
+        Returns
+        -------
+            Command arguments to execute in the PMP environment
+        """
+        input_datasets = definition.metric_dataset[SourceDatasetType.CMIP6]
+        source_id = input_datasets["source_id"].unique()[0]
+        experiment_id = input_datasets["experiment_id"].unique()[0]
+        member_id = input_datasets["member_id"].unique()[0]
+        variable_id = input_datasets["variable_id"].unique()[0]
+
+        logger.debug(f"input_datasets['source_id'].unique(): {input_datasets['source_id'].unique()}")
+        logger.debug(f"input_datasets['experiment_id'].unique(): {input_datasets['experiment_id'].unique()}")
+        logger.debug(f"input_datasets['member_id'].unique(): {input_datasets['member_id'].unique()}")
+        logger.debug(f"input_datasets['variable_id'].unique(): {input_datasets['variable_id'].unique()}")
+
+        model_files_raw = input_datasets.path.to_list()
+        if len(model_files_raw) == 1:
+            model_files = model_files_raw[0]  # If only one file, use it directly
+        elif len(model_files_raw) > 1:
+            model_files = build_glob_pattern(model_files_raw)  # If multiple files, build a glob pattern
+        else:
+            raise ValueError("No model files found")
+
+        logger.debug("build_cmd start")
+
+        logger.debug(f"input_datasets: {input_datasets}")
+        logger.debug(f"input_datasets.keys(): {input_datasets.keys()}")
+        logger.debug(f"input_datasets['variable_id']: {input_datasets['variable_id']}")
+
+        logger.debug(f"source_id: {source_id}")
+        logger.debug(f"experiment_id: {experiment_id}")
+        logger.debug(f"member_id: {member_id}")
+        logger.debug(f"variable_id: {variable_id}")
+
+        reference_dataset = definition.metric_dataset[SourceDatasetType.PMPClimatology]
+        reference_dataset_name = reference_dataset["source_id"].unique()[0]
+        reference_dataset_path = reference_dataset.datasets.iloc[0]["path"]
+
+        logger.debug(f"reference_dataset.datasets: {reference_dataset.datasets}")
+        logger.debug(f"reference_dataset['source_id']: {reference_dataset['source_id']}")
+        logger.debug(
+            f"reference_dataset.datasets.iloc[0]['path']: {reference_dataset.datasets.iloc[0]['path']}"
+        )
+
+        logger.debug(f"reference_dataset_name: {reference_dataset_name}")
+        logger.debug(f"reference_dataset_path: {reference_dataset_path}")
+
+        output_directory_path = str(definition.output_directory)
+
+        cmds = []
+
+        # ----------------------------------------------
+        # PART 1: Build the command to get climatologies
+        # ----------------------------------------------
+        # Model
+        data_name = f"{source_id}_{experiment_id}_{member_id}"
+        data_path = model_files
+        params = {
+            "driver_file": "mean_climate/pcmdi_compute_climatologies.py",
+            "parameter_file": self.parameter_file_1,
+            "vars": variable_id,
+            "infile": data_path,
+            "outfile": f"{output_directory_path}/{variable_id}_{data_name}_clims.nc",
+        }
+
+        cmds.append(build_pmp_command(**params))
+
+        # ----------------------------------------------
+        # PART 2: Build the command to calculate metrics
+        # ----------------------------------------------
+
+        # Reference
+        obs_dict = {
+            variable_id: {
+                reference_dataset_name: {
+                    "template": reference_dataset_path,
+                },
+                "default": reference_dataset_name,
+            }
+        }
+
+        # Generate a JSON file based on the obs_dict
+        with open(f"{output_directory_path}/obs_dict.json", "w") as f:
+            json.dump(obs_dict, f)
+
+        date = datetime.datetime.now().strftime("%Y%m%d")
+
+        params = {
+            "driver_file": "mean_climate/mean_climate_driver.py",
+            "parameter_file": self.parameter_file_2,
+            "vars": variable_id,
+            "custom_observations": f"{output_directory_path}/obs_dict.json",
+            "test_data_path": output_directory_path,
+            "test_data_set": source_id,
+            "realization": member_id,
+            "filename_template": f"{variable_id}_{data_name}_clims.198101-200512.AC.v{date}.nc",
+            "metrics_output_path": output_directory_path,
+            "cmec": "",
+        }
+
+        cmds.append(build_pmp_command(**params))
+
+        return cmds
+
+    def build_metric_result(self, definition: MetricExecutionDefinition) -> MetricExecutionResult:
+        """
+        Build a metric result from the output of the PMP driver
+
+        Parameters
+        ----------
+        definition
+            Definition of the metric execution
+
+        Returns
+        -------
+            Result of the metric execution
+        """
+        input_datasets = definition.metric_dataset[SourceDatasetType.CMIP6]
+        variable_id = input_datasets["variable_id"].unique()[0]
+
+        results_directory = definition.output_directory
+        png_directory = results_directory / variable_id
+        data_directory = results_directory / variable_id
+
+        results_files = list(results_directory.glob("*_cmec.json"))
+        if len(results_files) != 1:  # pragma: no cover
+            return MetricExecutionResult.build_from_failure(definition)
+
+        # Find the other outputs
+        png_files = list(png_directory.glob("*.png"))
+        data_files = list(data_directory.glob("*.nc"))
+
+        cmec_output, cmec_metric = process_json_result(results_files[0], png_files, data_files)
+
+        return MetricExecutionResult.build_from_output_bundle(
+            definition,
+            cmec_output_bundle=cmec_output,
+            cmec_metric_bundle=cmec_metric,
+        )
+
+    def run(self, definition: MetricExecutionDefinition) -> MetricExecutionResult:
+        """
+        Run the metric on the given configuration.
+
+        Parameters
+        ----------
+        definition : MetricExecutionDefinition
+            The configuration to run the metric on.
+
+        Returns
+        -------
+        :
+            The result of running the metric.
+        """
+        logger.debug("PMP annual cycle run start")
+        cmds = self.build_cmds(definition)
+
+        runs = [self.provider.run(cmd) for cmd in cmds]
+        logger.debug(f"runs: {runs}")
+
+        return self.build_metric_result(definition)

packages/ref-metrics-pmp/src/cmip_ref_metrics_pmp/params/pmp_param_annualcycle_1-clims.py

@@ -0,0 +1,21 @@
+#
+#  OPTIONS ARE SET BY USER IN THIS FILE AS INDICATED BELOW BY:
+#
+#
+
+# VARIABLES TO USE
+# vars = ["rlut"]
+
+# START AND END DATES FOR CLIMATOLOGY
+start = "1981-01"
+end = "2005-12"
+
+# INPUT DATASET - CAN BE MODEL OR OBSERVATIONS
+infile = (
+    "obs4MIPs_PCMDI_monthly/NASA-LaRC/CERES-EBAF-4-1"
+    "/mon/rlut/gn/v20210727"
+    "/rlut_mon_CERES-EBAF-4-1_PCMDI_gn_200301-201812.nc"
+)
+
+# DIRECTORY WHERE TO PUT RESULTS
+outfile = "climo/rlut_mon_CERES-EBAF-4-1_BE_gn.nc"

packages/ref-metrics-pmp/src/cmip_ref_metrics_pmp/params/pmp_param_annualcycle_2-metrics.py

@@ -0,0 +1,54 @@
+import os
+
+#
+#  OPTIONS ARE SET BY USER IN THIS FILE AS INDICATED BELOW BY:
+#
+#
+
+# RUN IDENTIFICATION
+# DEFINES A SUBDIRECTORY TO METRICS OUTPUT RESULTS SO MULTIPLE CASES CAN
+# BE COMPARED
+case_id = "basicTest"
+
+# LIST OF MODEL VERSIONS TO BE TESTED - WHICH ARE EXPECTED TO BE PART OF
+# CLIMATOLOGY FILENAME
+# test_data_set = ["ACCESS1-0", "CanCM4"]
+
+
+# VARIABLES TO USE
+# vars = ["rlut"]
+
+
+# Observations to use at the moment "default" or "alternate"
+reference_data_set = ["all"]
+# ext = '.nc'
+
+# INTERPOLATION OPTIONS
+target_grid = "2.5x2.5"  # OPTIONS: '2.5x2.5' or an actual cdms2 grid object
+regrid_tool = "regrid2"  # 'regrid2' # OPTIONS: 'regrid2','esmf'
+# -- OPTIONS: 'linear','conservative', only if tool is esmf
+regrid_method = "linear"
+regrid_tool_ocn = "esmf"  # OPTIONS: "regrid2","esmf"
+# -- OPTIONS: 'linear','conservative', only if tool is esmf
+regrid_method_ocn = "linear"
+
+# Templates for climatology files
+# %(param) will subsitute param with values in this file
+filename_template = "cmip5.historical.%(model_version).r1i1p1.mon.%(variable).198101-200512.AC.v20200426.nc"
+
+# filename template for landsea masks ('sftlf')
+sftlf_filename_template = "sftlf_%(model_version).nc"
+generate_sftlf = True  # if land surface type mask cannot be found, generate one
+
+# Region
+regions = {"rlut": ["Global"]}
+
+# ROOT PATH FOR MODELS CLIMATOLOGIES
+test_data_path = "demo_data_tmp/CMIP5_demo_clims/"
+# ROOT PATH FOR OBSERVATIONS
+# Note that atm/mo/%(variable)/ac will be added to this
+reference_data_path = ""
+custom_observations = "obs_dict.json"
+
+# DIRECTORY WHERE TO PUT RESULTS
+metrics_output_path = os.path.join("demo_output_tmp", "%(case_id)")