Coverage for src/CSET/operators/constraints.py: 93%

1# © Crown copyright, Met Office (2022-2025) and CSET contributors.

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

7# http://www.apache.org/licenses/LICENSE-2.0

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 if isinstance(varname, str): 80 ↛ 82line 80 didn't jump to line 82 because the condition on line 80 was always true

81 varname = [varname]

82 varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"])

83 varname.extend(["u_wind_at_10m", "v_wind_at_10m"])

85 # Case 2: Multiple varnames

86 if isinstance(varname, (list, tuple)):

87 varname_constraint = iris.Constraint(

88 cube_func=lambda cube: (

89 cube.long_name in varname

90 or cube.standard_name in varname

91 or cube.var_name in varname

92 )

93 )

95 else:

96 varname_constraint = iris.Constraint(name=varname)

98 return varname_constraint

100

101def generate_level_constraint(

102 coordinate: str, levels: int | list[int] | str, **kwargs

103) -> iris.Constraint:

104 """Generate constraint for particular levels on the specified coordinate.

105

106 Operator that generates a constraint to constrain to specific model or

107 pressure levels. If no levels are specified then any cube with the specified

108 coordinate is rejected.

109

110 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"``

111 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic.

112

113 Arguments

114 ---------

115 coordinate: str

116 Level coordinate name about which to constraint.

117 levels: int | list[int] | str

118 CF compliant level points, ``"*"`` for retrieving all levels, or

119 ``[]`` for no levels.

120

121 Returns

122 -------

123 constraint: iris.Constraint

124

125 Notes

126 -----

127 Due to the specification of ``coordinate`` as an argument any iterable

128 coordinate can be stratified with this function. Therefore,

129 ``"realization"`` is a valid option. Subsequently, ``levels`` specifies the

130 ensemble members, or group of ensemble members you wish to constrain your

131 results over.

132 """

133 # If asterisks, then return all levels for given coordinate.

134 if levels == "*":

135 return iris.Constraint(**{coordinate: lambda cell: True})

136 else:

137 # Ensure is iterable.

138 if not isinstance(levels, Iterable):

139 levels = [levels]

140

141 # When no levels specified reject cube with level coordinate.

142 if len(levels) == 0:

143

144 def no_levels(cube):

145 # Reject cubes for which coordinate exists.

146 return not cube.coords(coordinate)

147

148 return iris.Constraint(cube_func=no_levels)

149

150 # Filter the coordinate to the desired levels.

151 # Dictionary unpacking is used to provide programmatic keyword arguments.

152 return iris.Constraint(**{coordinate: levels})

153

154

155def generate_remove_single_level_constraint(

156 coord: str, level: int = 0, **kwargs

157) -> iris.Constraint:

158 """

159 Generate a constraint to remove a single model level number.

160

161 Operator that returns a constraint to remove the given level. By

162 default the first level is removed (assumed to be

163 level zero). However, any level can be removed.

164

165 Arguments

166 ---------

167 coord: str

168 The coordinate for which the level is to be removed.

169 level: int

170 Default is 0. The model level number to remove.

171

172 Returns

173 -------

174 iris.Constraint

175

176 Notes

177 -----

178 This operator is primarily used to ensure the levels are consistent

179 as some level sets (e.g. specific humidity) will be on the same level set

180 but have a different number of levels (e.g 71 instead of expected 70).

181 """

182 return iris.Constraint(**{coord: lambda m: m.point != level})

183

184

185def generate_cell_methods_constraint(

186 cell_methods: list,

187 varname: str | None = None,

188 coord: iris.coords.Coord | None = None,

189 interval: str | None = None,

190 comment: str | None = None,

191 **kwargs,

192) -> iris.Constraint:

193 """Generate constraint from cell methods.

194

195 Operator that takes a list of cell methods and generates a constraint from

196 that. Use [] to specify non-aggregated data.

197

198 Arguments

199 ---------

200 cell_methods: list

201 cube.cell_methods for filtering.

202 varname: str, optional

203 CF compliant name of variable.

204 coord: iris.coords.Coord, optional

205 iris.coords.Coord to which the cell method is applied to.

206 interval: str, optional

207 interval over which the cell method is applied to (e.g. 1 hour).

208 comment: str, optional

209 any comments in Cube meta data associated with the cell method.

210

211 Returns

212 -------

213 cell_method_constraint: iris.Constraint

214 """

215 if len(cell_methods) == 0:

216

217 def check_no_aggregation(cube: iris.cube.Cube) -> bool:

218 """Check that any cell methods are "point", meaning no aggregation."""

219 return set(cm.method for cm in cube.cell_methods) <= {"point"}

220

221 def check_cell_sum(cube: iris.cube.Cube) -> bool:

222 """Check that any cell methods are "sum"."""

223 return set(cm.method for cm in cube.cell_methods) == {"sum"}

224

225 def check_cell_mean(cube: iris.cube.Cube) -> bool:

226 """Check that any cell methods are "mean"."""

227 return set(cm.method for cm in cube.cell_methods) == {"mean"}

228

229 if varname:

230 # Require number_of_lightning_flashes to be "sum" cell_method input.

231 # Require surface_microphyisical_rainfall_amount and surface_microphysical_snowfall_amount to be "sum" cell_method inputs.

232 if ("lightning" in varname) or (

233 "surface_microphysical" in varname and "amount" in varname

234 ):

235 cell_methods_constraint = iris.Constraint(cube_func=check_cell_sum)

236 return cell_methods_constraint

237 # Require climatological ancillary as time-average mean.

238 if ("albedo" in varname) or ( 238 ↛ 245line 238 didn't jump to line 245 because the condition on line 238 was always true

239 "ocean" in varname and "chlorophyll" in varname

240 ):

241 cell_methods_constraint = iris.Constraint(cube_func=check_cell_mean)

242 return cell_methods_constraint

243

244 # If no variable name set, assume require instantaneous cube.

245 cell_methods_constraint = iris.Constraint(cube_func=check_no_aggregation)

246

247 else:

248 # If cell_method constraint set in recipe, check for required input.

249 def check_cell_methods(cube: iris.cube.Cube) -> bool:

250 return all(

251 iris.coords.CellMethod(

252 method=cm, coords=coord, intervals=interval, comments=comment

253 )

254 in cube.cell_methods

255 for cm in cell_methods

256 )

257

258 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods)

259

260 return cell_methods_constraint

261

262

263def generate_time_constraint(

264 time_start: str, time_end: str = None, **kwargs

265) -> iris.Constraint:

266 """Generate constraint between times.

267

268 Operator that takes one or two ISO 8601 date strings, and returns a

269 constraint that selects values between those dates (inclusive).

270

271 Arguments

272 ---------

273 time_start: str | datetime.datetime | cftime.datetime

274 ISO date for lower bound

275

276 time_end: str | datetime.datetime | cftime.datetime

277 ISO date for upper bound. If omitted it defaults to the same as

278 time_start

279

280 Returns

281 -------

282 time_constraint: iris.Constraint

283 """

284 if isinstance(time_start, str):

285 pdt_start, offset_start = operator_utils.pdt_fromisoformat(time_start)

286 else:

287 pdt_start, offset_start = time_start, timedelta(0)

288

289 if time_end is None:

290 pdt_end, offset_end = time_start, offset_start

291 elif isinstance(time_end, str):

292 pdt_end, offset_end = operator_utils.pdt_fromisoformat(time_end)

293 print(pdt_end)

294 print(offset_end)

295 else:

296 pdt_end, offset_end = time_end, timedelta(0)

297

298 if offset_start is None:

299 offset_start = timedelta(0)

300 if offset_end is None:

301 offset_end = timedelta(0)

302

303 time_constraint = iris.Constraint(

304 time=lambda t: (

305 (pdt_start <= (t.point - offset_start))

306 and ((t.point - offset_end) <= pdt_end)

307 )

308 )

309

310 return time_constraint

311

312

313def generate_area_constraint(

314 lat_start: float | None,

315 lat_end: float | None,

316 lon_start: float | None,

317 lon_end: float | None,

318 **kwargs,

319) -> iris.Constraint:

320 """Generate an area constraint between latitude/longitude limits.

321

322 Operator that takes a set of latitude and longitude limits and returns a

323 constraint that selects grid values only inside that area. Works with the

324 data's native grid so is defined within the rotated pole CRS.

325

326 Alternatively, all arguments may be None to indicate the area should not be

327 constrained. This is useful to allow making subsetting an optional step in a

328 processing pipeline.

329

330 Arguments

331 ---------

332 lat_start: float | None

333 Latitude value for lower bound

334 lat_end: float | None

335 Latitude value for top bound

336 lon_start: float | None

337 Longitude value for left bound

338 lon_end: float | None

339 Longitude value for right bound

340

341 Returns

342 -------

343 area_constraint: iris.Constraint

344 """

345 # Check all arguments are defined, or all are None.

346 if not (

347 all(

348 (

349 isinstance(lat_start, numbers.Real),

350 isinstance(lat_end, numbers.Real),

351 isinstance(lon_start, numbers.Real),

352 isinstance(lon_end, numbers.Real),

353 )

354 )

355 or all((lat_start is None, lat_end is None, lon_start is None, lon_end is None))

356 ):

357 raise TypeError("Bounds must real numbers, or all None.")

358

359 # Don't constrain area if all arguments are None.

360 if lat_start is None: # Only need to check once, as they will be the same.

361 # An empty constraint allows everything.

362 return iris.Constraint()

363

364 # Handle bounds crossing the date line.

365 if lon_end < lon_start: 365 ↛ 366line 365 didn't jump to line 366 because the condition on line 365 was never true

366 lon_end = lon_end + 360

367

368 def bound_lat(cell: iris.coords.Cell) -> bool:

369 return lat_start < cell < lat_end

370

371 def bound_lon(cell: iris.coords.Cell) -> bool:

372 # Adjust cell values to handle crossing the date line.

373 if cell < lon_start:

374 cell = cell + 360

375 return lon_start < cell < lon_end

376

377 area_constraint = iris.Constraint(

378 coord_values={"grid_latitude": bound_lat, "grid_longitude": bound_lon}

379 )

380 return area_constraint

381

382

383def generate_remove_single_ensemble_member_constraint(

384 ensemble_member: int = 0, **kwargs

385) -> iris.Constraint:

386 """

387 Generate a constraint to remove a single ensemble member.

388

389 Operator that returns a constraint to remove the given ensemble member. By

390 default the ensemble member removed is the control member (assumed to have

391 a realization of zero). However, any ensemble member can be removed, thus

392 allowing a non-zero control member to be removed if the control is a

393 different member.

394

395 Arguments

396 ---------

397 ensemble_member: int

398 Default is 0. The ensemble member realization to remove.

399

400 Returns

401 -------

402 iris.Constraint

403

404 Notes

405 -----

406 This operator is primarily used to remove the control member to allow

407 ensemble metrics to be calculated without the control member. For

408 example, the ensemble mean is not normally calculated including the

409 control member. It is particularly useful to remove the control member

410 when it is not an equally-likely member of the ensemble.

411 """

412 return iris.Constraint(realization=lambda m: m.point != ensemble_member)

413

414

415def generate_realization_constraint(

416 ensemble_members: int | list[int], **kwargs

417) -> iris.Constraint:

418 """

419 Generate a constraint to subset ensemble members.

420

421 Operator that is given a list of ensemble members and returns a constraint

422 to select those ensemble members. This operator is particularly useful for

423 subsetting ensembles.

424

425 Arguments

426 ---------

427 ensemble_members: int | list[int]

428 The ensemble members to be subsetted over.

429

430 Returns

431 -------

432 iris.Constraint

433 """

434 # Ensure ensemble_members is iterable.

435 ensemble_members = iter_maybe(ensemble_members)

436 return iris.Constraint(realization=ensemble_members)

437

438

439def generate_hour_constraint(

440 hour_start: int,

441 hour_end: int = None,

442 **kwargs,

443) -> iris.Constraint:

444 """Generate an hour constraint between hour of day limits.

445

446 Operator that takes a set of hour of day limits and returns a constraint that

447 selects only hours within that time frame regardless of day.

448

449 Alternatively, the result can be constrained to a single hour by just entering

450 a starting hour.

451

452 Should any sub-hourly data be given these will have the same hour coordinate

453 (e.g., 12:00 and 12:05 both have an hour coordinate of 12) all

454 times will be selected with this constraint.

455

456 Arguments

457 ---------

458 hour_start: int

459 The hour of day for the lower bound, within 0 to 23.

460 hour_end: int | None

461 The hour of day for the upper bound, within 0 to 23. Alternatively,

462 set to None if only one hour required.

463

464 Returns

465 -------

466 hour_constraint: iris.Constraint

467

468 Raises

469 ------

470 ValueError

471 If the provided arguments are outside of the range 0 to 23.

472 """

473 if hour_end is None:

474 hour_end = hour_start

475

476 if (hour_start < 0) or (hour_start > 23) or (hour_end < 0) or (hour_end > 23):

477 raise ValueError("Hours must be between 0 and 23 inclusive.")

478

479 hour_constraint = iris.Constraint(hour=lambda h: hour_start <= h.point <= hour_end)

480 return hour_constraint

481

482

483def combine_constraints(

484 constraint: iris.Constraint = None, **kwargs

485) -> iris.Constraint:

486 """

487 Operator that combines multiple constraints into one.

488

489 Arguments

490 ---------

491 constraint: iris.Constraint

492 First constraint to combine.

493 additional_constraint_1: iris.Constraint

494 Second constraint to combine. This must be a named argument.

495 additional_constraint_2: iris.Constraint

496 There can be any number of additional constraint, they just need unique

497 names.

498 ...

499

500 Returns

501 -------

502 combined_constraint: iris.Constraint

503

504 Raises

505 ------

506 TypeError

507 If the provided arguments are not constraints.

508 """

509 # If the first argument is not a constraint, it is ignored. This handles the

510 # automatic passing of the previous step's output.

511 if isinstance(constraint, iris.Constraint):

512 combined_constraint = constraint

513 else:

514 combined_constraint = iris.Constraint()

515

516 for constr in kwargs.values():

517 combined_constraint = combined_constraint & constr

518 return combined_constraint

519

520

521def generate_attribute_constraint(

522 attribute: str, value: str = None, **kwargs

523) -> iris.AttributeConstraint:

524 """Generate constraint on cube attributes.

525

526 Constrains based on the presence of an attribute, and that attribute having

527 a particular value.

528

529 Arguments

530 ---------

531 attribute: str

532 Attribute to constraint on.

533

534 value: str

535 Attribute value to constrain on. If omitted the constraint merely checks

536 for the presence of an attribute.

537

538 Returns

539 -------

540 attribute_constraint: iris.Constraint

541 """

542 if value is None:

543 attribute_constraint = iris.Constraint(

544 cube_func=lambda cube: attribute in cube.attributes

545 )

546 else:

547 attribute_constraint = iris.AttributeConstraint(**{attribute: value})

548 return attribute_constraint