Coverage for src/CSET/operators/constraints.py: 96%
54 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 08:37 +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 datetime import datetime
18from typing import Union
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_model_level_constraint(
68 model_level_number: Union[int, str], **kwargs
69) -> iris.Constraint:
70 """Generate constraint for a particular model level number.
72 Operator that takes a CF compliant model_level_number string, and uses iris to
73 generate a constraint to be passed into the read operator to minimize the
74 CubeList the read operator loads and speed up loading.
76 Arguments
77 ---------
78 model_level_number: str
79 CF compliant model level number.
81 Returns
82 -------
83 model_level_number_constraint: iris.Constraint
84 """
85 # Cast to int in case a string is given.
86 model_level_number = int(model_level_number)
87 model_level_number_constraint = iris.Constraint(
88 model_level_number=model_level_number
89 )
90 return model_level_number_constraint
93def generate_pressure_level_constraint(
94 pressure_levels: Union[int, list[int]], **kwargs
95) -> iris.Constraint:
96 """Generate constraint for the specified pressure_levels.
98 If no pressure levels are specified then any cube with a pressure coordinate
99 is rejected.
101 Arguments
102 ---------
103 pressure_levels: int|list
104 List of integer pressure levels in hPa either as single integer
105 for a single level or a list of multiple integers.
107 Returns
108 -------
109 pressure_constraint: iris.Constraint
110 """
111 # If pressure_level is an integer it is converted into a list.
112 if isinstance(pressure_levels, int):
113 pressure_levels = [pressure_levels]
114 if len(pressure_levels) == 0:
115 # If none specified reject cubes with pressure level coordinate.
116 def no_pressure_coordinate(cube):
117 try:
118 cube.coord("pressure")
119 except iris.exceptions.CoordinateNotFoundError:
120 return True
121 return False
123 pressure_constraint = iris.Constraint(cube_func=no_pressure_coordinate)
124 else:
125 pressure_constraint = iris.Constraint(pressure=pressure_levels)
127 return pressure_constraint
130def generate_cell_methods_constraint(cell_methods: list, **kwargs) -> iris.Constraint:
131 """Generate constraint from cell methods.
133 Operator that takes a list of cell methods and generates a constraint from
134 that.
136 Arguments
137 ---------
138 cell_methods: list
139 cube.cell_methods for filtering
141 Returns
142 -------
143 cell_method_constraint: iris.Constraint
144 """
146 def check_cell_methods(cube: iris.cube.Cube):
147 return cube.cell_methods == tuple(cell_methods)
149 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)
150 return cell_methods_constraint
153def generate_time_constraint(
154 time_start: str, time_end: str = None, **kwargs
155) -> iris.AttributeConstraint:
156 """Generate constraint between times.
158 Operator that takes one or two ISO 8601 date strings, and returns a
159 constraint that selects values between those dates (inclusive).
161 Arguments
162 ---------
163 time_start: str | datetime.datetime
164 ISO date for lower bound
166 time_end: str | datetime.datetime
167 ISO date for upper bound. If omitted it defaults to the same as
168 time_start
170 Returns
171 -------
172 time_constraint: iris.Constraint
173 """
174 if isinstance(time_start, str):
175 time_start = datetime.fromisoformat(time_start)
176 if time_end is None:
177 time_end = time_start
178 elif isinstance(time_end, str):
179 time_end = datetime.fromisoformat(time_end)
180 time_constraint = iris.Constraint(time=lambda t: time_start <= t.point <= time_end)
181 return time_constraint
184def generate_area_constraint(
185 lat_start: float | str,
186 lat_end: float | str,
187 lon_start: float | str,
188 lon_end: float | str,
189 **kwargs,
190) -> iris.Constraint:
191 """Generate an area constraint between latitude/longitude limits.
193 Operator that takes a set of latitude and longitude limits and returns a
194 constraint that selects grid values only inside that area. Works with the
195 data's native grid so is defined within the rotated pole CRS.
197 Arguments
198 ---------
199 lat_start: float
200 Latitude value for lower bound
201 lat_end: float
202 Latitude value for top bound
203 lon_start: float
204 Longitude value for left bound
205 lon_end: float
206 Longitude value for right bound
208 Returns
209 -------
210 area_constraint: iris.Constraint
211 """
212 if lat_start == "None": 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true
213 return iris.Constraint()
215 area_constraint = iris.Constraint( 215 ↛ exitline 215 didn't jump to the function exit
216 coord_values={
217 "grid_latitude": lambda cell: lat_start < cell < lat_end,
218 "grid_longitude": lambda cell: lon_start < cell < lon_end,
219 }
220 )
221 return area_constraint
224def combine_constraints(
225 constraint: iris.Constraint = None, **kwargs
226) -> iris.Constraint:
227 """
228 Operator that combines multiple constraints into one.
230 Arguments
231 ---------
232 constraint: iris.Constraint
233 First constraint to combine.
234 additional_constraint_1: iris.Constraint
235 Second constraint to combine. This must be a named argument.
236 additional_constraint_2: iris.Constraint
237 There can be any number of additional constraint, they just need unique
238 names.
239 ...
241 Returns
242 -------
243 combined_constraint: iris.Constraint
245 Raises
246 ------
247 TypeError
248 If the provided arguments are not constraints.
249 """
250 # If the first argument is not a constraint, it is ignored. This handles the
251 # automatic passing of the previous step's output.
252 if isinstance(constraint, iris.Constraint):
253 combined_constraint = constraint
254 else:
255 combined_constraint = iris.Constraint()
257 for constr in kwargs.values():
258 combined_constraint = combined_constraint & constr
259 return combined_constraint