Alter site filtering for site EMOS

improver/calibration/__init__.py

@@ -533,6 +533,38 @@ 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.
+
+    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)
+    wmo_ids = list(set.intersection(*map(set, wmo_ids)))
+    constr = iris.Constraint(wmo_id=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
 ):

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,

improver_tests/calibration/test_init.py

@@ -18,6 +18,7 @@
 
 from improver.calibration import (
     add_warning_comment,
+    get_common_wmo_ids,
     identify_parquet_type,
     split_cubes_for_samos,
     split_forecasts_and_bias_files,
@@ -29,6 +30,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
@@ -870,6 +872,65 @@ def test_split_cubes_for_samos_basic(
     assert result_prob_template is None
 
 
+@pytest.mark.parametrize(
+    "situation",
+    [
+        "all_matching",
+        "fewer_in_forecast",
+        "fewer_in_truth",
+        "fewer_in_additional_predictors",
+        "no_additional_predictors",
+        "mixture",
+    ],
+)
+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 == "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]
+
+    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)
+    data = np.ones(len(additional_wmo_ids), dtype=np.float32)
+    additional_predictors = None
+    if additional_wmo_ids:
+        additional_predictors = CubeList(
+            [set_up_spot_variable_cube(data, wmo_ids=additional_wmo_ids)]
+        )
+
+    forecast_result, truth_result, additional_predictor_result = get_common_wmo_ids(
+        forecast_cube, truth_cube, additional_predictors
+    )
+
+    if situation in ["all_matching", "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:
+        assert (
+            additional_predictor_result[0].coord("wmo_id").points.tolist() == expected
+        )
+
+
 @pytest.mark.parametrize(
     "provide_emos_coefficients,expect_emos_coefficients,provide_emos_additional_fields,expect_emos_additional_fields",
     [