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

1 import numba

2 import numpy as np

3 import xarray as xr

4

5 import imod

6 from imod.util.imports import MissingOptionalModule

7

8 try:

9 import shapely.geometry as sg

10 except ImportError:

11 sg = MissingOptionalModule("shapely")

12

13

14@numba.njit

15 def _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

23 def _increment(x0, x1):

24 if x1 > x0:

25 return 1

26 else:

27 return -1

28

29

30@numba.njit

31 def _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

245 def _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

253 def _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

343 def 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

385 def 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)