diff --git a/src/CSET/operators/constraints.py b/src/CSET/operators/constraints.py index bb0281c55..950c7d66a 100644 --- a/src/CSET/operators/constraints.py +++ b/src/CSET/operators/constraints.py @@ -26,6 +26,9 @@ import CSET.operators._utils as operator_utils from CSET._common import iter_maybe +# STASH code pattern: mXXsXXiXXX where X is a digit +_STASH_RE = re.compile(r"^m\d{2}s\d{2}i\d{3}$") + def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint: """Generate constraint from STASH code. @@ -52,33 +55,44 @@ def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint: def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint: """Generate constraint from variable name or STASH code. - Operator that takes a CF compliant variable name string, and generates an + Operator that takes a CF compliant variable name string or list of names, and generates an iris constraint to be passed into the read or filter operator. Can also be passed a STASH code to generate a STASH constraint. Arguments --------- - varname: str - CF compliant name of variable, or a UM STASH code such as "m01s03i236". + varname: str | list[str] + CF compliant name(s) of variable, or a UM STASH code such as "m01s03i236". Returns ------- varname_constraint: iris.Constraint + If a single UM STASHcode is requested, varname constraint is by STASHcode + If a single variable name is requested, constraint by varname + If multiple variable names are requested, constrain by list of variables. """ - if re.match(r"m[0-9]{2}s[0-9]{2}i[0-9]{3}$", varname): - varname_constraint = iris.AttributeConstraint(STASH=varname) - else: - varname_constraint = iris.Constraint(name=varname) + # Case 1: UM STASHcode input + if isinstance(varname, str) and _STASH_RE.match(varname): + return iris.AttributeConstraint(STASH=varname) # Ensure access to variable vector components for computed fields - if varname == "wind_speed_at_10m": + if "wind_speed_at_10m" in iter_maybe(varname): + varname = [varname] + varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"]) + + # Case 2: Multiple varnames + if isinstance(varname, (list, tuple)): varname_constraint = iris.Constraint( cube_func=lambda cube: ( - cube.long_name - in ["wind_at_10m", "eastward_wind_at_10m", "northward_wind_at_10m"] + cube.long_name in varname + or cube.standard_name in varname + or cube.var_name in varname ) ) + else: + varname_constraint = iris.Constraint(name=varname) + return varname_constraint diff --git a/tests/operators/test_constraints.py b/tests/operators/test_constraints.py index 0e12bae49..43a8be117 100644 --- a/tests/operators/test_constraints.py +++ b/tests/operators/test_constraints.py @@ -16,6 +16,8 @@ from datetime import datetime +import iris +import numpy as np import pytest from CSET.operators import constraints @@ -288,6 +290,33 @@ def test_generate_attribute_constraint_with_value(): assert expected_attr_constraint in repr(attr_constraint) +def test_generate_var_constraint_multiple_names(): + """Test constraint works for multiple variable names.""" + # Create two cubes with different names + cube1 = iris.cube.Cube(np.arange(5), long_name="temperature_long") + cube1.var_name = "var_temperature" + cube2 = iris.cube.Cube(np.arange(5), standard_name="wind_speed") + # Third cube that should NOT match + cube3 = iris.cube.Cube(np.arange(5), long_name="surface_pressure") + # Generate constraint with multiple names + constraint = constraints.generate_var_constraint(["var_temperature", "wind_speed"]) + # Apply constraint + cubes = iris.cube.CubeList([cube1, cube2, cube3]) + result = cubes.extract(constraint) + # Check correct cubes are selected + result_names = [c.name() for c in result] + + assert cube1 in result + assert cube2 in result + assert cube3 not in result + + assert "temperature_long" in result_names + assert "wind_speed" in result_names + assert "surface_pressure" not in result_names + # Should only return 2 cubes + assert len(result) == 2 + + def test_generate_remove_single_level_constraint(): """Tests constraint to remove default model_level_number of zero.""" remove_level_constraint = constraints.generate_remove_single_level_constraint(