Patch the reading in of windspeed

src/CSET/operators/constraints.py

@@ -52,7 +52,7 @@ def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint:
     return stash_constraint
 
 
-def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
+def generate_var_constraint(varname: str | list[str], **kwargs) -> iris.Constraint:
     """Generate constraint from variable name or STASH code.
 
     Operator that takes a CF compliant variable name string or list of names, and generates an
@@ -77,12 +77,13 @@ def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
         return iris.AttributeConstraint(STASH=varname)
 
     # Ensure access to variable vector components for computed fields
-    if "wind_speed_at_10m" in iter_maybe(varname):
+    varname_copy = iter_maybe(varname)[:]
+
+    if "wind_speed_at_10m" in varname_copy:
         if isinstance(varname, str):
             varname = [varname]
         varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"])
         varname.extend(["u_wind_at_10m", "v_wind_at_10m"])
-
     # Case 2: Multiple varnames
     if isinstance(varname, (list, tuple)):
         varname_constraint = iris.Constraint(
@@ -96,6 +97,8 @@ def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
     else:
         varname_constraint = iris.Constraint(name=varname)
 
+    varname_constraint.varname = varname_copy
+
     return varname_constraint
 
 

src/CSET/operators/read.py

@@ -195,7 +195,6 @@ def read_cubes(
     for cube in cubes:
         # Use 1 to indicate True, as booleans can't be saved in NetCDF attributes.
         cube.attributes["cset_comparison_base"] = 1
-
     # Load the rest of the models.
     cubes.extend(itertools.chain.from_iterable(model_cubes))
 
@@ -247,7 +246,7 @@ def _load_model(
     logger.debug("Constraint: %s", constraint)
     cubes = iris.load(input_files, constraint, callback=_loading_callback)
     # If required, compute wind_speed from components.
-    cubes = _compute_winds(cubes)
+    cubes = _compute_winds(cubes, constraint)
 
     # Add model_name attribute to each cube to make it available at any further
     # step without needing to pass it as function parameter.
@@ -900,7 +899,17 @@ def _fix_lfric_cloud_base_altitude(cube: iris.cube.Cube):
         cube.data = dask.array.ma.masked_greater(cube.core_data(), 144.0)
 
 
-def _compute_winds(cubes: iris.cube.CubeList):
+def get_filter_windspeed(constraint: iris.Constraint):
+    """Get the windspeed filter by using the hijacked constraint."""
+    if hasattr(constraint, "varname"):
+        return constraint.varname
+    else:
+        return None
+
+
+def _compute_winds(
+    cubes: iris.cube.CubeList, constraint: iris.Constraint | None = None
+):
     """To compute wind_speed from vector components if not available as diagnostic.
 
     Diagnostics of wind are also not always consistent between the UM
@@ -915,26 +924,38 @@ def _compute_winds(cubes: iris.cube.CubeList):
     # the cell methods, but it may not be warranted.
     #
     # A check on UM STASH attributes is also conducted to adjust directions.
+
+    if constraint is None:
+        return cubes
+
+    filter_windspeed = get_filter_windspeed(constraint)
+
     u_constr = iris.Constraint("eastward_wind_at_10m")
     v_constr = iris.Constraint("northward_wind_at_10m")
-    speed_constr = iris.Constraint("wind_speed_at_10m")
+    sp_constr = iris.Constraint("wind_speed_at_10m")
+
     try:
-        if cubes.extract(u_constr) and cubes.extract(v_constr):
-            if len(cubes) == 2:
-                wind_only = True
-            else:
-                wind_only = False
-            if len(cubes.extract(u_constr)) == 1 and not cubes.extract(speed_constr):
+        if (
+            cubes.extract(u_constr)
+            and cubes.extract(v_constr)
+            and not cubes.extract(sp_constr)
+        ):
+            if "wind_speed_at_10m" in constraint.varname:
                 _add_wind_speed_um(cubes)
             # Convert winds in the UM to be relative to true east and true north.
-            if cubes.extract(u_constr) and cubes.extract(v_constr):
-                _convert_wind_true_dirn_um(cubes)
-            # Return only wind_speed cube
-            if wind_only:
-                cubes = cubes.extract(speed_constr)
+            _convert_wind_true_dirn_um(cubes)
     except (KeyError, AttributeError):
         pass
 
+    if filter_windspeed:
+        filter_windspeed_constraint = iris.Constraint(
+            cube_func=lambda cube: (
+                cube.long_name in filter_windspeed
+                or cube.standard_name in filter_windspeed
+                or cube.var_name in filter_windspeed
+            )
+        )
+        cubes = cubes.extract(filter_windspeed_constraint)
     return cubes
 
 
@@ -1109,7 +1130,7 @@ def _fix_no_time_coords_callback(cube: iris.cube.Cube):
 
 
 def _normalise_longname(cube: iris.cube.Cube):
-    """Normalise long_name to the LFRic standard list."""
+    """Fix plev variable names to standard names."""
     if cube.coords("pressure"):
         if cube.name() == "x_wind":
             cube.long_name = "zonal_wind_at_pressure_levels"
@@ -1126,8 +1147,6 @@ def _normalise_longname(cube: iris.cube.Cube):
             cube.long_name = "eastward_wind_at_10m"
         if cube.name() == "y_wind" and cube.var_name == "v_wind_at_10m":
             cube.long_name = "northward_wind_at_10m"
-    if cube.name() == "air_pressure_at_sea_level":
-        cube.long_name = "air_pressure_at_mean_sea_level"
 
 
 def _check_combine_point_observations(cubes: iris.cube.CubeList):

tests/conftest.py

@@ -1355,3 +1355,55 @@ def _make_cube(data, long_name, model_name=None):
         return cube
 
     return _make_cube
+
+
+@pytest.fixture()
+def wind_cubelist_um():
+    """UM wind cube list."""
+    wind_cubelist = CubeList()
+
+    wind_cubelist.append(
+        Cube(
+            3.5, standard_name="x_wind", long_name="eastward_wind_at_10m", units="m s-1"
+        )
+    )
+    wind_cubelist.append(
+        Cube(
+            3.5,
+            standard_name="y_wind",
+            long_name="northward_wind_at_10m",
+            units="m s-1",
+        )
+    )
+
+    return wind_cubelist
+
+
+@pytest.fixture()
+def wind_cubelist_lfric():
+    """LFRic wind cube list."""
+    wind_cubelist = CubeList()
+
+    wind_cubelist.append(
+        Cube(3.5, var_name="wspd10m", long_name="wind_speed_at_10m", units="m s-1")
+    )
+    wind_cubelist.append(
+        Cube(3.5, var_name="u10m", long_name="eastward_wind_at_10m", units="m s-1")
+    )
+    wind_cubelist.append(
+        Cube(3.5, var_name="v10m", long_name="northward_wind_at_10m", units="m s-1")
+    )
+
+    return wind_cubelist
+
+
+@pytest.fixture()
+def wind_cubelist_observed():
+    """Observed wind cube list."""
+    wind_cubelist = CubeList()
+
+    wind_cubelist.append(
+        Cube(3.5, long_name="observed_wind_speed_at_10m", units="m s-1")
+    )
+
+    return wind_cubelist

tests/operators/test_read.py

@@ -159,7 +159,6 @@ def test_read_cube():
         constraint=constraints.generate_cell_methods_constraint([]),
     )
     assert isinstance(cube, iris.cube.Cube)
-
     # Global [latitude, longitude] example file
     cube = read.read_cube(
         "tests/test_data/air_temperature_global.nc",
@@ -1307,15 +1306,14 @@ def test_check_combine_point_observations_multiple_obs(cube):
 
 def test_compute_winds(vector_cubes, tmp_working_dir):
     """Ensure _compute_winds calculates wind_speed from component inputs."""
+    constraint = constraints.generate_var_constraint("wind_speed_at_10m")
     assert len(vector_cubes) == 2
     vector_cubes[0].rename("eastward_wind_at_10m")
     vector_cubes[1].rename("northward_wind_at_10m")
-    output_cubes = read._compute_winds(vector_cubes)
-    assert len(vector_cubes) == 3
+    output_cubes = read._compute_winds(vector_cubes, constraint=constraint)
     assert len(output_cubes) == 1
     assert output_cubes.extract(iris.Constraint("wind_speed_at_10m"))
     assert output_cubes.extract(iris.Constraint("wind_speed_at_10m"))[0].units == "ms-1"
-
     u = vector_cubes[0].data
     v = vector_cubes[1].data
     expected_wind = (u**2 + v**2) ** 0.5
@@ -1333,10 +1331,20 @@ def test_compute_winds_multi_vars(cube, vector_cubes, tmp_working_dir):
     """Ensure _compute_winds extends multiple input variable cubelist."""
     # Add cube of air_temperature to CubeList containing wind components.
     assert len(vector_cubes) == 2
+    constraint = constraints.generate_var_constraint(
+        [
+            "air_temperature",
+            "wind_speed_at_10m",
+            "eastward_wind_at_10m",
+            "northward_wind_at_10m",
+        ]
+    )
     vector_cubes[0].rename("eastward_wind_at_10m")
     vector_cubes[1].rename("northward_wind_at_10m")
+
     output_cubes = read._compute_winds(
-        iris.cube.CubeList([cube, vector_cubes[0], vector_cubes[1]])
+        iris.cube.CubeList([cube, vector_cubes[0], vector_cubes[1]]),
+        constraint=constraint,
     )
     assert len(vector_cubes) == 2
     assert len(output_cubes) == 4
@@ -1351,3 +1359,141 @@ def test_compute_winds_multi_vars(cube, vector_cubes, tmp_working_dir):
     expected_wind = (u**2 + v**2) ** 0.5
     output_wind = output_cubes.extract(iris.Constraint("wind_speed_at_10m"))[0]
     assert np.allclose(output_wind.data, expected_wind, rtol=1e-6, atol=1e-2)
+
+
+def test_compute_wind_no_constraint(wind_cubelist_um):
+    """No constraint given should return the cubelist unchanged."""
+    cubes = read._compute_winds(wind_cubelist_um)
+    assert cubes == wind_cubelist_um
+
+
+def test_wind_um(wind_cubelist_um):
+    """UM cubes filtered to wind speed only."""
+    cubes = wind_cubelist_um.copy()
+    constraint = constraints.generate_var_constraint("wind_speed_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    speed = cubes.extract_cube(iris.Constraint("wind_speed_at_10m"))
+    assert speed.standard_name == "wind_speed"
+
+
+def test_wind_um_eastward(wind_cubelist_um):
+    """UM cubes filtered to eastward wind only."""
+    cubes = wind_cubelist_um.copy()
+    constraint = constraints.generate_var_constraint("eastward_wind_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    assert eastward.standard_name == "x_wind"
+
+
+def test_wind_um_northward(wind_cubelist_um):
+    """UM cubes filtered to northward wind only."""
+    cubes = wind_cubelist_um.copy()
+    constraint = constraints.generate_var_constraint("northward_wind_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert northward.standard_name == "y_wind"
+
+
+def test_wind_um_northward_and_eastward(wind_cubelist_um):
+    """UM cubes filtered to both wind components."""
+    cubes = wind_cubelist_um.copy()
+    constraint = constraints.generate_var_constraint(
+        ["eastward_wind_at_10m", "northward_wind_at_10m"]
+    )
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 2
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert eastward.standard_name == "x_wind"
+    assert northward.standard_name == "y_wind"
+
+
+def test_wind_um_all(wind_cubelist_um):
+    """UM cubes filtered to wind speed and both wind components."""
+    cubes = wind_cubelist_um.copy()
+    constraint = constraints.generate_var_constraint(
+        ["wind_speed_at_10m", "eastward_wind_at_10m", "northward_wind_at_10m"]
+    )
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 3
+    speed = cubes.extract_cube(iris.Constraint("wind_speed_at_10m"))
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert speed.standard_name == "wind_speed"
+    assert eastward.standard_name == "x_wind"
+    assert northward.standard_name == "y_wind"
+
+
+def test_wind_lfric(wind_cubelist_lfric):
+    """LFRic cubes filtered to wind speed only."""
+    cubes = wind_cubelist_lfric.copy()
+    constraint = constraints.generate_var_constraint("wind_speed_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    speed = cubes.extract_cube(iris.Constraint("wind_speed_at_10m"))
+    assert speed.var_name == "wspd10m"
+
+
+def test_wind_lfric_eastward(wind_cubelist_lfric):
+    """LFRic cubes filtered to eastward wind only."""
+    cubes = wind_cubelist_lfric.copy()
+    constraint = constraints.generate_var_constraint("eastward_wind_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    assert eastward.var_name == "u10m"
+
+
+def test_wind_lfric_northward(wind_cubelist_lfric):
+    """LFRic cubes filtered to northward wind only."""
+    cubes = wind_cubelist_lfric.copy()
+    constraint = constraints.generate_var_constraint("northward_wind_at_10m")
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert northward.var_name == "v10m"
+
+
+def test_wind_lfric_northward_and_eastward(wind_cubelist_lfric):
+    """LFRic cubes filtered to both wind components."""
+    cubes = wind_cubelist_lfric.copy()
+    constraint = constraints.generate_var_constraint(
+        ["eastward_wind_at_10m", "northward_wind_at_10m"]
+    )
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 2
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert eastward.var_name == "u10m"
+    assert northward.var_name == "v10m"
+
+
+def test_wind_lfric_all(wind_cubelist_lfric):
+    """LFRic cubes filtered to speed and both wind components."""
+    cubes = wind_cubelist_lfric.copy()
+    constraint = constraints.generate_var_constraint(
+        ["wind_speed_at_10m", "eastward_wind_at_10m", "northward_wind_at_10m"]
+    )
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 3
+    speed = cubes.extract_cube(iris.Constraint("wind_speed_at_10m"))
+    eastward = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
+    northward = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
+    assert speed.var_name == "wspd10m"
+    assert eastward.var_name == "u10m"
+    assert northward.var_name == "v10m"
+
+
+def test_wind_observed(wind_cubelist_observed):
+    """Observed cubes filtered to speed."""
+    cubes = wind_cubelist_observed.copy()
+    constraint = constraints.generate_var_constraint(
+        ["observed_wind_speed_at_10m", "wind_speed_at_10m"]
+    )
+    cubes = read._compute_winds(cubes, constraint=constraint)
+    assert len(cubes) == 1
+    observed_speed = cubes.extract_cube(iris.Constraint("observed_wind_speed_at_10m"))
+    assert observed_speed