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