diff --git a/src/CSET/operators/constraints.py b/src/CSET/operators/constraints.py index 1900394e9..f03ff352e 100644 --- a/src/CSET/operators/constraints.py +++ b/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 diff --git a/src/CSET/operators/read.py b/src/CSET/operators/read.py index 9a3c6e435..b8dda708e 100644 --- a/src/CSET/operators/read.py +++ b/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): diff --git a/tests/conftest.py b/tests/conftest.py index bbd604b56..89f9d07f0 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/operators/test_read.py b/tests/operators/test_read.py index c59a39580..85fe95ef2 100644 --- a/tests/operators/test_read.py +++ b/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