Coverage for src/CSET/operators/constraints.py: 93%
113 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-13 14:57 +0000
1# © Crown copyright, Met Office (2022-2025) and CSET contributors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Operators to generate constraints to filter with."""
17import numbers
18import re
19from collections.abc import Iterable
20from datetime import timedelta
22import iris
23import iris.coords
24import iris.cube
26import CSET.operators._utils as operator_utils
27from CSET._common import iter_maybe
29# STASH code pattern: mXXsXXiXXX where X is a digit
30_STASH_RE = re.compile(r"^m\d{2}s\d{2}i\d{3}$")
33def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint:
34 """Generate constraint from STASH code.
36 Operator that takes a stash string, and uses iris to generate a constraint
37 to be passed into the read operator to minimize the CubeList the read
38 operator loads and speed up loading.
40 Arguments
41 ---------
42 stash: str
43 stash code to build iris constraint, such as "m01s03i236"
45 Returns
46 -------
47 stash_constraint: iris.AttributeConstraint
48 """
49 # At a later stage str list an option to combine constraints. Arguments
50 # could be a list of stash codes that combined build the constraint.
51 stash_constraint = iris.AttributeConstraint(STASH=stash)
52 return stash_constraint
55def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
56 """Generate constraint from variable name or STASH code.
58 Operator that takes a CF compliant variable name string or list of names, and generates an
59 iris constraint to be passed into the read or filter operator. Can also be
60 passed a STASH code to generate a STASH constraint.
62 Arguments
63 ---------
64 varname: str | list[str]
65 CF compliant name(s) of variable, or a UM STASH code such as "m01s03i236".
67 Returns
68 -------
69 varname_constraint: iris.Constraint
70 If a single UM STASHcode is requested, varname constraint is by STASHcode
71 If a single variable name is requested, constraint by varname
72 If multiple variable names are requested, constrain by list of variables.
73 """
74 # Case 1: UM STASHcode input
76 if isinstance(varname, str) and _STASH_RE.match(varname):
77 return iris.AttributeConstraint(STASH=varname)
79 # Ensure access to variable vector components for computed fields
80 if "wind_speed_at_10m" in iter_maybe(varname):
81 if isinstance(varname, str): 81 ↛ 83line 81 didn't jump to line 83 because the condition on line 81 was always true
82 varname = [varname]
83 varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"])
84 varname.extend(["u_wind_at_10m", "v_wind_at_10m"])
86 # Case 2: Multiple varnames
87 if isinstance(varname, (list, tuple)):
88 varname_constraint = iris.Constraint(
89 cube_func=lambda cube: (
90 cube.long_name in varname
91 or cube.standard_name in varname
92 or cube.var_name in varname
93 )
94 )
96 else:
97 varname_constraint = iris.Constraint(name=varname)
99 return varname_constraint
102def generate_level_constraint(
103 coordinate: str, levels: int | list[int] | str, **kwargs
104) -> iris.Constraint:
105 """Generate constraint for particular levels on the specified coordinate.
107 Operator that generates a constraint to constrain to specific model or
108 pressure levels. If no levels are specified then any cube with the specified
109 coordinate is rejected.
111 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"``
112 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic.
114 Arguments
115 ---------
116 coordinate: str
117 Level coordinate name about which to constraint.
118 levels: int | list[int] | str
119 CF compliant level points, ``"*"`` for retrieving all levels, or
120 ``[]`` for no levels.
122 Returns
123 -------
124 constraint: iris.Constraint
126 Notes
127 -----
128 Due to the specification of ``coordinate`` as an argument any iterable
129 coordinate can be stratified with this function. Therefore,
130 ``"realization"`` is a valid option. Subsequently, ``levels`` specifies the
131 ensemble members, or group of ensemble members you wish to constrain your
132 results over.
133 """
134 # If asterisks, then return all levels for given coordinate.
135 if levels == "*":
136 return iris.Constraint(**{coordinate: lambda cell: True})
137 else:
138 # Ensure is iterable.
139 if not isinstance(levels, Iterable):
140 levels = [levels]
142 # When no levels specified reject cube with level coordinate.
143 if len(levels) == 0:
145 def no_levels(cube):
146 # Reject cubes for which coordinate exists.
147 return not cube.coords(coordinate)
149 return iris.Constraint(cube_func=no_levels)
151 # Filter the coordinate to the desired levels.
152 # Dictionary unpacking is used to provide programmatic keyword arguments.
153 return iris.Constraint(**{coordinate: levels})
156def generate_remove_single_level_constraint(
157 coord: str, level: int = 0, **kwargs
158) -> iris.Constraint:
159 """
160 Generate a constraint to remove a single model level number.
162 Operator that returns a constraint to remove the given level. By
163 default the first level is removed (assumed to be
164 level zero). However, any level can be removed.
166 Arguments
167 ---------
168 coord: str
169 The coordinate for which the level is to be removed.
170 level: int
171 Default is 0. The model level number to remove.
173 Returns
174 -------
175 iris.Constraint
177 Notes
178 -----
179 This operator is primarily used to ensure the levels are consistent
180 as some level sets (e.g. specific humidity) will be on the same level set
181 but have a different number of levels (e.g 71 instead of expected 70).
182 """
183 return iris.Constraint(**{coord: lambda m: m.point != level})
186def generate_cell_methods_constraint(
187 cell_methods: list,
188 varname: str | None = None,
189 coord: iris.coords.Coord | None = None,
190 interval: str | None = None,
191 comment: str | None = None,
192 **kwargs,
193) -> iris.Constraint:
194 """Generate constraint from cell methods.
196 Operator that takes a list of cell methods and generates a constraint from
197 that. Use [] to specify non-aggregated data.
199 Arguments
200 ---------
201 cell_methods: list
202 cube.cell_methods for filtering.
203 varname: str, optional
204 CF compliant name of variable.
205 coord: iris.coords.Coord, optional
206 iris.coords.Coord to which the cell method is applied to.
207 interval: str, optional
208 interval over which the cell method is applied to (e.g. 1 hour).
209 comment: str, optional
210 any comments in Cube meta data associated with the cell method.
212 Returns
213 -------
214 cell_method_constraint: iris.Constraint
215 """
216 if len(cell_methods) == 0:
218 def check_no_aggregation(cube: iris.cube.Cube) -> bool:
219 """Check that any cell methods are "point", meaning no aggregation."""
220 return {cm.method for cm in cube.cell_methods} <= {"point"}
222 def check_cell_sum(cube: iris.cube.Cube) -> bool:
223 """Check that any cell methods are "sum"."""
224 return {cm.method for cm in cube.cell_methods} == {"sum"}
226 def check_cell_mean(cube: iris.cube.Cube) -> bool:
227 """Check that any cell methods are "mean"."""
228 return {cm.method for cm in cube.cell_methods} == {"mean"}
230 if varname:
231 # Require number_of_lightning_flashes to be "sum" cell_method input.
232 # Require surface_microphyisical_rainfall_amount and surface_microphysical_snowfall_amount to be "sum" cell_method inputs.
233 if ("lightning" in varname) or (
234 "surface_microphysical" in varname and "amount" in varname
235 ):
236 cell_methods_constraint = iris.Constraint(cube_func=check_cell_sum)
237 return cell_methods_constraint
238 # Require climatological ancillary as time-average mean.
239 if ("albedo" in varname) or ( 239 ↛ 246line 239 didn't jump to line 246 because the condition on line 239 was always true
240 "ocean" in varname and "chlorophyll" in varname
241 ):
242 cell_methods_constraint = iris.Constraint(cube_func=check_cell_mean)
243 return cell_methods_constraint
245 # If no variable name set, assume require instantaneous cube.
246 cell_methods_constraint = iris.Constraint(cube_func=check_no_aggregation)
248 else:
249 # If cell_method constraint set in recipe, check for required input.
250 def check_cell_methods(cube: iris.cube.Cube) -> bool:
251 return all(
252 iris.coords.CellMethod(
253 method=cm, coords=coord, intervals=interval, comments=comment
254 )
255 in cube.cell_methods
256 for cm in cell_methods
257 )
259 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)
261 return cell_methods_constraint
264def generate_time_constraint(
265 time_start: str, time_end: str | None = None, **kwargs
266) -> iris.Constraint:
267 """Generate constraint between times.
269 Operator that takes one or two ISO 8601 date strings, and returns a
270 constraint that selects values between those dates (inclusive).
272 Arguments
273 ---------
274 time_start: str | datetime.datetime | cftime.datetime
275 ISO date for lower bound
277 time_end: str | datetime.datetime | cftime.datetime
278 ISO date for upper bound. If omitted it defaults to the same as
279 time_start
281 Returns
282 -------
283 time_constraint: iris.Constraint
284 """
285 if isinstance(time_start, str):
286 pdt_start, offset_start = operator_utils.pdt_fromisoformat(time_start)
287 else:
288 pdt_start, offset_start = time_start, timedelta(0)
290 if time_end is None:
291 pdt_end, offset_end = time_start, offset_start
292 elif isinstance(time_end, str):
293 pdt_end, offset_end = operator_utils.pdt_fromisoformat(time_end)
294 print(pdt_end)
295 print(offset_end)
296 else:
297 pdt_end, offset_end = time_end, timedelta(0)
299 if offset_start is None:
300 offset_start = timedelta(0)
301 if offset_end is None:
302 offset_end = timedelta(0)
304 time_constraint = iris.Constraint(
305 time=lambda t: (
306 (pdt_start <= (t.point - offset_start))
307 and ((t.point - offset_end) <= pdt_end)
308 )
309 )
311 return time_constraint
314def generate_area_constraint(
315 lat_start: float | None,
316 lat_end: float | None,
317 lon_start: float | None,
318 lon_end: float | None,
319 **kwargs,
320) -> iris.Constraint:
321 """Generate an area constraint between latitude/longitude limits.
323 Operator that takes a set of latitude and longitude limits and returns a
324 constraint that selects grid values only inside that area. Works with the
325 data's native grid so is defined within the rotated pole CRS.
327 Alternatively, all arguments may be None to indicate the area should not be
328 constrained. This is useful to allow making subsetting an optional step in a
329 processing pipeline.
331 Arguments
332 ---------
333 lat_start: float | None
334 Latitude value for lower bound
335 lat_end: float | None
336 Latitude value for top bound
337 lon_start: float | None
338 Longitude value for left bound
339 lon_end: float | None
340 Longitude value for right bound
342 Returns
343 -------
344 area_constraint: iris.Constraint
345 """
346 # Check all arguments are defined, or all are None.
347 if not (
348 all(
349 (
350 isinstance(lat_start, numbers.Real),
351 isinstance(lat_end, numbers.Real),
352 isinstance(lon_start, numbers.Real),
353 isinstance(lon_end, numbers.Real),
354 )
355 )
356 or all((lat_start is None, lat_end is None, lon_start is None, lon_end is None))
357 ):
358 raise TypeError("Bounds must real numbers, or all None.")
360 # Don't constrain area if all arguments are None.
361 if lat_start is None: # Only need to check once, as they will be the same.
362 # An empty constraint allows everything.
363 return iris.Constraint()
365 # Handle bounds crossing the date line.
366 if lon_end < lon_start: 366 ↛ 367line 366 didn't jump to line 367 because the condition on line 366 was never true
367 lon_end = lon_end + 360
369 def bound_lat(cell: iris.coords.Cell) -> bool:
370 return lat_start < cell < lat_end
372 def bound_lon(cell: iris.coords.Cell) -> bool:
373 # Adjust cell values to handle crossing the date line.
374 if cell < lon_start:
375 cell = cell + 360
376 return lon_start < cell < lon_end
378 area_constraint = iris.Constraint(
379 coord_values={"grid_latitude": bound_lat, "grid_longitude": bound_lon}
380 )
381 return area_constraint
384def generate_remove_single_ensemble_member_constraint(
385 ensemble_member: int = 0, **kwargs
386) -> iris.Constraint:
387 """
388 Generate a constraint to remove a single ensemble member.
390 Operator that returns a constraint to remove the given ensemble member. By
391 default the ensemble member removed is the control member (assumed to have
392 a realization of zero). However, any ensemble member can be removed, thus
393 allowing a non-zero control member to be removed if the control is a
394 different member.
396 Arguments
397 ---------
398 ensemble_member: int
399 Default is 0. The ensemble member realization to remove.
401 Returns
402 -------
403 iris.Constraint
405 Notes
406 -----
407 This operator is primarily used to remove the control member to allow
408 ensemble metrics to be calculated without the control member. For
409 example, the ensemble mean is not normally calculated including the
410 control member. It is particularly useful to remove the control member
411 when it is not an equally-likely member of the ensemble.
412 """
413 return iris.Constraint(realization=lambda m: m.point != ensemble_member)
416def generate_realization_constraint(
417 ensemble_members: int | list[int], **kwargs
418) -> iris.Constraint:
419 """
420 Generate a constraint to subset ensemble members.
422 Operator that is given a list of ensemble members and returns a constraint
423 to select those ensemble members. This operator is particularly useful for
424 subsetting ensembles.
426 Arguments
427 ---------
428 ensemble_members: int | list[int]
429 The ensemble members to be subsetted over.
431 Returns
432 -------
433 iris.Constraint
434 """
435 # Ensure ensemble_members is iterable.
436 ensemble_members = iter_maybe(ensemble_members)
437 return iris.Constraint(realization=ensemble_members)
440def generate_hour_constraint(
441 hour_start: int,
442 hour_end: int | None = None,
443 **kwargs,
444) -> iris.Constraint:
445 """Generate an hour constraint between hour of day limits.
447 Operator that takes a set of hour of day limits and returns a constraint that
448 selects only hours within that time frame regardless of day.
450 Alternatively, the result can be constrained to a single hour by just entering
451 a starting hour.
453 Should any sub-hourly data be given these will have the same hour coordinate
454 (e.g., 12:00 and 12:05 both have an hour coordinate of 12) all
455 times will be selected with this constraint.
457 Arguments
458 ---------
459 hour_start: int
460 The hour of day for the lower bound, within 0 to 23.
461 hour_end: int | None
462 The hour of day for the upper bound, within 0 to 23. Alternatively,
463 set to None if only one hour required.
465 Returns
466 -------
467 hour_constraint: iris.Constraint
469 Raises
470 ------
471 ValueError
472 If the provided arguments are outside of the range 0 to 23.
473 """
474 if hour_end is None:
475 hour_end = hour_start
477 if (hour_start < 0) or (hour_start > 23) or (hour_end < 0) or (hour_end > 23):
478 raise ValueError("Hours must be between 0 and 23 inclusive.")
480 hour_constraint = iris.Constraint(hour=lambda h: hour_start <= h.point <= hour_end)
481 return hour_constraint
484def combine_constraints(
485 constraint: iris.Constraint = None, **kwargs
486) -> iris.Constraint:
487 """
488 Operator that combines multiple constraints into one.
490 Arguments
491 ---------
492 constraint: iris.Constraint
493 First constraint to combine.
494 additional_constraint_1: iris.Constraint
495 Second constraint to combine. This must be a named argument.
496 additional_constraint_2: iris.Constraint
497 There can be any number of additional constraint, they just need unique
498 names.
499 ...
501 Returns
502 -------
503 combined_constraint: iris.Constraint
505 Raises
506 ------
507 TypeError
508 If the provided arguments are not constraints.
509 """
510 # If the first argument is not a constraint, it is ignored. This handles the
511 # automatic passing of the previous step's output.
512 if isinstance(constraint, iris.Constraint):
513 combined_constraint = constraint
514 else:
515 combined_constraint = iris.Constraint()
517 for constr in kwargs.values():
518 combined_constraint = combined_constraint & constr
519 return combined_constraint
522def generate_attribute_constraint(
523 attribute: str, value: str | None = None, **kwargs
524) -> iris.AttributeConstraint:
525 """Generate constraint on cube attributes.
527 Constrains based on the presence of an attribute, and that attribute having
528 a particular value.
530 Arguments
531 ---------
532 attribute: str
533 Attribute to constraint on.
535 value: str
536 Attribute value to constrain on. If omitted the constraint merely checks
537 for the presence of an attribute.
539 Returns
540 -------
541 attribute_constraint: iris.Constraint
542 """
543 if value is None:
544 attribute_constraint = iris.Constraint(
545 cube_func=lambda cube: attribute in cube.attributes
546 )
547 else:
548 attribute_constraint = iris.AttributeConstraint(**{attribute: value})
549 return attribute_constraint