Coverage for src/CSET/operators/filters.py: 100%

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. 

14 

15"""Operators to perform various kind of filtering.""" 

16 

17from typing import Union 

18 

19import iris 

20import iris.cube 

21import iris.exceptions 

22 

23 

24def filter_cubes( 

25 cube: Union[iris.cube.Cube, iris.cube.CubeList], 

26 constraint: iris.Constraint, 

27 **kwargs, 

28) -> iris.cube.Cube: 

29 """Filter a CubeList down to a single Cube based on a constraint. 

30 

31 Arguments 

32 --------- 

33 cube: iris.cube.Cube | iris.cube.CubeList 

34 Cube(s) to filter 

35 constraint: iris.Constraint 

36 Constraint to extract 

37 

38 Returns 

39 ------- 

40 iris.cube.Cube 

41 

42 Raises 

43 ------ 

44 ValueError 

45 If the constraint doesn't produce a single cube. 

46 """ 

47 filtered_cubes = cube.extract(constraint) 

48 # Return directly if already a cube. 

49 if isinstance(filtered_cubes, iris.cube.Cube): 

50 return filtered_cubes 

51 # Check filtered cubes is a CubeList containing one cube. 

52 if isinstance(filtered_cubes, iris.cube.CubeList) and len(filtered_cubes) == 1: 

53 return filtered_cubes[0] 

54 else: 

55 raise ValueError( 

56 f"Constraint doesn't produce single cube. Constraint: {constraint}" 

57 f"\nSource: {cube}\nResult: {filtered_cubes}" 

58 ) 

59 

60 

61def filter_multiple_cubes( 

62 cubes: Union[iris.cube.Cube, iris.cube.CubeList], 

63 **kwargs, 

64) -> iris.cube.CubeList: 

65 """Filter a CubeList on multiple constraints, returning another CubeList. 

66 

67 Arguments 

68 --------- 

69 cube: iris.cube.Cube | iris.cube.CubeList 

70 Cube(s) to filter 

71 constraint: iris.Constraint 

72 Constraint to extract. This must be a named argument. There can be any 

73 number of additional constraints, they just need unique names. 

74 

75 Returns 

76 ------- 

77 iris.cube.CubeList 

78 

79 Raises 

80 ------ 

81 ValueError 

82 The constraints don't produce a single cube per constraint. 

83 """ 

84 # Ensure input is a CubeList. 

85 if isinstance(cubes, iris.cube.Cube): 

86 cubes = iris.cube.CubeList((cubes,)) 

87 if len(kwargs) < 1: 

88 raise ValueError("Must have at least one constraint.") 

89 try: 

90 filtered_cubes = cubes.extract_cubes(kwargs.values()) 

91 except iris.exceptions.ConstraintMismatchError as err: 

92 raise ValueError( 

93 "The constraints don't produce a single cube per constraint." 

94 ) from err 

95 return filtered_cubes