Implement calculate wind_speed from components

src/CSET/operators/constraints.py

@@ -69,6 +69,16 @@ def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
         varname_constraint = iris.AttributeConstraint(STASH=varname)
     else:
         varname_constraint = iris.Constraint(name=varname)
+
+    # Ensure access to variable vector components for computed fields
+    if varname == "wind_speed_at_10m":
+        varname_constraint = iris.Constraint(
+            cube_func=lambda cube: (
+                cube.long_name
+                in ["wind_at_10m", "eastward_wind_at_10m", "northward_wind_at_10m"]
+            )
+        )
+
     return varname_constraint
 
 

src/CSET/operators/plot.py

@@ -1151,10 +1151,13 @@ def _plot_and_save_vector_plot(
         Plot title.
     """
     fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
-
     # Create a cube containing the magnitude of the vector field.
     cube_vec_mag = (cube_u**2 + cube_v**2) ** 0.5
-    cube_vec_mag.rename(f"{cube_u.name()}_{cube_v.name()}_magnitude")
+    cube_vec_mag.rename(f"{cube_u.long_name}_{cube_v.long_name}_magnitude")
+    if "eastward_wind" in cube_u.long_name and "northward_wind" in cube_v.long_name:
+        cube_vec_mag.rename(
+            "wind_speed" + cube_u.long_name.replace("eastward_wind", "")
+        )
 
     # Specify the color bar
     cmap, levels, norm = colorbar_map_levels(cube_vec_mag)
@@ -2496,7 +2499,7 @@ def vector_plot(
             cube_v_slice,
             filename=plot_filename,
             title=plot_title,
-            method="contourf",
+            method="pcolormesh",
         )
         plot_index.append(plot_filename)
 

src/CSET/operators/read.py

@@ -232,7 +232,7 @@ def _load_model(
     logging.debug("Constraint: %s", constraint)
     cubes = iris.load(input_files, constraint, callback=_loading_callback)
     # Make the UM's winds consistent with LFRic.
-    _fix_um_winds(cubes)
+    cubes = _fix_um_winds(cubes)
 
     # Add model_name attribute to each cube to make it available at any further
     # step without needing to pass it as function parameter.
@@ -896,19 +896,28 @@ def _fix_um_winds(cubes: iris.cube.CubeList):
     speed_constr = iris.AttributeConstraint(STASH="m01s03i227")
     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):
                 _add_wind_speed_um(cubes)
             # Convert winds in the UM to be relative to true east and true north.
             _convert_wind_true_dirn_um(cubes)
+            # Return only wind_speed cube
+            if wind_only:
+                cubes = cubes.extract(speed_constr)
     except (KeyError, AttributeError):
         pass
 
+    return cubes
+
 
 def _add_wind_speed_um(cubes: iris.cube.CubeList):
     """Add windspeeds to cubes from the UM."""
     wspd10 = (
-        cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i225"))[0] ** 2
-        + cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i226"))[0] ** 2
+        cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i225")) ** 2
+        + cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i226")) ** 2
     ) ** 0.5
     wspd10.attributes["STASH"] = "m01s03i227"
     wspd10.standard_name = "wind_speed"

tests/operators/test_constraints.py

@@ -35,6 +35,15 @@ def test_generate_var_constraint():
     assert repr(var_constraint) == expected_var_constraint
 
 
+def test_generate_var_constraint_wind():
+    """Generate iris cube constraint for wind speed."""
+    var_constraint = constraints.generate_var_constraint("wind_speed_at_10m")
+    expected_var_constraint = (
+        "Constraint(cube_func=<function generate_var_constraint.<locals>.<lambda>"
+    )
+    assert expected_var_constraint in repr(var_constraint)
+
+
 def test_generate_var_constraint_stash():
     """Generate iris cube constraint for UM STASH code with var constraint."""
     var_constraint = constraints.generate_var_constraint("m01s03i236")