Rewrite to use pandas dataframes more widely in QRF

improver/calibration/__init__.py

@@ -24,7 +24,7 @@
         ("altitude", pa.float32()),
         ("blend_time", pa.timestamp("s", "utc")),
         ("forecast_period", pa.int64()),
-        ("forecast_reference_time", pa.int64()),
+        ("forecast_reference_time", pa.timestamp("s", "utc")),
         ("latitude", pa.float32()),
         ("longitude", pa.float32()),
         ("time", pa.timestamp("s", "utc")),
@@ -44,7 +44,7 @@
         ("latitude", pa.float32()),
         ("longitude", pa.float32()),
         ("altitude", pa.float32()),
-        ("time", pa.int64()),
+        ("time", pa.timestamp("s", "utc")),
         ("wmo_id", pa.string()),
         ("ob_value", pa.float32()),
     ]

improver/calibration/dataframe_utilities.py

@@ -76,8 +76,7 @@ def _dataframe_column_check(df: DataFrame, compulsory_columns: Sequence) -> None
     if not set(compulsory_columns).issubset(df.columns):
         diff = set(compulsory_columns).difference(df.columns)
         msg = (
-            "The following compulsory column(s) are missing from the "
-            f"DataFrame: {diff}"
+            f"The following compulsory column(s) are missing from the DataFrame: {diff}"
         )
         raise ValueError(msg)
 
@@ -183,11 +182,7 @@ def _fill_missing_entries(df, combi_cols, static_cols, site_id_col):
     df = df.set_index(combi_cols).reindex(new_index).reset_index(level=combi_cols)
 
     # Fill the NaNs within the static columns for each wmo_id.
-    filled_df = (
-        df.groupby(site_id_col)[combi_cols + static_cols]
-        .fillna(method="ffill")
-        .fillna(method="bfill")
-    )
+    filled_df = df.groupby(site_id_col)[combi_cols + static_cols].ffill().bfill()
     df = df.drop(columns=static_cols)
     df = df.merge(filled_df, on=combi_cols)
 
@@ -361,8 +356,7 @@ def get_forecast_representation(df: DataFrame) -> str:
         )
     if len(representations) == 0:
         raise ValueError(
-            f"None of the columns {REPRESENTATION_COLUMNS} "
-            "exist in the input dataset"
+            f"None of the columns {REPRESENTATION_COLUMNS} exist in the input dataset"
         )
     return representations.pop()
 

improver/calibration/load_and_apply_quantile_regression_random_forest.py

@@ -12,7 +12,9 @@
 import iris
 import joblib
 import numpy as np
+import pandas as pd
 from iris.cube import Cube, CubeList
+from iris.pandas import as_data_frame
 from quantile_forest import RandomForestQuantileRegressor
 
 from improver import PostProcessingPlugin
@@ -22,19 +24,20 @@
 )
 from improver.ensemble_copula_coupling.ensemble_copula_coupling import (
     RebadgePercentilesAsRealizations,
-    RebadgeRealizationsAsPercentiles,
 )
 from improver.ensemble_copula_coupling.utilities import choose_set_of_percentiles
 from improver.utilities.cube_checker import assert_spatial_coords_match
 
+iris.FUTURE.pandas_ndim = True
+
 
 class LoadAndApplyQRF(PostProcessingPlugin):
     """Load and apply the trained Quantile Regression Random Forest (QRF) model."""
 
     def __init__(
         self,
         feature_config: dict[str, list[str]],
-        target_cube_name: str,
+        target_cf_name: str,
         transformation: Optional[str] = None,
         pre_transform_addition: Optional[float] = None,
     ):
@@ -52,24 +55,24 @@ def __init__(
                 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:
+                "DYNAMIC_VARIABLE_CF_NAME": ["FEATURE1", "FEATURE2"] e.g:
                 {
                 "air_temperature": ["mean", "std", "altitude"],
                 "visibility_at_screen_level": ["mean", "std"]
                 "distance_to_water": ["static"],
                 }
-            target_cube_name (str):
-                A string containing the cube name of the forecast to be
-                calibrated. This will be used to separate it from the rest of the
-                dynamic predictors, if present.
+            target_cf_name (str):
+                A string containing the CF name of diagnostic to be calibrated. This
+                will be used to separate it from the rest of the dynamic predictors,
+                if present.
             transformation (str):
                 Transformation to be applied to the data before fitting.
             pre_transform_addition (float):
                 Value to be added before transformation.
 
         """
         self.feature_config = feature_config
-        self.target_cube_name = target_cube_name
+        self.target_cf_name = target_cf_name
         self.transformation = transformation
         self.pre_transform_addition = pre_transform_addition
 
@@ -110,7 +113,7 @@ def _get_inputs(
 
         # Extract all additional cubes which are associated with a feature in the
         # feature_config.
-        forecast_constraint = iris.Constraint(name=self.target_cube_name)
+        forecast_constraint = iris.Constraint(name=self.target_cf_name)
         forecast_cube = cube_inputs.extract(forecast_constraint)
 
         if forecast_cube:
@@ -119,7 +122,7 @@ def _get_inputs(
             msg = (
                 "No target forecast provided. An input file representing the target "
                 "must be provided, even if the target will not be used as a feature. "
-                f"The target is '{self.target_cube_name}'."
+                f"The target is '{self.target_cf_name}'."
             )
             raise ValueError(msg)
 
@@ -141,7 +144,7 @@ def _get_inputs(
             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():
+        if self.target_cf_name not in self.feature_config.keys():
             cube_inputs.remove(forecast_cube)
 
         return cube_inputs, forecast_cube, qrf_model
@@ -189,42 +192,40 @@ def _percentiles_to_realizations(cube_inputs: CubeList) -> CubeList:
         return realization_cube_inputs
 
     @staticmethod
-    def _organise_cubes(
-        cube_inputs: CubeList, forecast_cube: Cube
-    ) -> tuple[CubeList, Cube]:
-        """Promote the forecast period and forecast reference time coordinates to be
-        dimension coordinates, if present, on the feature cubes and the template
-        forecast cube.
+    def _cube_to_dataframe(cube_inputs: CubeList) -> pd.DataFrame:
+        """Convert cube inputs to a pandas DataFrame.
 
         Args:
-            cube_inputs: CubeList of feature cubes, which may include the forecast to be
-            forecast_cube: Forecast cube for use as a template.
-
+            cube_inputs: List of cubes containing the features and the forecast to be
+                calibrated.
         Returns:
-            Feature cubes and template cube with forecast period and
-            forecast reference time promoted to dimension coordinates.
+            DataFrame containing the data from the cubes, with auxiliary coordinates
+            included as columns.
         """
-        # Ensure that forecast_period is a dimension on all cubes.
-        fp_dim_cube_inputs = iris.cube.CubeList([])
-        for feature_cube in cube_inputs:
-            if feature_cube.coords("forecast_period", dim_coords=False):
-                feature_cube = iris.util.new_axis(feature_cube, "forecast_period")
-            if feature_cube.coords("forecast_reference_time", dim_coords=False):
-                feature_cube = iris.util.new_axis(
-                    feature_cube, "forecast_reference_time"
-                )
-            fp_dim_cube_inputs.append(feature_cube)
-        cube_inputs = fp_dim_cube_inputs
-
-        # Ensure the forecast cube has the same dimensions as the features
-        template_forecast_cube = iris.util.new_axis(forecast_cube, "forecast_period")
-        template_forecast_cube = iris.util.new_axis(
-            template_forecast_cube, "forecast_reference_time"
-        )
+        # Convert the first cube to a DataFrame.
+        df = as_data_frame(cube_inputs[0], add_aux_coords=True).reset_index()
+
+        # Iteratively convert remaining cubes to DataFrame and merge.
+        for cube in cube_inputs[1:]:
+            temporary_df = as_data_frame(cube, add_aux_coords=True).reset_index()
+            possible_columns = [
+                "wmo_id",
+                "time",
+                "forecast_reference_time",
+                "forecast_period",
+            ]
+            merge_columns = [
+                col for col in possible_columns if col in temporary_df.columns
+            ]
+            df = df.merge(
+                temporary_df[merge_columns + [cube.name()]],
+                on=merge_columns,
+                how="left",
+            )
 
-        # Check that the grids are the same for all dynamic predictors and the forecast
-        assert_spatial_coords_match(cube_inputs)
-        return cube_inputs, template_forecast_cube
+        for column in ["forecast_reference_time", "time"]:
+            df[column] = df[column].apply(lambda x: x._to_real_datetime())
+        return df
 
     def process(
         self,
@@ -254,24 +255,21 @@ def process(
         template_forecast_cube = forecast_cube.copy()
         if forecast_cube.coords("percentile"):
             percentiles = self._compute_percentiles(forecast_cube.copy(), "percentile")
-            cube_inputs = self._percentiles_to_realizations(cube_inputs.copy())
-            template_forecast_cube = RebadgePercentilesAsRealizations()(
-                template_forecast_cube
-            )
         elif forecast_cube.coords("realization"):
             percentiles = self._compute_percentiles(forecast_cube.copy(), "realization")
 
-        cube_inputs, template_forecast_cube = self._organise_cubes(
-            cube_inputs, template_forecast_cube
-        )
+        assert_spatial_coords_match(cube_inputs)
+        df = self._cube_to_dataframe(cube_inputs)
 
-        result = ApplyQuantileRegressionRandomForests(
+        calibrated_forecast = ApplyQuantileRegressionRandomForests(
+            target_name=self.target_cf_name,
             feature_config=self.feature_config,
             quantiles=percentiles,
             transformation=self.transformation,
             pre_transform_addition=self.pre_transform_addition,
-        )(qrf_model, cube_inputs, template_forecast_cube)
+        )(qrf_model, df)
+        calibrated_forecast_cube = template_forecast_cube.copy(
+            data=np.broadcast_to(calibrated_forecast.T, template_forecast_cube.shape)
+        )
 
-        if forecast_cube.coords("percentile"):
-            result = RebadgeRealizationsAsPercentiles()(result)
-        return result
+        return calibrated_forecast_cube