Coverage for src/CSET/operators/constraints.py: 93%
115 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 13: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 | list[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 varname_copy = tuple(iter_maybe(varname))
82 if "wind_speed_at_10m" in varname_copy:
83 if isinstance(varname, str):
84 varname = [varname]
85 varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"])
86 varname.extend(["u_wind_at_10m", "v_wind_at_10m"])
87 # Case 2: Multiple varnames
88 if isinstance(varname, (list, tuple)):
89 varname_constraint = iris.Constraint(
90 cube_func=lambda cube: (
91 cube.long_name in varname
92 or cube.standard_name in varname
93 or cube.var_name in varname
94 )
95 )
97 else:
98 varname_constraint = iris.Constraint(name=varname)
99 varname_constraint.varname = varname_copy
100 return varname_constraint
103def generate_level_constraint(
104 coordinate: str, levels: int | list[int] | str, **kwargs
105) -> iris.Constraint:
106 """Generate constraint for particular levels on the specified coordinate.
108 Operator that generates a constraint to constrain to specific model or
109 pressure levels. If no levels are specified then any cube with the specified
110 coordinate is rejected.
112 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"``
113 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic.
115 Arguments
116 ---------
117 coordinate: str
118 Level coordinate name about which to constraint.
119 levels: int | list[int] | str
120 CF compliant level points, ``"*"`` for retrieving all levels, or
121 ``[]`` for no levels.
123 Returns
124 -------
125 constraint: iris.Constraint
127 Notes
128 -----
129 Due to the specification of ``coordinate`` as an argument any iterable
130 coordinate can be stratified with this function. Therefore,
131 ``"realization"`` is a valid option. Subsequently, ``levels`` specifies the
132 ensemble members, or group of ensemble members you wish to constrain your
133 results over.
134 """
135 # If asterisks, then return all levels for given coordinate.
136 if levels == "*":
137 return iris.Constraint(**{coordinate: lambda cell: True})
138 else:
139 # Ensure is iterable.
140 if not isinstance(levels, Iterable):
141 levels = [levels]
143 # When no levels specified reject cube with level coordinate.
144 if len(levels) == 0:
146 def no_levels(cube):
147 # Reject cubes for which coordinate exists.
148 return not cube.coords(coordinate)
150 return iris.Constraint(cube_func=no_levels)
152 # Filter the coordinate to the desired levels.
153 # Dictionary unpacking is used to provide programmatic keyword arguments.
154 return iris.Constraint(**{coordinate: levels})
157def generate_remove_single_level_constraint(
158 coord: str, level: int = 0, **kwargs
159) -> iris.Constraint:
160 """
161 Generate a constraint to remove a single model level number.
163 Operator that returns a constraint to remove the given level. By
164 default the first level is removed (assumed to be
165 level zero). However, any level can be removed.
167 Arguments
168 ---------
169 coord: str
170 The coordinate for which the level is to be removed.
171 level: int
172 Default is 0. The model level number to remove.
174 Returns
175 -------
176 iris.Constraint
178 Notes
179 -----
180 This operator is primarily used to ensure the levels are consistent
181 as some level sets (e.g. specific humidity) will be on the same level set
182 but have a different number of levels (e.g 71 instead of expected 70).
183 """
184 return iris.Constraint(**{coord: lambda m: m.point != level})
187def generate_cell_methods_constraint(
188 cell_methods: list,
189 varname: str | None = None,
190 coord: iris.coords.Coord | None = None,
191 interval: str | None = None,
192 comment: str | None = None,
193 **kwargs,
194) -> iris.Constraint:
195 """Generate constraint from cell methods.
197 Operator that takes a list of cell methods and generates a constraint from
198 that. Use [] to specify non-aggregated data.
200 Arguments
201 ---------
202 cell_methods: list
203 cube.cell_methods for filtering.
204 varname: str, optional
205 CF compliant name of variable.
206 coord: iris.coords.Coord, optional
207 iris.coords.Coord to which the cell method is applied to.
208 interval: str, optional
209 interval over which the cell method is applied to (e.g. 1 hour).
210 comment: str, optional
211 any comments in Cube meta data associated with the cell method.
213 Returns
214 -------
215 cell_method_constraint: iris.Constraint
216 """
217 if len(cell_methods) == 0:
219 def check_no_aggregation(cube: iris.cube.Cube) -> bool:
220 """Check that any cell methods are "point", meaning no aggregation."""
221 return {cm.method for cm in cube.cell_methods} <= {"point"}
223 def check_cell_sum(cube: iris.cube.Cube) -> bool:
224 """Check that any cell methods are "sum"."""
225 return {cm.method for cm in cube.cell_methods} == {"sum"}
227 def check_cell_mean(cube: iris.cube.Cube) -> bool:
228 """Check that any cell methods are "mean"."""
229 return {cm.method for cm in cube.cell_methods} == {"mean"}
231 if varname:
232 # Require number_of_lightning_flashes to be "sum" cell_method input.
233 # Require surface_microphyisical_rainfall_amount and surface_microphysical_snowfall_amount to be "sum" cell_method inputs.
234 if ("lightning" in varname) or (
235 "surface_microphysical" in varname and "amount" in varname
236 ):
237 cell_methods_constraint = iris.Constraint(cube_func=check_cell_sum)
238 return cell_methods_constraint
239 # Require climatological ancillary as time-average mean.
240 if ("albedo" in varname) or ( 240 ↛ 247line 240 didn't jump to line 247 because the condition on line 240 was always true
241 "ocean" in varname and "chlorophyll" in varname
242 ):
243 cell_methods_constraint = iris.Constraint(cube_func=check_cell_mean)
244 return cell_methods_constraint
246 # If no variable name set, assume require instantaneous cube.
247 cell_methods_constraint = iris.Constraint(cube_func=check_no_aggregation)
249 else:
250 # If cell_method constraint set in recipe, check for required input.
251 def check_cell_methods(cube: iris.cube.Cube) -> bool:
252 return all(
253 iris.coords.CellMethod(
254 method=cm, coords=coord, intervals=interval, comments=comment
255 )
256 in cube.cell_methods
257 for cm in cell_methods
258 )
260 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)
262 return cell_methods_constraint
265def generate_time_constraint(
266 time_start: str, time_end: str | None = None, **kwargs
267) -> iris.Constraint:
268 """Generate constraint between times.
270 Operator that takes one or two ISO 8601 date strings, and returns a
271 constraint that selects values between those dates (inclusive).
273 Arguments
274 ---------
275 time_start: str | datetime.datetime | cftime.datetime
276 ISO date for lower bound
278 time_end: str | datetime.datetime | cftime.datetime
279 ISO date for upper bound. If omitted it defaults to the same as
280 time_start
282 Returns
283 -------
284 time_constraint: iris.Constraint
285 """
286 if isinstance(time_start, str):
287 pdt_start, offset_start = operator_utils.pdt_fromisoformat(time_start)
288 else:
289 pdt_start, offset_start = time_start, timedelta(0)
291 if time_end is None:
292 pdt_end, offset_end = time_start, offset_start
293 elif isinstance(time_end, str):
294 pdt_end, offset_end = operator_utils.pdt_fromisoformat(time_end)
295 print(pdt_end)
296 print(offset_end)
297 else:
298 pdt_end, offset_end = time_end, timedelta(0)
300 if offset_start is None:
301 offset_start = timedelta(0)
302 if offset_end is None:
303 offset_end = timedelta(0)
305 time_constraint = iris.Constraint(
306 time=lambda t: (
307 (pdt_start <= (t.point - offset_start))
308 and ((t.point - offset_end) <= pdt_end)
309 )
310 )
312 return time_constraint
315def generate_area_constraint(
316 lat_start: float | None,
317 lat_end: float | None,
318 lon_start: float | None,
319 lon_end: float | None,
320 **kwargs,
321) -> iris.Constraint:
322 """Generate an area constraint between latitude/longitude limits.
324 Operator that takes a set of latitude and longitude limits and returns a
325 constraint that selects grid values only inside that area. Works with the
326 data's native grid so is defined within the rotated pole CRS.
328 Alternatively, all arguments may be None to indicate the area should not be
329 constrained. This is useful to allow making subsetting an optional step in a
330 processing pipeline.
332 Arguments
333 ---------
334 lat_start: float | None
335 Latitude value for lower bound
336 lat_end: float | None
337 Latitude value for top bound
338 lon_start: float | None
339 Longitude value for left bound
340 lon_end: float | None
341 Longitude value for right bound
343 Returns
344 -------
345 area_constraint: iris.Constraint
346 """
347 # Check all arguments are defined, or all are None.
348 if not (
349 all(
350 (
351 isinstance(lat_start, numbers.Real),
352 isinstance(lat_end, numbers.Real),
353 isinstance(lon_start, numbers.Real),
354 isinstance(lon_end, numbers.Real),
355 )
356 )
357 or all((lat_start is None, lat_end is None, lon_start is None, lon_end is None))
358 ):
359 raise TypeError("Bounds must real numbers, or all None.")
361 # Don't constrain area if all arguments are None.
362 if lat_start is None: # Only need to check once, as they will be the same.
363 # An empty constraint allows everything.
364 return iris.Constraint()
366 # Handle bounds crossing the date line.
367 if lon_end < lon_start: 367 ↛ 368line 367 didn't jump to line 368 because the condition on line 367 was never true
368 lon_end = lon_end + 360
370 def bound_lat(cell: iris.coords.Cell) -> bool:
371 return lat_start < cell < lat_end
373 def bound_lon(cell: iris.coords.Cell) -> bool:
374 # Adjust cell values to handle crossing the date line.
375 if cell < lon_start:
376 cell = cell + 360
377 return lon_start < cell < lon_end
379 area_constraint = iris.Constraint(
380 coord_values={"grid_latitude": bound_lat, "grid_longitude": bound_lon}
381 )
382 return area_constraint
385def generate_remove_single_ensemble_member_constraint(
386 ensemble_member: int = 0, **kwargs
387) -> iris.Constraint:
388 """
389 Generate a constraint to remove a single ensemble member.
391 Operator that returns a constraint to remove the given ensemble member. By
392 default the ensemble member removed is the control member (assumed to have
393 a realization of zero). However, any ensemble member can be removed, thus
394 allowing a non-zero control member to be removed if the control is a
395 different member.
397 Arguments
398 ---------
399 ensemble_member: int
400 Default is 0. The ensemble member realization to remove.
402 Returns
403 -------
404 iris.Constraint
406 Notes
407 -----
408 This operator is primarily used to remove the control member to allow
409 ensemble metrics to be calculated without the control member. For
410 example, the ensemble mean is not normally calculated including the
411 control member. It is particularly useful to remove the control member
412 when it is not an equally-likely member of the ensemble.
413 """
414 return iris.Constraint(realization=lambda m: m.point != ensemble_member)
417def generate_realization_constraint(
418 ensemble_members: int | list[int], **kwargs
419) -> iris.Constraint:
420 """
421 Generate a constraint to subset ensemble members.
423 Operator that is given a list of ensemble members and returns a constraint
424 to select those ensemble members. This operator is particularly useful for
425 subsetting ensembles.
427 Arguments
428 ---------
429 ensemble_members: int | list[int]
430 The ensemble members to be subsetted over.
432 Returns
433 -------
434 iris.Constraint
435 """
436 # Ensure ensemble_members is iterable.
437 ensemble_members = iter_maybe(ensemble_members)
438 return iris.Constraint(realization=ensemble_members)
441def generate_hour_constraint(
442 hour_start: int,
443 hour_end: int | None = None,
444 **kwargs,
445) -> iris.Constraint:
446 """Generate an hour constraint between hour of day limits.
448 Operator that takes a set of hour of day limits and returns a constraint that
449 selects only hours within that time frame regardless of day.
451 Alternatively, the result can be constrained to a single hour by just entering
452 a starting hour.
454 Should any sub-hourly data be given these will have the same hour coordinate
455 (e.g., 12:00 and 12:05 both have an hour coordinate of 12) all
456 times will be selected with this constraint.
458 Arguments
459 ---------
460 hour_start: int
461 The hour of day for the lower bound, within 0 to 23.
462 hour_end: int | None
463 The hour of day for the upper bound, within 0 to 23. Alternatively,
464 set to None if only one hour required.
466 Returns
467 -------
468 hour_constraint: iris.Constraint
470 Raises
471 ------
472 ValueError
473 If the provided arguments are outside of the range 0 to 23.
474 """
475 if hour_end is None:
476 hour_end = hour_start
478 if (hour_start < 0) or (hour_start > 23) or (hour_end < 0) or (hour_end > 23):
479 raise ValueError("Hours must be between 0 and 23 inclusive.")
481 hour_constraint = iris.Constraint(hour=lambda h: hour_start <= h.point <= hour_end)
482 return hour_constraint
485def combine_constraints(
486 constraint: iris.Constraint = None, **kwargs
487) -> iris.Constraint:
488 """
489 Operator that combines multiple constraints into one.
491 Arguments
492 ---------
493 constraint: iris.Constraint
494 First constraint to combine.
495 additional_constraint_1: iris.Constraint
496 Second constraint to combine. This must be a named argument.
497 additional_constraint_2: iris.Constraint
498 There can be any number of additional constraint, they just need unique
499 names.
500 ...
502 Returns
503 -------
504 combined_constraint: iris.Constraint
506 Raises
507 ------
508 TypeError
509 If the provided arguments are not constraints.
510 """
511 # If the first argument is not a constraint, it is ignored. This handles the
512 # automatic passing of the previous step's output.
513 if isinstance(constraint, iris.Constraint):
514 combined_constraint = constraint
515 else:
516 combined_constraint = iris.Constraint()
518 for constr in kwargs.values():
519 combined_constraint = combined_constraint & constr
520 return combined_constraint
523def generate_attribute_constraint(
524 attribute: str, value: str | None = None, **kwargs
525) -> iris.AttributeConstraint:
526 """Generate constraint on cube attributes.
528 Constrains based on the presence of an attribute, and that attribute having
529 a particular value.
531 Arguments
532 ---------
533 attribute: str
534 Attribute to constraint on.
536 value: str
537 Attribute value to constrain on. If omitted the constraint merely checks
538 for the presence of an attribute.
540 Returns
541 -------
542 attribute_constraint: iris.Constraint
543 """
544 if value is None:
545 attribute_constraint = iris.Constraint(
546 cube_func=lambda cube: attribute in cube.attributes
547 )
548 else:
549 attribute_constraint = iris.AttributeConstraint(**{attribute: value})
550 return attribute_constraint