diff --git a/src/CSET/operators/constraints.py b/src/CSET/operators/constraints.py index 694bdc5f0..bb0281c55 100644 --- a/src/CSET/operators/constraints.py +++ b/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 diff --git a/src/CSET/operators/plot.py b/src/CSET/operators/plot.py index 5ec8d7315..ff2ca6198 100644 --- a/src/CSET/operators/plot.py +++ b/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) diff --git a/src/CSET/operators/read.py b/src/CSET/operators/read.py index c9747a6f5..ddd0b888f 100644 --- a/src/CSET/operators/read.py +++ b/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" diff --git a/tests/operators/test_constraints.py b/tests/operators/test_constraints.py index c5112ae00..0e12bae49 100644 --- a/tests/operators/test_constraints.py +++ b/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=." + ) + 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")