Coverage for src/CSET/operators/constraints.py: 93%
115 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 09:36 +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 = 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)
100 varname_constraint.varname = varname_copy
102 return varname_constraint
105def generate_level_constraint(
106 coordinate: str, levels: int | list[int] | str, **kwargs
107) -> iris.Constraint:
108 """Generate constraint for particular levels on the specified coordinate.
110 Operator that generates a constraint to constrain to specific model or
111 pressure levels. If no levels are specified then any cube with the specified
112 coordinate is rejected.
114 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"``
115 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic.
117 Arguments
118 ---------
119 coordinate: str
120 Level coordinate name about which to constraint.
121 levels: int | list[int] | str
122 CF compliant level points, ``"*"`` for retrieving all levels, or
123 ``[]`` for no levels.
125 Returns
126 -------
127 constraint: iris.Constraint
129 Notes
130 -----
131 Due to the specification of ``coordinate`` as an argument any iterable
132 coordinate can be stratified with this function. Therefore,
133 ``"realization"`` is a valid option. Subsequently, ``levels`` specifies the
134 ensemble members, or group of ensemble members you wish to constrain your
135 results over.
136 """
137 # If asterisks, then return all levels for given coordinate.
138 if levels == "*":
139 return iris.Constraint(**{coordinate: lambda cell: True})
140 else:
141 # Ensure is iterable.
142 if not isinstance(levels, Iterable):
143 levels = [levels]
145 # When no levels specified reject cube with level coordinate.
146 if len(levels) == 0:
148 def no_levels(cube):
149 # Reject cubes for which coordinate exists.
150 return not cube.coords(coordinate)
152 return iris.Constraint(cube_func=no_levels)
154 # Filter the coordinate to the desired levels.
155 # Dictionary unpacking is used to provide programmatic keyword arguments.
156 return iris.Constraint(**{coordinate: levels})
159def generate_remove_single_level_constraint(
160 coord: str, level: int = 0, **kwargs
161) -> iris.Constraint:
162 """
163 Generate a constraint to remove a single model level number.
165 Operator that returns a constraint to remove the given level. By
166 default the first level is removed (assumed to be
167 level zero). However, any level can be removed.
169 Arguments
170 ---------
171 coord: str
172 The coordinate for which the level is to be removed.
173 level: int
174 Default is 0. The model level number to remove.
176 Returns
177 -------
178 iris.Constraint
180 Notes
181 -----
182 This operator is primarily used to ensure the levels are consistent
183 as some level sets (e.g. specific humidity) will be on the same level set
184 but have a different number of levels (e.g 71 instead of expected 70).
185 """
186 return iris.Constraint(**{coord: lambda m: m.point != level})
189def generate_cell_methods_constraint(
190 cell_methods: list,
191 varname: str | None = None,
192 coord: iris.coords.Coord | None = None,
193 interval: str | None = None,
194 comment: str | None = None,
195 **kwargs,
196) -> iris.Constraint:
197 """Generate constraint from cell methods.
199 Operator that takes a list of cell methods and generates a constraint from
200 that. Use [] to specify non-aggregated data.
202 Arguments
203 ---------
204 cell_methods: list
205 cube.cell_methods for filtering.
206 varname: str, optional
207 CF compliant name of variable.
208 coord: iris.coords.Coord, optional
209 iris.coords.Coord to which the cell method is applied to.
210 interval: str, optional
211 interval over which the cell method is applied to (e.g. 1 hour).
212 comment: str, optional
213 any comments in Cube meta data associated with the cell method.
215 Returns
216 -------
217 cell_method_constraint: iris.Constraint
218 """
219 if len(cell_methods) == 0:
221 def check_no_aggregation(cube: iris.cube.Cube) -> bool:
222 """Check that any cell methods are "point", meaning no aggregation."""
223 return {cm.method for cm in cube.cell_methods} <= {"point"}
225 def check_cell_sum(cube: iris.cube.Cube) -> bool:
226 """Check that any cell methods are "sum"."""
227 return {cm.method for cm in cube.cell_methods} == {"sum"}
229 def check_cell_mean(cube: iris.cube.Cube) -> bool:
230 """Check that any cell methods are "mean"."""
231 return {cm.method for cm in cube.cell_methods} == {"mean"}
233 if varname:
234 # Require number_of_lightning_flashes to be "sum" cell_method input.
235 # Require surface_microphyisical_rainfall_amount and surface_microphysical_snowfall_amount to be "sum" cell_method inputs.
236 if ("lightning" in varname) or (
237 "surface_microphysical" in varname and "amount" in varname
238 ):
239 cell_methods_constraint = iris.Constraint(cube_func=check_cell_sum)
240 return cell_methods_constraint
241 # Require climatological ancillary as time-average mean.
242 if ("albedo" in varname) or ( 242 ↛ 249line 242 didn't jump to line 249 because the condition on line 242 was always true
243 "ocean" in varname and "chlorophyll" in varname
244 ):
245 cell_methods_constraint = iris.Constraint(cube_func=check_cell_mean)
246 return cell_methods_constraint
248 # If no variable name set, assume require instantaneous cube.
249 cell_methods_constraint = iris.Constraint(cube_func=check_no_aggregation)
251 else:
252 # If cell_method constraint set in recipe, check for required input.
253 def check_cell_methods(cube: iris.cube.Cube) -> bool:
254 return all(
255 iris.coords.CellMethod(
256 method=cm, coords=coord, intervals=interval, comments=comment
257 )
258 in cube.cell_methods
259 for cm in cell_methods
260 )
262 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)
264 return cell_methods_constraint
267def generate_time_constraint(
268 time_start: str, time_end: str | None = None, **kwargs
269) -> iris.Constraint:
270 """Generate constraint between times.
272 Operator that takes one or two ISO 8601 date strings, and returns a
273 constraint that selects values between those dates (inclusive).
275 Arguments
276 ---------
277 time_start: str | datetime.datetime | cftime.datetime
278 ISO date for lower bound
280 time_end: str | datetime.datetime | cftime.datetime
281 ISO date for upper bound. If omitted it defaults to the same as
282 time_start
284 Returns
285 -------
286 time_constraint: iris.Constraint
287 """
288 if isinstance(time_start, str):
289 pdt_start, offset_start = operator_utils.pdt_fromisoformat(time_start)
290 else:
291 pdt_start, offset_start = time_start, timedelta(0)
293 if time_end is None:
294 pdt_end, offset_end = time_start, offset_start
295 elif isinstance(time_end, str):
296 pdt_end, offset_end = operator_utils.pdt_fromisoformat(time_end)
297 print(pdt_end)
298 print(offset_end)
299 else:
300 pdt_end, offset_end = time_end, timedelta(0)
302 if offset_start is None:
303 offset_start = timedelta(0)
304 if offset_end is None:
305 offset_end = timedelta(0)
307 time_constraint = iris.Constraint(
308 time=lambda t: (
309 (pdt_start <= (t.point - offset_start))
310 and ((t.point - offset_end) <= pdt_end)
311 )
312 )
314 return time_constraint
317def generate_area_constraint(
318 lat_start: float | None,
319 lat_end: float | None,
320 lon_start: float | None,
321 lon_end: float | None,
322 **kwargs,
323) -> iris.Constraint:
324 """Generate an area constraint between latitude/longitude limits.
326 Operator that takes a set of latitude and longitude limits and returns a
327 constraint that selects grid values only inside that area. Works with the
328 data's native grid so is defined within the rotated pole CRS.
330 Alternatively, all arguments may be None to indicate the area should not be
331 constrained. This is useful to allow making subsetting an optional step in a
332 processing pipeline.
334 Arguments
335 ---------
336 lat_start: float | None
337 Latitude value for lower bound
338 lat_end: float | None
339 Latitude value for top bound
340 lon_start: float | None
341 Longitude value for left bound
342 lon_end: float | None
343 Longitude value for right bound
345 Returns
346 -------
347 area_constraint: iris.Constraint
348 """
349 # Check all arguments are defined, or all are None.
350 if not (
351 all(
352 (
353 isinstance(lat_start, numbers.Real),
354 isinstance(lat_end, numbers.Real),
355 isinstance(lon_start, numbers.Real),
356 isinstance(lon_end, numbers.Real),
357 )
358 )
359 or all((lat_start is None, lat_end is None, lon_start is None, lon_end is None))
360 ):
361 raise TypeError("Bounds must real numbers, or all None.")
363 # Don't constrain area if all arguments are None.
364 if lat_start is None: # Only need to check once, as they will be the same.
365 # An empty constraint allows everything.
366 return iris.Constraint()
368 # Handle bounds crossing the date line.
369 if lon_end < lon_start: 369 ↛ 370line 369 didn't jump to line 370 because the condition on line 369 was never true
370 lon_end = lon_end + 360
372 def bound_lat(cell: iris.coords.Cell) -> bool:
373 return lat_start < cell < lat_end
375 def bound_lon(cell: iris.coords.Cell) -> bool:
376 # Adjust cell values to handle crossing the date line.
377 if cell < lon_start:
378 cell = cell + 360
379 return lon_start < cell < lon_end
381 area_constraint = iris.Constraint(
382 coord_values={"grid_latitude": bound_lat, "grid_longitude": bound_lon}
383 )
384 return area_constraint
387def generate_remove_single_ensemble_member_constraint(
388 ensemble_member: int = 0, **kwargs
389) -> iris.Constraint:
390 """
391 Generate a constraint to remove a single ensemble member.
393 Operator that returns a constraint to remove the given ensemble member. By
394 default the ensemble member removed is the control member (assumed to have
395 a realization of zero). However, any ensemble member can be removed, thus
396 allowing a non-zero control member to be removed if the control is a
397 different member.
399 Arguments
400 ---------
401 ensemble_member: int
402 Default is 0. The ensemble member realization to remove.
404 Returns
405 -------
406 iris.Constraint
408 Notes
409 -----
410 This operator is primarily used to remove the control member to allow
411 ensemble metrics to be calculated without the control member. For
412 example, the ensemble mean is not normally calculated including the
413 control member. It is particularly useful to remove the control member
414 when it is not an equally-likely member of the ensemble.
415 """
416 return iris.Constraint(realization=lambda m: m.point != ensemble_member)
419def generate_realization_constraint(
420 ensemble_members: int | list[int], **kwargs
421) -> iris.Constraint:
422 """
423 Generate a constraint to subset ensemble members.
425 Operator that is given a list of ensemble members and returns a constraint
426 to select those ensemble members. This operator is particularly useful for
427 subsetting ensembles.
429 Arguments
430 ---------
431 ensemble_members: int | list[int]
432 The ensemble members to be subsetted over.
434 Returns
435 -------
436 iris.Constraint
437 """
438 # Ensure ensemble_members is iterable.
439 ensemble_members = iter_maybe(ensemble_members)
440 return iris.Constraint(realization=ensemble_members)
443def generate_hour_constraint(
444 hour_start: int,
445 hour_end: int | None = None,
446 **kwargs,
447) -> iris.Constraint:
448 """Generate an hour constraint between hour of day limits.
450 Operator that takes a set of hour of day limits and returns a constraint that
451 selects only hours within that time frame regardless of day.
453 Alternatively, the result can be constrained to a single hour by just entering
454 a starting hour.
456 Should any sub-hourly data be given these will have the same hour coordinate
457 (e.g., 12:00 and 12:05 both have an hour coordinate of 12) all
458 times will be selected with this constraint.
460 Arguments
461 ---------
462 hour_start: int
463 The hour of day for the lower bound, within 0 to 23.
464 hour_end: int | None
465 The hour of day for the upper bound, within 0 to 23. Alternatively,
466 set to None if only one hour required.
468 Returns
469 -------
470 hour_constraint: iris.Constraint
472 Raises
473 ------
474 ValueError
475 If the provided arguments are outside of the range 0 to 23.
476 """
477 if hour_end is None:
478 hour_end = hour_start
480 if (hour_start < 0) or (hour_start > 23) or (hour_end < 0) or (hour_end > 23):
481 raise ValueError("Hours must be between 0 and 23 inclusive.")
483 hour_constraint = iris.Constraint(hour=lambda h: hour_start <= h.point <= hour_end)
484 return hour_constraint
487def combine_constraints(
488 constraint: iris.Constraint = None, **kwargs
489) -> iris.Constraint:
490 """
491 Operator that combines multiple constraints into one.
493 Arguments
494 ---------
495 constraint: iris.Constraint
496 First constraint to combine.
497 additional_constraint_1: iris.Constraint
498 Second constraint to combine. This must be a named argument.
499 additional_constraint_2: iris.Constraint
500 There can be any number of additional constraint, they just need unique
501 names.
502 ...
504 Returns
505 -------
506 combined_constraint: iris.Constraint
508 Raises
509 ------
510 TypeError
511 If the provided arguments are not constraints.
512 """
513 # If the first argument is not a constraint, it is ignored. This handles the
514 # automatic passing of the previous step's output.
515 if isinstance(constraint, iris.Constraint):
516 combined_constraint = constraint
517 else:
518 combined_constraint = iris.Constraint()
520 for constr in kwargs.values():
521 combined_constraint = combined_constraint & constr
522 return combined_constraint
525def generate_attribute_constraint(
526 attribute: str, value: str | None = None, **kwargs
527) -> iris.AttributeConstraint:
528 """Generate constraint on cube attributes.
530 Constrains based on the presence of an attribute, and that attribute having
531 a particular value.
533 Arguments
534 ---------
535 attribute: str
536 Attribute to constraint on.
538 value: str
539 Attribute value to constrain on. If omitted the constraint merely checks
540 for the presence of an attribute.
542 Returns
543 -------
544 attribute_constraint: iris.Constraint
545 """
546 if value is None:
547 attribute_constraint = iris.Constraint(
548 cube_func=lambda cube: attribute in cube.attributes
549 )
550 else:
551 attribute_constraint = iris.AttributeConstraint(**{attribute: value})
552 return attribute_constraint