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

1# © Crown copyright, Met Office (2022-2024) 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. 

14 

15"""Operators to extract a transect given a tuple of xy coords to start/finish.""" 

16 

17import logging 

18 

19import iris 

20import numpy as np 

21 

22from CSET.operators._utils import get_cube_yxcoordname 

23 

24 

25def _check_within_bounds(point: tuple[float, float], lat_coord, lon_coord): 

26 """Check if the point (lat, lon) is within the bounds of the data.""" 

27 lon_min = min(lon_coord.points) 

28 lon_max = max(lon_coord.points) 

29 lat_min = min(lat_coord.points) 

30 lat_max = max(lat_coord.points) 

31 if ( 

32 lat_min > point[0] 

33 or lat_max < point[0] 

34 or lon_min > point[1] 

35 or lon_max < point[1] 

36 ): 

37 raise IndexError( 

38 f"Point {point} not between {(lat_min, lon_min)} and {(lat_max, lon_max)}" 

39 ) 

40 

41 

42def calc_transect( 

43 input: iris.cube.Cube | iris.cube.CubeList, startcoords: tuple, endcoords: tuple 

44): 

45 """Compute transect between startcoords and endcoords. 

46 

47 Computes a transect for a given cube/cubelist containing at least latitude 

48 and longitude coordinates, using an appropriate sampling interval along the 

49 transect based on grid spacing. Also decides a suitable x coordinate to plot along 

50 the transect - see notes for details on this. 

51 

52 Arguments 

53 --------- 

54 

55 input: Cube | CubeList 

56 An iris cube or cubelist containing latitude and longitude coordinate dimensions, 

57 to compute the transect on. 

58 startcoords: tuple 

59 A tuple containing the start coordinates for the transect using model coordinates, 

60 ordered (latitude,longitude). 

61 endcoords: tuple 

62 A tuple containing the end coordinates for the transect using model coordinates, 

63 ordered (latitude,longitude). 

64 

65 Returns 

66 ------- 

67 output: Cube | CubeList 

68 A cube containing at least pressure and the coordinate specified by coord, for 

69 the transect specified between startcoords and endcoords. 

70 

71 Notes 

72 ----- 

73 This operator uses the iris.Nearest method to interpolate the specific point along 

74 the transect. 

75 Identification of an appropriate coordinate to plot along the x axis is done by 

76 determining the largest distance out of latitude and longitude. For example, if 

77 the transect is 90 degrees (west to east), then delta latitude is zero, so it will 

78 always plot as a function of longitude. Vice versa if the transect is 180 degrees 

79 (south to north). If a transect is 45 degrees, and delta latitude and longitude 

80 are the same, it will default to plotting as a function of longitude. Note this 

81 doesn't affect the transect plot, its purely for interpretation with some appropriate 

82 x axis labelling/points. 

83 """ 

84 # If function is passed a single cube, make this iterable. 

85 if type(input) is iris.cube.Cube: 

86 input = iris.cube.CubeList([input]) 

87 

88 # To store final transects 

89 output = iris.cube.CubeList() 

90 

91 for cube in input: 

92 # Find out xy coord name 

93 lat_name, lon_name = get_cube_yxcoordname(cube) 

94 

95 lon_coord = cube.coord(lon_name) 

96 lat_coord = cube.coord(lat_name) 

97 

98 _check_within_bounds(startcoords, lat_coord, lon_coord) 

99 _check_within_bounds(endcoords, lat_coord, lon_coord) 

100 

101 # Compute vector distance between start and end points in degrees. 

102 dist_deg = np.sqrt( 

103 (startcoords[0] - endcoords[0]) ** 2 + (startcoords[1] - endcoords[1]) ** 2 

104 ) 

105 

106 # Compute minimum gap between x/y spatial coords. 

107 lon_min = np.abs(np.min(lon_coord.points[1:] - lon_coord.points[:-1])) 

108 lat_min = np.abs(np.min(lat_coord.points[1:] - lat_coord.points[:-1])) 

109 

110 # For scenarios where coord is at 90 degree to the grid (i.e. no 

111 # latitude/longitude change). Only xmin or ymin will be zero, not both 

112 # (otherwise startcoords and endcoords the same). 

113 if startcoords[1] == endcoords[1]: 

114 # Along latitude. 

115 lon_pnts = np.repeat(startcoords[1], int(dist_deg / lat_min)) 

116 lat_pnts = np.linspace( 

117 startcoords[0], endcoords[0], int(dist_deg / lat_min) 

118 ) 

119 transect_coord = "latitude" 

120 elif startcoords[0] == endcoords[0]: 

121 # Along longitude. 

122 lon_pnts = np.linspace( 

123 startcoords[1], endcoords[1], int(dist_deg / lon_min) 

124 ) 

125 lat_pnts = np.repeat(startcoords[0], int(dist_deg / lon_min)) 

126 transect_coord = "longitude" 

127 else: 

128 # Else use the smallest grid space in x or y direction. 

129 number_of_points = int(dist_deg / np.min([lon_min, lat_min])) 

130 lon_pnts = np.linspace(startcoords[1], endcoords[1], number_of_points) 

131 lat_pnts = np.linspace(startcoords[0], endcoords[0], number_of_points) 

132 

133 # If change in latitude larger than change in longitude: 

134 if abs(startcoords[0] - endcoords[0]) > abs(startcoords[1] - endcoords[1]): 

135 transect_coord = "latitude" 

136 else: 

137 transect_coord = "longitude" 

138 

139 # Create cubelist to store interpolated points along transect. 

140 interpolated_cubes = iris.cube.CubeList() 

141 

142 # Iterate over all points along transect, lon_pnts will be the same shape as 

143 # lat_pnts so we can use either to iterate over. 

144 for i in range(0, lon_pnts.shape[0]): 

145 logging.info("%s/%s", i + 1, lon_pnts.shape[0]) 

146 

147 # Get point along transect. 

148 cube_slice = cube.interpolate( 

149 [(lon_name, lon_pnts[i]), (lat_name, lat_pnts[i])], 

150 iris.analysis.Nearest(), 

151 ) 

152 

153 # Remove existing coordinates ready to add one single map coordinate 

154 # Note latitude/longitude cubes may have additional AuxCoord to remove 

155 for coord_name in [ 

156 "latitude", 

157 "longitude", 

158 "grid_latitude", 

159 "grid_longitude", 

160 ]: 

161 if cube_slice.coords(coord_name): 

162 cube_slice.remove_coord(coord_name) 

163 

164 if transect_coord == "latitude": 

165 dist_coord = iris.coords.DimCoord( 

166 lat_pnts[i], long_name="latitude", units="degrees" 

167 ) 

168 cube_slice.add_aux_coord(dist_coord) 

169 cube_slice = iris.util.new_axis(cube_slice, scalar_coord="latitude") 

170 else: 

171 dist_coord = iris.coords.DimCoord( 

172 lon_pnts[i], long_name="longitude", units="degrees" 

173 ) 

174 cube_slice.add_aux_coord(dist_coord) 

175 cube_slice = iris.util.new_axis(cube_slice, scalar_coord="longitude") 

176 

177 interpolated_cubes.append(cube_slice) 

178 

179 # Concatenate into single cube. 

180 interpolated_cubes = interpolated_cubes.concatenate() 

181 

182 # Add metadata to interpolated cubes showing coordinates. 

183 interpolated_cubes[0].attributes["transect_coords"] = ( 

184 f"{startcoords[0]}_{startcoords[1]}_{endcoords[0]}_{endcoords[1]}" 

185 ) 

186 

187 # If concatenation successful, should be CubeList with one cube left. 

188 assert len(interpolated_cubes) == 1, ( 

189 f"len(interpolated_cubes) = {len(interpolated_cubes)}" 

190 ) 

191 

192 # Carry over useful attributes to transect, like model identifier. 

193 for key, value in cube.attributes.items(): 

194 interpolated_cubes[0].attributes.setdefault(key, value) 

195 

196 output.append(interpolated_cubes[0]) 

197 

198 if len(output) == 1: 

199 return output[0] 

200 else: 

201 return output