Merge pull request #11 from gavinevans/mobt_877_implement_qrf2

Add acceptance tests for QRF

Author: Kat-90

improver/calibration/load_and_apply_quantile_regression_random_forest.py

@@ -3,9 +3,11 @@
 #
 # This file is part of 'IMPROVER' and is released under the BSD 3-Clause license.
 # See LICENSE in the root of the repository for full licensing details.
-"""Script to load and apply the trained Quantile Regression Random Forest (QRF) model."""
+"""Script to load and apply the trained Quantile Regression Random Forest (QRF)
+model."""
 
 import pathlib
+from typing import Optional
 
 import iris
 import joblib
@@ -14,6 +16,7 @@
 from quantile_forest import RandomForestQuantileRegressor
 
 from improver import PostProcessingPlugin
+from improver.calibration import add_warning_comment
 from improver.calibration.quantile_regression_random_forest import (
     ApplyQuantileRegressionRandomForests,
 )
@@ -31,8 +34,8 @@ def __init__(
         self,
         feature_config: dict[str, list[str]],
         target_cube_name: str,
-        transformation: str = None,
-        pre_transform_addition: float = None,
+        transformation: Optional[str] = None,
+        pre_transform_addition: Optional[float] = None,
     ):
         """Initialise the plugin.
 
@@ -70,13 +73,13 @@ def __init__(
         self.pre_transform_addition = pre_transform_addition
 
     def _get_inputs(
-        self, file_paths: pathlib.Path
+        self, file_paths: list[pathlib.Path]
     ) -> tuple[CubeList, Cube, RandomForestQuantileRegressor]:
         """Get inputs from disk and separate the model and the features.
 
         Args:
-            file_paths: Path to the trained QRF model and the forecast to be calibrated,
-                and the features, as required.
+            file_paths: List of paths to the trained QRF model and the forecast to be
+                calibrated and the features, as required.
 
         Returns:
             CubeList of the features cubes, the forecast cube, and the
@@ -120,6 +123,7 @@ def _get_inputs(
             raise ValueError(msg)
 
         if not qrf_model:
+            forecast_cube = add_warning_comment(forecast_cube)
             return None, forecast_cube, None
 
         if len(cube_inputs) != len(self.feature_config.keys()):
@@ -134,7 +138,7 @@ def _get_inputs(
         if not qrf_model:
             # The specified model doesn't exist and the forecast will not be calibrated
             return forecast_cube
-        
+
         # If target diagnostic not a feature in the training then remove.
         if self.target_cube_name not in self.feature_config.keys():
             cube_inputs.remove(forecast_cube)
@@ -160,7 +164,7 @@ def _compute_percentiles(forecast_cube: Cube, coord: str) -> list[float]:
         return percentiles
 
     @staticmethod
-    def _percentiles_to_realizations(cube_inputs: Cube) -> CubeList:
+    def _percentiles_to_realizations(cube_inputs: CubeList) -> CubeList:
         """Convert percentiles to realizations. The input forecasts are expected to
         be percentiles but these percentiles are rebadged as realizations.
 
@@ -174,8 +178,8 @@ def _percentiles_to_realizations(cube_inputs: Cube) -> CubeList:
                 where appropriate
         """
 
-        # Ensure there is a realization dimension on all cubes. This assumes a percentile
-        # dimension is present.
+        # Ensure there is a realization dimension on all cubes. This assumes a
+        # percentile dimension is present.
         realization_cube_inputs = iris.cube.CubeList([])
         for feature_cube in cube_inputs:
             if feature_cube.coords("percentile"):
@@ -224,22 +228,20 @@ def _organise_cubes(
 
     def process(
         self,
-        file_paths: pathlib.Path,
+        file_paths: list[pathlib.Path],
     ) -> Cube:
-        """Loading and applying the trained model for Quantile Regression Random Forest.
-
-        Load in the previously trained model for Quantile Regression Random
-        Forest (QRF). The model is applied to the forecast that is supplied,
-        so as to calibrate the forecast. The calibrated forecast is written
-        to a cube. If no model is provided the input forecast is returned unchanged.
+        """Load and applying the trained Quantile Regression Random Forest (QRF) model.
+        The model is applied to the forecast supplied to calibrate the forecast.
+        The calibrated forecast is written to a cube. If no model is provided the
+        input forecast is returned unchanged.
 
         Args:
             file_paths (cli.inputpaths):
                 A list of input paths containing:
                 - The path to a QRF trained model in pickle file format to be used
                 for calibration.
                 - The path to a NetCDF file containing the forecast to be calibrated.
-                - Optionally, paths to NetCDF files containing additional preictors.
+                - Optionally, paths to NetCDF files containing additional predictors.
 
         Returns:
             iris.cube.Cube:

improver/calibration/load_and_train_quantile_regression_random_forest.py

@@ -2,10 +2,12 @@
 #
 # This file is part of 'IMPROVER' and is released under the BSD 3-Clause license.
 # See LICENSE in the root of the repository for full licensing details.
-"""Script to load inputs and train a model using Quantile Regression Random Forest (QRF)."""
+"""Script to load inputs and train a model using Quantile Regression Random Forest
+(QRF)."""
 
 import pathlib
 from pathlib import Path
+from typing import Optional
 
 import iris
 import numpy as np
@@ -56,8 +58,24 @@ def __init__(
         self.pre_transform_addition = pre_transform_addition
         self.compression = compression
 
-    def _split_cubes_and_parquet_files(self, file_paths):
-        """Split the input file paths into cubes and parquet files."""
+    def _split_cubes_and_parquet_files(
+        self, file_paths: list[pathlib.Path | str]
+    ) -> tuple[Optional[pathlib.Path], Optional[pathlib.Path], iris.cube.CubeList]:
+        """Split the input file paths into cubes and parquet files.
+
+        Args:
+            file_paths: List of file paths.
+
+        Returns:
+            Tuple containing the items below if found:
+                - Path to the forecast parquet file.
+                - Path to the truth parquet file.
+                - List of cubes loaded from the NetCDF files.
+
+        Raises:
+            ValueError: If the number of cubes loaded does not match the number of
+                features expected.
+        """
 
         forecast_table_path = None
         truth_table_path = None
@@ -69,7 +87,8 @@ def _split_cubes_and_parquet_files(self, file_paths):
                 cube = load_cube(str(file_path))
                 cube_inputs.append(cube)
             except IsADirectoryError:
-                # For loop here because the read_schema must read a .parquet file rather than a directory.
+                # For loop here because the read_schema must read a .parquet file
+                # rather than a directory.
                 for file in Path(file_path).glob("**/*.parquet"):
                     try:
                         pq.read_schema(file).field("forecast_period")
@@ -79,17 +98,18 @@ def _split_cubes_and_parquet_files(self, file_paths):
                     if forecast_table_path and truth_table_path:
                         break
             except OSError:
-                print("The directory doesn't exist, calibration is skipped for this cycle")
-                return
+                # This will occur when the filepath does not exist. In this case,
+                # return None.
+                return None, None, None
 
         if len(self.feature_config.keys()) not in [
             len(cube_inputs),
             len(cube_inputs) + 1,
         ]:
             msg = (
                 "The number of cubes loaded does not match the number of features "
-                "expected. These can mismatch if the some features are coming from the "
-                "historic forecast. The number of cubes loaded was: "
+                "expected. These can mismatch if some features are coming from the "
+                "historic forecast table. The number of cubes loaded was: "
                 f"{len(cube_inputs)}. The number of features expected was: "
                 f"{len(self.feature_config.keys())}."
             )
@@ -98,9 +118,29 @@ def _split_cubes_and_parquet_files(self, file_paths):
         return forecast_table_path, truth_table_path, cube_inputs
 
     def _read_parquet_files(
-        self, forecast_table_path, truth_table_path, forecast_periods
-    ):
-        """Read the forecast and truth data from parquet files."""
+        self,
+        forecast_table_path: pathlib.Path | str,
+        truth_table_path: pathlib.Path | str,
+        forecast_periods: list[int],
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
+        """Read the forecast and truth data from parquet files.
+
+        Args:
+            forecast_table_path: Path to the forecast parquet file.
+            truth_table_path: Path to the truth parquet file.
+            forecast_periods: List of forecast periods in seconds.
+
+        Returns:
+            Tuple containing:
+                - DataFrame containing the forecast data.
+                - DataFrame containing the truth data.
+
+        Raises:
+            ValueError: If the forecast parquet file does not contain the expected
+                fields.
+            ValueError: If the truth parquet file does not contain the expected
+                fields.
+        """
         cycletimes = []
 
         for forecast_period in forecast_periods:
@@ -151,7 +191,8 @@ def _read_parquet_files(
             engine="pyarrow",
         )
 
-        # Convert df columns from ms to pandas timestamp object to work with existing code
+        # Convert df columns from ms to pandas timestamp object to work with existing
+        # code
         for column in ["time", "forecast_reference_time", "blend_time"]:
             forecast_df[column] = pd.to_datetime(
                 forecast_df[column], unit="ns", utc=True
@@ -166,7 +207,7 @@ def _read_parquet_files(
         truth_df = pd.read_parquet(
             truth_table_path, filters=filters, schema=TRUTH_SCHEMA, engine="pyarrow"
         )
-        
+
         truth_df["time"] = pd.to_datetime(truth_df["time"], unit="ns", utc=True)
 
         if truth_df.empty:
@@ -177,9 +218,28 @@ def _read_parquet_files(
             raise IOError(msg)
         return forecast_df, truth_df
 
-    def _dataframe_to_cubes(self, forecast_df, truth_df, forecast_periods):
+    def _dataframe_to_cubes(
+        self,
+        forecast_df: pd.DataFrame,
+        truth_df: pd.DataFrame,
+        forecast_periods: list[int],
+    ) -> tuple[iris.cube.Cube, iris.cube.Cube]:
         """Convert the forecast and truth dataframes to cubes at each forecast period
-        required."""
+        required.
+
+        Args:
+            forecast_df: DataFrame containing the forecast data.
+            truth_df: DataFrame containing the truth data.
+            forecast_periods: List of forecast periods in seconds.
+
+        Returns:
+            Tuple containing:
+                - Cube containing the forecast data.
+                - Cube containing the truth data.
+
+        Raises:
+            ValueError: The forecast has failed to concatenate into a single cube.
+        """
         forecast_cubes = iris.cube.CubeList([])
         truth_cubes = iris.cube.CubeList([])
 
@@ -226,14 +286,28 @@ def _dataframe_to_cubes(self, forecast_df, truth_df, forecast_periods):
         return forecast_cube, truth_cube
 
     @staticmethod
-    def filter_bad_sites(forecast_cube, truth_cube, cube_inputs):
-        """Remove sites that have NaNs in the data."""
-        bad_site_ids = []
+    def filter_bad_sites(
+        forecast_cube: iris.cube.Cube,
+        truth_cube: iris.cube.Cube,
+        cube_inputs: iris.cube.CubeList,
+    ) -> tuple[iris.cube.Cube, iris.cube.Cube, iris.cube.CubeList]:
+        """Remove sites that have NaNs in the data.
+
+        Args:
+            forecast_cube: Cube containing the forecast data.
+            truth_cube: Cube containing the truth data.
+            cube_inputs: List of additional feature cubes.
+
+        Returns:
+            Tuple containing:
+                - Cube containing the forecast data with bad sites removed.
+                - Cube containing the truth data with bad sites removed.
+                - List of additional feature cubes with bad sites removed.
+        """
         nan_mask = np.any(np.isnan(truth_cube.data), axis=truth_cube.coord_dims("time"))
         all_site_ids = truth_cube.coord("wmo_id").points
         bad_site_ids = all_site_ids[nan_mask]
-        wmo_ids = set(all_site_ids) - set(bad_site_ids)
-        constr = iris.Constraint(wmo_id=lambda cell: cell in wmo_ids)
+        constr = iris.Constraint(wmo_id=lambda cell: cell not in bad_site_ids.tolist())
         truth_cube = truth_cube.extract(constr)
         forecast_cube = forecast_cube.extract(constr)
         feature_cube_inputs = iris.cube.CubeList([])
@@ -247,13 +321,11 @@ def process(
         self,
         file_paths: list[pathlib.Path | str],
         model_output: str = None,
-    ):
-        """Loading input files and training a model using Quantile Regression Random Forest.
-
-        Loads in arguments for training a Quantile Regression Random Forest (QRF)
-        model which can later be applied to calibrate the forecast.
-        Two sources of input data must be provided: historical forecasts and
-        historical truth data (to use in calibration). The model is output as a pickle file.
+    ) -> None:
+        """Load input files and training a Quantile Regression Random Forest (QRF)
+        model. This model can be applied later to calibrate the forecast. Two sources
+        of input data must be provided: historical forecasts and historical truth data
+        (to use in calibration). The model is output as a pickle file.
 
         Args:
             file_paths (cli.inputpaths):
@@ -266,15 +338,16 @@ def process(
                 for calibration.
                 - Optionally, paths to NetCDF files containing additional predictors.
             feature_config (dict):
-                Feature configuration defining the features to be used for quantile regression.
-                    The configuration is a dictionary of strings, where the keys are the names of
-                    the input cube(s) supplied, and the values are a list. This list can contain both
-                    computed features, such as the mean or standard deviation (std), or static
-                    features, such as the altitude. The computed features will be computed using
-                    the cube defined in the dictionary key. If the key is the feature itself e.g.
-                    a distance to water cube, then the value should state "static". This will ensure
-                    the cube's data is used as the feature.
-                    The config will have the structure:
+                Feature configuration defining the features to be used for quantile
+                regression. The configuration is a dictionary of strings, where the
+                keys are the names of the input cube(s) supplied, and the values are
+                a list. This list can contain both computed features, such as the mean
+                or standard deviation (std), or static features, such as the altitude.
+                The computed features will be computed using the cube defined in the
+                dictionary key. If the key is the feature itself e.g. a distance to
+                water cube, then the value should state "static". This will ensure
+                the cube's data is used as the feature. The config will have the
+                structure:
                     "DYNAMIC_VARIABLE_NAME": ["FEATURE1", "FEATURE2"] e.g:
                     {
                     "air_temperature": ["mean", "std", "altitude"],
@@ -286,8 +359,10 @@ def process(
                 calibrated. This will be used to filter the target forecast and truth
                 dataframes.
             forecast_period (int):
-                Range of forecast periods to be including in training in hours in the
-                form: "start:end:interval" e.g. "6:18:6".
+                Range of forecast periods to be calibrated in hours in the form:
+                "start:end:interval" e.g. "6:18:6" or a single forecast period e.g. "6".
+                The end value is exclusive, so "6:18:6" will calibrate the 6 and 12
+                hours.
             cycletime (str):
                 Cycletime of the forecast to be calibrated in a format similar to
                 20170109T0000Z. This is used to filter the correct blendtimes from
@@ -321,8 +396,19 @@ def process(
         if not forecast_table_path or not truth_table_path:
             return None
 
-        forecast_periods = list(range(*map(int, self.forecast_periods.split(":"))))
-        forecast_periods = [fp * 3600 for fp in forecast_periods]
+        if ":" in self.forecast_periods:
+            forecast_periods = list(range(*map(int, self.forecast_periods.split(":"))))
+            forecast_periods = [fp * 3600 for fp in forecast_periods]
+        else:
+            try:
+                forecast_periods = [int(self.forecast_periods) * 3600]
+            except ValueError:
+                msg = (
+                    "The forecast_periods argument must be a single integer or "
+                    "a range in the form 'start:end:interval'. The forecast period"
+                    f"provided was: {self.forecast_periods}."
+                )
+                raise ValueError(msg)
 
         forecast_df, truth_df = self._read_parquet_files(
             forecast_table_path, truth_table_path, forecast_periods