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
« 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
5import imod
6from imod.util.imports import MissingOptionalModule
8try:
9 import shapely.geometry as sg
10except ImportError:
11 sg = MissingOptionalModule("shapely")
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
22@numba.njit
23def _increment(x0, x1):
24 if x1 > x0:
25 return 1
26 else:
27 return -1
30@numba.njit
31def _draw_line(xs, ys, x0, x1, y0, y1, xmin, xmax, ymin, ymax):
32 """
33 Generate line cell coordinates.
35 Based on "A Fast Voxel Traversal Algorithm for Ray Tracing" by Amanatides
36 & Woo, 1987.
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.
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
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.
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
97 # Avoid ZeroDivision
98 if no_dx:
99 t_x = 0.0
100 else:
101 t_x = 1.0 / a_x
103 if no_dy:
104 t_y = 0.0
105 else:
106 t_y = 1.0 / a_y
108 # Vector equations
109 def x(t):
110 return x0 + a_x * t
112 def y(t):
113 return y0 + a_y * t
115 # Set increments; 1 or -1
116 x_increment = _increment(x0, x1)
117 y_increment = _increment(y0, y1)
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
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
143 # Collection of results
144 ixs = []
145 iys = []
146 segment_length = []
147 dxs = []
148 dys = []
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))
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
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)
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]
184 if y_increment == -1:
185 cellboundary_y = ys[iy]
186 else:
187 cellboundary_y = ys[iy + y_increment]
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))
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
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
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))
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
242 return ixs[use], iys[use], segment_length[use], dxs[use], dys[use]
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])
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
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()
277 ixs = []
278 iys = []
279 sdxs = []
280 sdys = []
281 segments = []
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])
301 ixs.append(i)
302 iys.append(j)
303 sdxs.append(sdx)
304 sdys.append(sdy)
305 segments.append(segment_length)
307 if len(ixs) == 0:
308 raise ValueError("Linestring does not intersect data")
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)
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
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")
340 return section
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.
351 Adapted from Metpy:
352 https://github.com/Unidata/MetPy/blob/master/metpy/interpolate/slices.py
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``
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.
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")
381 linecoords = [start, end]
382 return _cross_section(data, linecoords)
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.
393 Adapted from Metpy:
394 https://github.com/Unidata/MetPy/blob/master/metpy/interpolate/slices.py
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``
403 coordinate.
404 linestring : shapely.geometry.LineString
405 Shapely geometry designating the linestring along which to sample the
406 cross section.
408 Note that a LineString can easily be taken from a geopandas.GeoDataFrame
409 using the .geometry attribute. Please refer to the examples.
411 Returns
412 -------
413 `xarray.DataArray` or `xarray.Dataset`
414 The interpolated cross section, with new index dimension along the
415 cross-section.
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.
423 >>> geodataframe = gpd.read_file("cross_section.shp")
424 >>> linestring = geodataframe.geometry[0]
425 >>> section = cross_section_linestring(data, linestring)
427 Or, construct the linestring directly in Python:
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)
433 If you have drawn multiple cross sections within a shapefile, simply loop
434 over the linestrings:
436 >>> sections = [cross_section_linestring(data, ls) for ls in geodataframe.geometry]
438 """
439 linecoords = np.array(linestring.coords)
440 return _cross_section(data, linecoords)