Coverage for src/CSET/operators/constraints.py: 95%
45 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 15:05 +0000
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 15:05 +0000
1# Copyright 2022 Met Office and 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."""
17from collections.abc import Iterable
18from datetime import datetime
20import iris
21import iris.cube
22import iris.exceptions
25def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint:
26 """Generate constraint from STASH code.
28 Operator that takes a stash string, and uses iris to generate a constraint
29 to be passed into the read operator to minimize the CubeList the read
30 operator loads and speed up loading.
32 Arguments
33 ---------
34 stash: str
35 stash code to build iris constraint, such as "m01s03i236"
37 Returns
38 -------
39 stash_constraint: iris.AttributeConstraint
40 """
41 # At a later stage str list an option to combine constraints. Arguments
42 # could be a list of stash codes that combined build the constraint.
43 stash_constraint = iris.AttributeConstraint(STASH=stash)
44 return stash_constraint
47def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
48 """Generate constraint from variable name.
50 Operator that takes a CF compliant variable name string, and uses iris to
51 generate a constraint to be passed into the read operator to minimize the
52 CubeList the read operator loads and speed up loading.
54 Arguments
55 ---------
56 varname: str
57 CF compliant name of variable. Needed later for LFRic.
59 Returns
60 -------
61 varname_constraint: iris.Constraint
62 """
63 varname_constraint = iris.Constraint(name=varname)
64 return varname_constraint
67def generate_level_constraint(
68 coordinate: str, levels: int | list[int], **kwargs
69) -> iris.Constraint:
70 """Generate constraint for particular levels on the specified coordinate.
72 Operator that generates a constraint to constrain to specific model or
73 pressure levels. If no levels are specified then any cube with the specified
74 coordinate is rejected.
76 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"``
77 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic.
79 Arguments
80 ---------
81 coordinate: str
82 Level coordinate name about which to constraint.
83 levels: int | list[int]
84 CF compliant levels.
86 Returns
87 -------
88 constraint: iris.Constraint
89 """
90 # Ensure is iterable.
91 if not isinstance(levels, Iterable):
92 levels = [levels]
94 # When no levels specified reject cube with level coordinate.
95 if len(levels) == 0:
97 def no_levels(cube):
98 # Reject cubes for which coordinate exists.
99 return not bool(cube.coords(coordinate))
101 return iris.Constraint(cube_func=no_levels)
103 # Filter the coordinate to the desired levels.
104 # Dictionary unpacking is used to provide programmatic keyword arguments.
105 return iris.Constraint(**{coordinate: levels})
108def generate_cell_methods_constraint(cell_methods: list, **kwargs) -> iris.Constraint:
109 """Generate constraint from cell methods.
111 Operator that takes a list of cell methods and generates a constraint from
112 that.
114 Arguments
115 ---------
116 cell_methods: list
117 cube.cell_methods for filtering
119 Returns
120 -------
121 cell_method_constraint: iris.Constraint
122 """
124 def check_cell_methods(cube: iris.cube.Cube):
125 return cube.cell_methods == tuple(cell_methods)
127 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)
128 return cell_methods_constraint
131def generate_time_constraint(
132 time_start: str, time_end: str = None, **kwargs
133) -> iris.AttributeConstraint:
134 """Generate constraint between times.
136 Operator that takes one or two ISO 8601 date strings, and returns a
137 constraint that selects values between those dates (inclusive).
139 Arguments
140 ---------
141 time_start: str | datetime.datetime
142 ISO date for lower bound
144 time_end: str | datetime.datetime
145 ISO date for upper bound. If omitted it defaults to the same as
146 time_start
148 Returns
149 -------
150 time_constraint: iris.Constraint
151 """
152 if isinstance(time_start, str):
153 time_start = datetime.fromisoformat(time_start)
154 if time_end is None:
155 time_end = time_start
156 elif isinstance(time_end, str):
157 time_end = datetime.fromisoformat(time_end)
158 time_constraint = iris.Constraint(time=lambda t: time_start <= t.point <= time_end)
159 return time_constraint
162def generate_area_constraint(
163 lat_start: float | str,
164 lat_end: float | str,
165 lon_start: float | str,
166 lon_end: float | str,
167 **kwargs,
168) -> iris.Constraint:
169 """Generate an area constraint between latitude/longitude limits.
171 Operator that takes a set of latitude and longitude limits and returns a
172 constraint that selects grid values only inside that area. Works with the
173 data's native grid so is defined within the rotated pole CRS.
175 Arguments
176 ---------
177 lat_start: float
178 Latitude value for lower bound
179 lat_end: float
180 Latitude value for top bound
181 lon_start: float
182 Longitude value for left bound
183 lon_end: float
184 Longitude value for right bound
186 Returns
187 -------
188 area_constraint: iris.Constraint
189 """
190 if lat_start == "None": 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 return iris.Constraint()
193 area_constraint = iris.Constraint( 193 ↛ exitline 193 didn't jump to the function exit
194 coord_values={
195 "grid_latitude": lambda cell: lat_start < cell < lat_end,
196 "grid_longitude": lambda cell: lon_start < cell < lon_end,
197 }
198 )
199 return area_constraint
202def combine_constraints(
203 constraint: iris.Constraint = None, **kwargs
204) -> iris.Constraint:
205 """
206 Operator that combines multiple constraints into one.
208 Arguments
209 ---------
210 constraint: iris.Constraint
211 First constraint to combine.
212 additional_constraint_1: iris.Constraint
213 Second constraint to combine. This must be a named argument.
214 additional_constraint_2: iris.Constraint
215 There can be any number of additional constraint, they just need unique
216 names.
217 ...
219 Returns
220 -------
221 combined_constraint: iris.Constraint
223 Raises
224 ------
225 TypeError
226 If the provided arguments are not constraints.
227 """
228 # If the first argument is not a constraint, it is ignored. This handles the
229 # automatic passing of the previous step's output.
230 if isinstance(constraint, iris.Constraint):
231 combined_constraint = constraint
232 else:
233 combined_constraint = iris.Constraint()
235 for constr in kwargs.values():
236 combined_constraint = combined_constraint & constr
237 return combined_constraint