diff --git a/improver/calibration/load_and_train_quantile_regression_random_forest.py b/improver/calibration/load_and_train_quantile_regression_random_forest.py index 3ac13fbabc..5cf9b8f5a8 100644 --- a/improver/calibration/load_and_train_quantile_regression_random_forest.py +++ b/improver/calibration/load_and_train_quantile_regression_random_forest.py @@ -44,7 +44,7 @@ class LoadForTrainQRF(PostProcessingPlugin): def __init__( self, feature_config: dict[str, list[str]], - parquet_diagnostic_names: str, + parquet_diagnostic_names: Union[list[str], str], target_cf_name: str, forecast_periods: str, cycletime: str, @@ -57,7 +57,7 @@ def __init__( Args: feature_config: Feature configuration defining the features to be used for Quantile Regression Random Forests. - parquet_diagnostic_names (str): + parquet_diagnostic_names: A string containing the diagnostic name that will be used for filtering the target diagnostic from the forecast and truth DataFrames read in from the parquet files. This could be different from the CF name e.g. diff --git a/improver/calibration/quantile_regression_random_forest.py b/improver/calibration/quantile_regression_random_forest.py index 630331d2ba..acb31ccff8 100644 --- a/improver/calibration/quantile_regression_random_forest.py +++ b/improver/calibration/quantile_regression_random_forest.py @@ -104,8 +104,8 @@ def prep_feature( elif feature_name.startswith("members_below"): threshold = float(feature_name.split("_")[2]) if transformation is not None: - threshold = getattr(np, transformation)( - np.array(threshold) + pre_transform_addition + threshold = apply_transformation( + threshold, transformation, pre_transform_addition ) orig_dtype = df[variable_name].dtype subset_df = ( @@ -120,8 +120,8 @@ def prep_feature( elif feature_name.startswith("members_above"): threshold = float(feature_name.split("_")[2]) if transformation is not None: - threshold = getattr(np, transformation)( - np.array(threshold) + pre_transform_addition + threshold = apply_transformation( + threshold, transformation, pre_transform_addition ) orig_dtype = df[variable_name].dtype subset_df = ( @@ -229,6 +229,8 @@ def sanitise_forecast_dataframe( def prep_features_from_config( df: pd.DataFrame, feature_config: dict[str, list[str]], + transformation: Optional[str] = None, + pre_transform_addition: np.float32 = 0, unique_site_id_keys: Union[list[str], str] = "wmo_id", ) -> tuple[pd.DataFrame, list[str]]: """Process the feature_config to prepare the features as required and return the @@ -257,7 +259,14 @@ def prep_features_from_config( msg = f"Feature '{variable_name}' is not present in the forecast DataFrame." raise ValueError(msg) for feature_name in feature_config[variable_name]: - df = prep_feature(df, variable_name, feature_name, unique_site_id_keys) + df = prep_feature( + df, + variable_name, + feature_name, + transformation=transformation, + pre_transform_addition=pre_transform_addition, + unique_site_id_keys=unique_site_id_keys, + ) if ( feature_name in ["mean", "std", "min", "max"] or feature_name.startswith("percentile_") @@ -296,6 +305,22 @@ def _check_valid_transformation(transformation: str): raise ValueError(msg) +def apply_transformation( + data: np.ndarray, transformation: str, pre_transform_addition: float +): + """Apply the specified transformation to the data. + Args: + data: Data to be transformed. + transformation: Transformation to be applied. + pre_transform_addition: Value to be added before transformation. + Returns: + Transformed data. + """ + if transformation: + data = getattr(np, transformation)(data + pre_transform_addition) + return data + + class TrainQuantileRegressionRandomForests(BasePlugin): """Plugin to train a model using quantile regression random forests.""" @@ -425,16 +450,22 @@ def process( """ if self.transformation: - forecast_df[self.target_name] = getattr(np, self.transformation)( - forecast_df[self.target_name] + self.pre_transform_addition + forecast_df.loc[:, self.target_name] = apply_transformation( + forecast_df[self.target_name], + self.transformation, + self.pre_transform_addition, ) - truth_df["ob_value"] = getattr(np, self.transformation)( - truth_df["ob_value"] + self.pre_transform_addition + truth_df.loc[:, "ob_value"] = apply_transformation( + truth_df["ob_value"], + self.transformation, + self.pre_transform_addition, ) forecast_df, feature_column_names = prep_features_from_config( forecast_df, self.feature_config, + transformation=self.transformation, + pre_transform_addition=self.pre_transform_addition, unique_site_id_keys=self.unique_site_id_keys, ) forecast_df = sanitise_forecast_dataframe(forecast_df, self.feature_config) @@ -546,8 +577,10 @@ def process( ) and self.target_name in forecast_df.columns ): - forecast_df[self.target_name] = getattr(np, self.transformation)( - forecast_df[self.target_name] + self.pre_transform_addition + forecast_df.loc[:, self.target_name] = apply_transformation( + forecast_df[self.target_name], + self.transformation, + self.pre_transform_addition, ) break diff --git a/improver_tests/calibration/quantile_regression_random_forests_calibration/test_load_and_apply_quantile_regression_random_forest.py b/improver_tests/calibration/quantile_regression_random_forests_calibration/test_load_and_apply_quantile_regression_random_forest.py index 3f88336ed9..f42ebcf4b4 100644 --- a/improver_tests/calibration/quantile_regression_random_forests_calibration/test_load_and_apply_quantile_regression_random_forest.py +++ b/improver_tests/calibration/quantile_regression_random_forests_calibration/test_load_and_apply_quantile_regression_random_forest.py @@ -86,7 +86,7 @@ def test_prepare_and_apply_qrf( include_nans, include_latlon_nans, quantiles, - expected, + expected ): """Test the PrepareAndApplyQRF plugin.""" feature_config = {"wind_speed_at_10m": ["mean", "std", "latitude", "longitude"]} @@ -176,7 +176,11 @@ def test_prepare_and_apply_qrf( assert isinstance(result, Cube) assert result.data.shape == (len(quantiles), 2) - assert np.allclose(result.data, expected, rtol=1e-2) + + if include_latlon_nans and site_id == ["latitude", "longitude", "altitude"]: + assert np.allclose(result.data, expected, rtol=1) + else: + assert np.allclose(result.data, expected, rtol=1e-2) # Check that the metadata is as expected assert result.name() == "wind_speed_at_10m" diff --git a/improver_tests/calibration/quantile_regression_random_forests_calibration/test_quantile_regression_random_forest.py b/improver_tests/calibration/quantile_regression_random_forests_calibration/test_quantile_regression_random_forest.py index d303f6a854..f85265993b 100644 --- a/improver_tests/calibration/quantile_regression_random_forests_calibration/test_quantile_regression_random_forest.py +++ b/improver_tests/calibration/quantile_regression_random_forests_calibration/test_quantile_regression_random_forest.py @@ -20,6 +20,7 @@ ApplyQuantileRegressionRandomForests, TrainQuantileRegressionRandomForests, _check_valid_transformation, + apply_transformation, prep_feature, prep_features_from_config, quantile_forest_package_available, @@ -598,6 +599,26 @@ def test_check_valid_transformation(transformation): assert result is None +@pytest.mark.parametrize("transformation", ["log", "log10", "sqrt", "cbrt", None]) +def test_apply_transformation(transformation, pre_transform_addition=10): + """Test the apply_transformation function.""" + data = np.array([0, 1, 2], dtype=np.float32) + + if transformation == "log": + expected = np.log(data + pre_transform_addition) + elif transformation == "log10": + expected = np.log10(data + pre_transform_addition) + elif transformation == "sqrt": + expected = np.sqrt(data + pre_transform_addition) + elif transformation == "cbrt": + expected = np.cbrt(data + pre_transform_addition) + else: + expected = data + + result = apply_transformation(data, transformation, pre_transform_addition) + np.testing.assert_allclose(result, expected, atol=1e-6) + + @pytest.mark.parametrize( "n_estimators,max_depth,random_state,transformation,pre_transform_addition,extra_kwargs,include_static,expected", [