diff --git a/improver/calibration/__init__.py b/improver/calibration/__init__.py index fa6c58af99..9b2985e26a 100644 --- a/improver/calibration/__init__.py +++ b/improver/calibration/__init__.py @@ -533,6 +533,43 @@ def add_warning_comment(forecast: Cube) -> Cube: return forecast +def get_common_wmo_ids( + forecast_cube: Cube, + truth_cube: Cube, + additional_predictors: Optional[CubeList] = None, +) -> Tuple[Cube, Cube, CubeList]: + """Extracts the common WMO IDs from the forecast, truth and any additional + predictor cubes. + + Args: + forecast_cube: Cube containing the forecast data. + truth_cube: Cube containing the truth data. + additional_predictors: CubeList containing any additional predictors. + + Raises: + IOError: If no common WMO IDs are found in the input cubes. + + Returns: + The forecast, truth and additional predictor cubes with only the common + WMO IDs retained. + """ + wmo_ids = [] + wmo_ids.append(forecast_cube.coord("wmo_id").points) + wmo_ids.append(truth_cube.coord("wmo_id").points) + if additional_predictors is not None: + for ap in additional_predictors: + wmo_ids.append(ap.coord("wmo_id").points) + common_wmo_ids = list(set.intersection(*map(set, wmo_ids))) + if not common_wmo_ids: + raise IOError("No common WMO IDs found in the input cubes.") + constr = iris.Constraint(wmo_id=common_wmo_ids) + truth_cube = truth_cube.extract(constr) + forecast_cube = forecast_cube.extract(constr) + if additional_predictors is not None: + additional_predictors = additional_predictors.extract(constr) + return forecast_cube, truth_cube, additional_predictors + + def get_training_period_cycles( cycletime: str, forecast_period: Union[int, str], training_length: int ): diff --git a/improver/cli/estimate_emos_coefficients_from_table.py b/improver/cli/estimate_emos_coefficients_from_table.py index 1c5e2873da..24000363aa 100755 --- a/improver/cli/estimate_emos_coefficients_from_table.py +++ b/improver/cli/estimate_emos_coefficients_from_table.py @@ -115,9 +115,7 @@ def process( CubeList containing the coefficients estimated using EMOS. Each coefficient is stored in a separate cube. """ - import iris - from iris.cube import CubeList - + from improver.calibration import get_common_wmo_ids from improver.calibration.emos_calibration import ( EstimateCoefficientsForEnsembleCalibration, ) @@ -138,11 +136,9 @@ def process( return # Extract WMO IDs from the additional predictors. - if additional_predictors: - constr = iris.Constraint(wmo_id=truth_cube.coord("wmo_id").points) - additional_predictors = CubeList( - [ap.extract(constr) for ap in additional_predictors] - ) + forecast_cube, truth_cube, additional_predictors = get_common_wmo_ids( + forecast_cube, truth_cube, additional_predictors + ) plugin = EstimateCoefficientsForEnsembleCalibration( distribution, diff --git a/improver_tests/calibration/test_init.py b/improver_tests/calibration/test_init.py index 7092cadcde..a76693a743 100644 --- a/improver_tests/calibration/test_init.py +++ b/improver_tests/calibration/test_init.py @@ -18,6 +18,7 @@ from improver.calibration import ( add_warning_comment, + get_common_wmo_ids, get_training_period_cycles, identify_parquet_type, split_cubes_for_samos, @@ -30,6 +31,7 @@ from improver.synthetic_data.set_up_test_cubes import ( set_up_percentile_cube, set_up_probability_cube, + set_up_spot_variable_cube, set_up_variable_cube, ) from improver.utilities.save import save_netcdf @@ -871,6 +873,90 @@ def test_split_cubes_for_samos_basic( assert result_prob_template is None +@pytest.mark.parametrize( + "situation", + [ + "all_matching", + "all_matching_with_multiple_additional_predictors", + "fewer_in_forecast", + "fewer_in_truth", + "fewer_in_additional_predictors", + "no_additional_predictors", + "mixture", + "no_overlapping_sites", + ], +) +def test_get_common_wmo_ids(situation): + """Test the get_common_wmo_ids function.""" + forecast_wmo_ids = [1, 2, 3, 4, 5] + truth_wmo_ids = [1, 2, 3, 4, 5] + additional_wmo_ids = [1, 2, 3, 4, 5] + + if situation == "all_matching_with_multiple_additional_predictors": + additional_wmo_ids = [1, 2, 3, 4, 5, 6] + # A second 'additional predictor' cube will be added later + elif situation == "fewer_in_forecast": + forecast_wmo_ids = [1, 2, 3] + elif situation == "fewer_in_truth": + truth_wmo_ids = [1, 2, 3] + elif situation == "fewer_in_additional_predictors": + additional_wmo_ids = [1, 2, 3] + elif situation == "no_additional_predictors": + additional_wmo_ids = [] + elif situation == "mixture": + forecast_wmo_ids = [1, 2, 3, 4] + truth_wmo_ids = [1, 2, 3, 5] + additional_wmo_ids = [1, 2, 3, 6] + elif situation == "no_overlapping_sites": + forecast_wmo_ids = [1, 2] + truth_wmo_ids = [3, 4] + additional_wmo_ids = [5, 6] + + data = np.ones(len(forecast_wmo_ids), dtype=np.float32) + forecast_cube = set_up_spot_variable_cube(data, wmo_ids=forecast_wmo_ids) + data = np.ones(len(truth_wmo_ids), dtype=np.float32) + truth_cube = set_up_spot_variable_cube(data, wmo_ids=truth_wmo_ids) + + additional_predictors = None + if additional_wmo_ids: + data = np.ones(len(additional_wmo_ids), dtype=np.float32) + additional_predictors = CubeList( + [set_up_spot_variable_cube(data, wmo_ids=additional_wmo_ids)] + ) + # Add a second additional predictor cube to the 'additional_predictors' list + if situation == "all_matching_with_multiple_additional_predictors": + additional_predictors.append( + set_up_spot_variable_cube(data, wmo_ids=additional_wmo_ids) + ) + + if situation == "no_overlapping_sites": + with pytest.raises( + IOError, match="No common WMO IDs found in the input cubes." + ): + get_common_wmo_ids(forecast_cube, truth_cube, additional_predictors) + return + + forecast_result, truth_result, additional_predictor_result = get_common_wmo_ids( + forecast_cube, truth_cube, additional_predictors + ) + + if situation in [ + "all_matching", + "all_matching_with_multiple_additional_predictors", + "no_additional_predictors", + ]: + expected = [f"{x:05}" for x in [1, 2, 3, 4, 5]] + else: + expected = [f"{x:05}" for x in [1, 2, 3]] + assert forecast_result.coord("wmo_id").points.tolist() == expected + assert truth_result.coord("wmo_id").points.tolist() == expected + if additional_predictors is None: + assert additional_predictor_result is None + else: + for cube in additional_predictor_result: + assert cube.coord("wmo_id").points.tolist() == expected + + @pytest.mark.parametrize( "provide_emos_coefficients,expect_emos_coefficients,provide_emos_additional_fields,expect_emos_additional_fields", [