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 

24logger = logging.getLogger(__name__) 

25 

26 

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

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

29 lon_min = min(lon_coord.points) 

30 lon_max = max(lon_coord.points) 

31 lat_min = min(lat_coord.points) 

32 lat_max = max(lat_coord.points) 

33 if ( 

34 lat_min > point[0] 

35 or lat_max < point[0] 

36 or lon_min > point[1] 

37 or lon_max < point[1] 

38 ): 

39 raise IndexError( 

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

41 ) 

42 

43 

44def calc_transect( 

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

46): 

47 """Compute transect between startcoords and endcoords. 

48 

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

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

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

52 the transect - see notes for details on this. 

53 

54 Arguments 

55 --------- 

56 

57 input: Cube | CubeList 

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

59 to compute the transect on. 

60 startcoords: tuple 

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

62 ordered (latitude,longitude). 

63 endcoords: tuple 

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

65 ordered (latitude,longitude). 

66 

67 Returns 

68 ------- 

69 output: Cube | CubeList 

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

71 the transect specified between startcoords and endcoords. 

72 

73 Notes 

74 ----- 

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

76 the transect. 

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

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

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

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

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

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

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

84 x axis labelling/points. 

85 """ 

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

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

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

89 

90 # To store final transects 

91 output = iris.cube.CubeList() 

92 

93 for cube in input: 

94 # Find out xy coord name 

95 lat_name, lon_name = get_cube_yxcoordname(cube) 

96 

97 lon_coord = cube.coord(lon_name) 

98 lat_coord = cube.coord(lat_name) 

99 

100 _check_within_bounds(startcoords, lat_coord, lon_coord) 

101 _check_within_bounds(endcoords, lat_coord, lon_coord) 

102 

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

104 dist_deg = np.sqrt( 

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

106 ) 

107 

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

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

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

111 

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

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

114 # (otherwise startcoords and endcoords the same). 

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

116 # Along latitude. 

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

118 lat_pnts = np.linspace( 

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

120 ) 

121 transect_coord = "latitude" 

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

123 # Along longitude. 

124 lon_pnts = np.linspace( 

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

126 ) 

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

128 transect_coord = "longitude" 

129 else: 

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

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

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

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

134 

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

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

137 transect_coord = "latitude" 

138 else: 

139 transect_coord = "longitude" 

140 

141 # Create cubelist to store interpolated points along transect. 

142 interpolated_cubes = iris.cube.CubeList() 

143 

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

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

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

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

148 

149 # Get point along transect. 

150 cube_slice = cube.interpolate( 

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

152 iris.analysis.Nearest(), 

153 ) 

154 

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

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

157 for coord_name in [ 

158 "latitude", 

159 "longitude", 

160 "grid_latitude", 

161 "grid_longitude", 

162 ]: 

163 if cube_slice.coords(coord_name): 

164 cube_slice.remove_coord(coord_name) 

165 

166 if transect_coord == "latitude": 

167 dist_coord = iris.coords.DimCoord( 

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

169 ) 

170 cube_slice.add_aux_coord(dist_coord) 

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

172 else: 

173 dist_coord = iris.coords.DimCoord( 

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

175 ) 

176 cube_slice.add_aux_coord(dist_coord) 

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

178 

179 interpolated_cubes.append(cube_slice) 

180 

181 # Concatenate into single cube. 

182 interpolated_cubes = interpolated_cubes.concatenate() 

183 

184 # Add metadata to interpolated cubes showing coordinates. 

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

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

187 ) 

188 

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

190 assert len(interpolated_cubes) == 1, ( 

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

192 ) 

193 

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

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

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

197 

198 output.append(interpolated_cubes[0]) 

199 

200 if len(output) == 1: 

201 return output[0] 

202 else: 

203 return output