Coverage for C:\src\imod-python\imod\select\cross_sections.py: 82%

207 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import numba 

2import numpy as np 

3import xarray as xr 

4 

5import imod 

6from imod.util.imports import MissingOptionalModule 

7 

8try: 

9 import shapely.geometry as sg 

10except ImportError: 

11 sg = MissingOptionalModule("shapely") 

12 

13 

14@numba.njit 

15def _index(edges, coord, coef): 

16 if coef < 0: 

17 return np.searchsorted(edges, coord, side="left") - 1 

18 else: 

19 return np.searchsorted(edges, coord, side="right") - 1 

20 

21 

22@numba.njit 

23def _increment(x0, x1): 

24 if x1 > x0: 

25 return 1 

26 else: 

27 return -1 

28 

29 

30@numba.njit 

31def _draw_line(xs, ys, x0, x1, y0, y1, xmin, xmax, ymin, ymax): 

32 """ 

33 Generate line cell coordinates. 

34 

35 Based on "A Fast Voxel Traversal Algorithm for Ray Tracing" by Amanatides 

36 & Woo, 1987. 

37 

38 Note: out of bound values are marked with -1. This might be slightly 

39 misleading, since -1 is valid indexing value in Python. However, it is 

40 actually desirable in this case, since the values will be taken from a 

41 DataArray/Dataset. In this case, the section will automatically have the 

42 right dimension size, even with skipped parts at the start and end of the 

43 cross section. 

44 

45 Parameters 

46 ---------- 

47 xs : np.array 

48 x coordinates of cell edges 

49 ys : np.array 

50 y coordinates of cell edges 

51 x0, y0 : float 

52 starting point coordinates 

53 x1, y1 : float 

54 end point coordinates 

55 xmin, xmax, ymin, ymax : float 

56 grid bounds 

57 

58 Returns 

59 ------- 

60 ixs : np.array 

61 column indices, out of bound values marked with -1. 

62 iys : np.array 

63 row indices, out of bound values marked with -1. 

64 segment_length : np.array 

65 length of segment per sampled cell 

66 dxs : np.array 

67 length along column of segment per sampled cell 

68 dys : np.array 

69 length along row of segment per sampled cell 

70 """ 

71 # Vector equations given by: 

72 # x = x0 + a_x * t 

73 # y = y0 + a_y * t 

74 # where t is the vector length; t_x and t_y are the inverse of a_x and a_y, 

75 # respectively. 

76 # 

77 # The algorithm works cell by cell. It computes the distance to cross the 

78 # cell boundary, in both x and y. We express this distance in the vector 

79 # space, tmax_x to the next x boundary, and tmax_y to the next y boundary. 

80 # We compare tmax_x and tmax_y and move the shortest distance. 

81 # We move by updating the indices, ix and iy, and taking a step along t. 

82 # 

83 # If we're outside the grid, we initialize the starting position on a cell 

84 # boundary. 

85 # 

86 # Addtionally, there's some logic to deal with start and end points that 

87 # fall outside of the grid bounding box. 

88 

89 dx = x1 - x0 

90 dy = y1 - y0 

91 length = np.sqrt(dx**2 + dy**2) 

92 a_x = dx / length 

93 a_y = dy / length 

94 no_dx = dx == 0.0 

95 no_dy = dy == 0.0 

96 

97 # Avoid ZeroDivision 

98 if no_dx: 

99 t_x = 0.0 

100 else: 

101 t_x = 1.0 / a_x 

102 

103 if no_dy: 

104 t_y = 0.0 

105 else: 

106 t_y = 1.0 / a_y 

107 

108 # Vector equations 

109 def x(t): 

110 return x0 + a_x * t 

111 

112 def y(t): 

113 return y0 + a_y * t 

114 

115 # Set increments; 1 or -1 

116 x_increment = _increment(x0, x1) 

117 y_increment = _increment(y0, y1) 

118 

119 # Initialize start position of t 

120 if x0 < xmin: 

121 t = t_x * (xmin - x0) 

122 elif x0 > xmax: 

123 t = t_x * (xmax - x0) 

124 elif y0 < ymin: 

125 t = t_y * (ymin - y0) 

126 elif y0 > ymax: 

127 t = t_y * (ymax - y0) 

128 else: # within grid bounds 

129 t = 0.0 

130 

131 # Initialize end position of t 

132 if x1 < xmin: 

133 t_end = t_x * (xmin - x0) 

134 elif x1 > xmax: 

135 t_end = t_x * (xmax - x0) 

136 elif y1 < ymin: 

137 t_end = t_y * (ymin - y0) 

138 elif y1 > ymax: 

139 t_end = t_y * (ymax - y0) 

140 else: # within grid bounds 

141 t_end = length 

142 

143 # Collection of results 

144 ixs = [] 

145 iys = [] 

146 segment_length = [] 

147 dxs = [] 

148 dys = [] 

149 

150 # Store how much of the cross-section has no data 

151 skipped_start = t 

152 skipped_end = length - t_end 

153 if skipped_start > 0.0: 

154 ixs.append(-1) 

155 iys.append(-1) 

156 segment_length.append(skipped_start) 

157 dxs.append(x(skipped_start) - x(0)) 

158 dys.append(y(skipped_start) - y(0)) 

159 

160 # Arbitrarily large number so it's always the largest one 

161 if no_dx: 

162 tmax_x = 1.0e20 

163 if no_dy: 

164 tmax_y = 1.0e20 

165 

166 # First step 

167 ix = _index(xs, x(t), a_x) 

168 iy = _index(ys, y(t), a_y) 

169 ixs.append(ix) 

170 iys.append(iy) 

171 

172 # Main loop, move through grid 

173 ncol = xs.size - 1 

174 nrow = ys.size - 1 

175 tstep = 0.0 

176 while ix < ncol and iy < nrow: 

177 # Compute distance to cell boundary 

178 # We need the start of the cell if we're moving in negative direction. 

179 if x_increment == -1: 

180 cellboundary_x = xs[ix] 

181 else: 

182 cellboundary_x = xs[ix + x_increment] 

183 

184 if y_increment == -1: 

185 cellboundary_y = ys[iy] 

186 else: 

187 cellboundary_y = ys[iy + y_increment] 

188 

189 # Compute max distance to move along t. 

190 # Checks for infinite slopes 

191 if not no_dx: 

192 # dx_t = cellboundary_x - x(t) 

193 tmax_x = t_x * (cellboundary_x - x(t)) 

194 if not no_dy: 

195 # dy_t = cellboundary_y - y(t) 

196 tmax_y = t_y * (cellboundary_y - y(t)) 

197 

198 # Find which dimension requires smallest step along t 

199 if tmax_x == tmax_y: 

200 ix += x_increment 

201 iy += y_increment 

202 tstep = tmax_x 

203 elif tmax_x < tmax_y: 

204 ix += x_increment 

205 tstep = tmax_x 

206 else: 

207 iy += y_increment 

208 tstep = tmax_y 

209 

210 if (t + tstep) < t_end: 

211 dxs.append(x(t + tstep) - x(t)) 

212 dys.append(y(t + tstep) - y(t)) 

213 t += tstep 

214 # Store 

215 ixs.append(ix) 

216 iys.append(iy) 

217 segment_length.append(tstep) 

218 else: 

219 tstep = t_end - t 

220 # Store final step 

221 dxs.append(x(t + tstep) - x(t)) 

222 dys.append(y(t + tstep) - y(t)) 

223 segment_length.append(tstep) 

224 break 

225 

226 if skipped_end > 0.0: 

227 segment_length.append(skipped_end) 

228 ixs.append(-1) 

229 iys.append(-1) 

230 dxs.append(x(length) - x(t_end)) 

231 dys.append(y(length) - y(t_end)) 

232 

233 # Because of numerical precision, extremely small segments might be 

234 # included. Those are filtered out here. 

235 ixs = np.array(ixs) 

236 iys = np.array(iys) 

237 dxs = np.array(dxs) 

238 dys = np.array(dys) 

239 segment_length = np.array(segment_length) 

240 use = np.abs(segment_length) > 1.0e-6 

241 

242 return ixs[use], iys[use], segment_length[use], dxs[use], dys[use] 

243 

244 

245def _bounding_box(xmin, xmax, ymin, ymax): 

246 a = (xmin, ymin) 

247 b = (xmax, ymin) 

248 c = (xmax, ymax) 

249 d = (xmin, ymax) 

250 return sg.Polygon([a, b, c, d]) 

251 

252 

253def _cross_section(data, linecoords): 

254 dx, xmin, xmax, dy, ymin, ymax = imod.util.spatial.spatial_reference(data) 

255 if isinstance(dx, float): 

256 dx = np.full(data.x.size, dx) 

257 if isinstance(dy, float): 

258 dy = np.full(data.y.size, dy) 

259 x_decreasing = data.indexes["x"].is_monotonic_decreasing 

260 y_decreasing = data.indexes["y"].is_monotonic_decreasing 

261 

262 # Create vertex edges 

263 nrow = data.y.size 

264 ncol = data.x.size 

265 ys = np.full(nrow + 1, ymin) 

266 xs = np.full(ncol + 1, xmin) 

267 # Always increasing 

268 if x_decreasing: 

269 xs[1:] += np.abs(dx[::-1]).cumsum() 

270 else: 

271 xs[1:] += np.abs(dx).cumsum() 

272 if y_decreasing: 

273 ys[1:] += np.abs(dy[::-1]).cumsum() 

274 else: 

275 ys[1:] += np.abs(dy).cumsum() 

276 

277 ixs = [] 

278 iys = [] 

279 sdxs = [] 

280 sdys = [] 

281 segments = [] 

282 

283 bounding_box = _bounding_box(xmin, xmax, ymin, ymax) 

284 for start, end in zip(linecoords[:-1], linecoords[1:]): 

285 linestring = sg.LineString([start, end]) 

286 if not linestring.length: 

287 continue 

288 if linestring.intersects(bounding_box): 

289 x0, y0 = start 

290 x1, y1 = end 

291 i, j, segment_length, sdx, sdy = _draw_line( 

292 xs, ys, x0, x1, y0, y1, xmin, xmax, ymin, ymax 

293 ) 

294 else: # append the linestring in full as nodata section 

295 i = np.array([-1]) 

296 j = np.array([-1]) 

297 sdx = np.array([-1]) 

298 sdy = np.array([-1]) 

299 segment_length = np.array([linestring.length]) 

300 

301 ixs.append(i) 

302 iys.append(j) 

303 sdxs.append(sdx) 

304 sdys.append(sdy) 

305 segments.append(segment_length) 

306 

307 if len(ixs) == 0: 

308 raise ValueError("Linestring does not intersect data") 

309 

310 # Concatenate into a single array 

311 ixs = np.concatenate(ixs) 

312 iys = np.concatenate(iys) 

313 sdxs = np.concatenate(sdxs) 

314 sdys = np.concatenate(sdys) 

315 segments = np.concatenate(segments) 

316 

317 # Flip around indexes 

318 if x_decreasing: 

319 ixs = ncol - 1 - ixs 

320 ixs[ixs >= ncol] = -1 

321 if y_decreasing: 

322 iys = nrow - 1 - iys 

323 iys[iys >= nrow] = -1 

324 

325 # Select data 

326 # use .where to get rid of out of nodata parts 

327 ind_x = xr.DataArray(ixs, dims=["s"]) 

328 ind_y = xr.DataArray(iys, dims=["s"]) 

329 section = data.isel(x=ind_x, y=ind_y).where(ind_x >= 0) 

330 # Set dimension values 

331 section.coords["s"] = segments.cumsum() - 0.5 * segments 

332 section = section.assign_coords(ds=("s", segments)) 

333 section = section.assign_coords(dx=("s", sdxs)) 

334 section = section.assign_coords(dy=("s", sdys)) 

335 # Without this sort, the is_increasing_monotonic property of the "s" index 

336 # in the DataArray returns False, and plotting the DataArray as a quadmesh 

337 # appears to fail. TODO: investigate, seems like an xarray issue. 

338 section = section.sortby("s") 

339 

340 return section 

341 

342 

343def cross_section_line(data, start, end): 

344 r""" 

345 Obtain an interpolated cross-sectional slice through gridded data. 

346 Utilizing the interpolation functionality in ``xarray``, this function 

347 takes a vertical cross-sectional slice along a line through the given 

348 data on a regular (possibly non-equidistant) grid, which is given as an 

349 `xarray.DataArray` so that we can utilize its coordinate data. 

350 

351 Adapted from Metpy: 

352 https://github.com/Unidata/MetPy/blob/master/metpy/interpolate/slices.py 

353 

354 Parameters 

355 ---------- 

356 data: `xarray.DataArray` or `xarray.Dataset` 

357 Three- (or higher) dimensional field(s) to interpolate. The DataArray 

358 (or each DataArray in the Dataset) must have been parsed by MetPy and 

359 include both an x and y coordinate dimension and the added ``crs`` 

360 

361 coordinate. 

362 start: (2, ) array_like 

363 A latitude-longitude pair designating the start point of the cross 

364 section. 

365 end: (2, ) array_like 

366 A latitude-longitude pair designating the end point of the cross 

367 section. 

368 

369 Returns 

370 ------- 

371 `xarray.DataArray` or `xarray.Dataset` 

372 The interpolated cross section, with new dimension "s" along the 

373 cross-section. The cellsizes along "s" are given in the "ds" coordinate. 

374 """ 

375 # Check for intersection 

376 _, xmin, xmax, _, ymin, ymax = imod.util.spatial.spatial_reference(data) 

377 bounding_box = _bounding_box(xmin, xmax, ymin, ymax) 

378 if not sg.LineString([start, end]).intersects(bounding_box): 

379 raise ValueError("Line does not intersect data") 

380 

381 linecoords = [start, end] 

382 return _cross_section(data, linecoords) 

383 

384 

385def cross_section_linestring(data, linestring): 

386 r""" 

387 Obtain an interpolated cross-sectional slice through gridded data. 

388 Utilizing the interpolation functionality in ``xarray``, this function 

389 takes a vertical cross-sectional slice along a linestring through the given 

390 data on a regular grid, which is given as an `xarray.DataArray` so that 

391 we can utilize its coordinate data. 

392 

393 Adapted from Metpy: 

394 https://github.com/Unidata/MetPy/blob/master/metpy/interpolate/slices.py 

395 

396 Parameters 

397 ---------- 

398 data: `xarray.DataArray` or `xarray.Dataset` 

399 Three- (or higher) dimensional field(s) to interpolate. The DataArray 

400 (or each DataArray in the Dataset) must have been parsed by MetPy and 

401 include both an x and y coordinate dimension and the added ``crs`` 

402 

403 coordinate. 

404 linestring : shapely.geometry.LineString 

405 Shapely geometry designating the linestring along which to sample the 

406 cross section. 

407 

408 Note that a LineString can easily be taken from a geopandas.GeoDataFrame 

409 using the .geometry attribute. Please refer to the examples. 

410 

411 Returns 

412 ------- 

413 `xarray.DataArray` or `xarray.Dataset` 

414 The interpolated cross section, with new index dimension along the 

415 cross-section. 

416 

417 Examples 

418 -------- 

419 Load a shapefile (that you might have drawn before using a GIS program), 

420 take a linestring from it, and use it to extract the data for a cross 

421 section. 

422 

423 >>> geodataframe = gpd.read_file("cross_section.shp") 

424 >>> linestring = geodataframe.geometry[0] 

425 >>> section = cross_section_linestring(data, linestring) 

426 

427 Or, construct the linestring directly in Python: 

428 

429 >>> import shapely.geometry as sg 

430 >>> linestring = sg.LineString([(0.0, 1.0), (5.0, 5.0), (7.5, 5.0)]) 

431 >>> section = cross_section_linestring(data, linestring) 

432 

433 If you have drawn multiple cross sections within a shapefile, simply loop 

434 over the linestrings: 

435 

436 >>> sections = [cross_section_linestring(data, ls) for ls in geodataframe.geometry] 

437 

438 """ 

439 linecoords = np.array(linestring.coords) 

440 return _cross_section(data, linecoords)