Coverage for C:\src\imod-python\imod\select\points.py: 77%
118 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 warnings
3import numpy as np
4import xarray as xr
5import xugrid as xu
7import imod
10def get_unstructured_cell2d_from_xy(uda, **points):
11 # Unstructured grids always require to be tested both on x and y coordinates
12 # to see if points are within bounds.
13 for coord in ["x", "y"]:
14 if coord not in points.keys():
15 raise KeyError(
16 f"Missing {coord} in point coordinates."
17 "Unstructured grids require both an x and y coordinate"
18 "to get cell indices."
19 )
20 xy = np.column_stack([points["x"], points["y"]])
21 return uda.ugrid.grid.locate_points(xy)
24def __check_and_get_points_shape(points) -> dict:
25 """Check whether points have the right shape"""
26 shapes = {}
27 for coord, value in points.items():
28 arr = np.atleast_1d(value)
29 points[coord] = arr
30 shape = arr.shape
31 if not len(shape) == 1:
32 raise ValueError(
33 f"Coordinate {coord} is not one-dimensional, but has shape: {shape}"
34 )
35 shapes[coord] = shape
36 return shapes
39def __check_point_shapes_consistency(shapes):
40 if not len(set(shapes.values())) == 1:
41 msg = "\n".join([f"{coord}: {shape}" for coord, shape in shapes.items()])
42 raise ValueError(f"Shapes of coordinates do match each other:\n{msg}")
45def _check_points(points):
46 """
47 Check whether the array with points has the right and consistent shape.
48 """
50 shapes = __check_and_get_points_shape(points)
51 __check_point_shapes_consistency(shapes)
54def __arr_like_points(points, fill_value):
55 """
56 Return array with the same shape as the first array provided in points.
57 """
58 first_value = next(iter(points.values()))
59 shape = np.atleast_1d(first_value).shape
61 return np.full(shape, fill_value)
64def points_in_bounds(da, **points):
65 """
66 Returns whether points specified by keyword arguments fall within the bounds
67 of ``da``.
69 Parameters
70 ----------
71 da : xr.DataArray
72 points : keyword arguments of coordinate=values
73 keyword arguments specifying coordinate and values. Please refer to the
74 examples.
76 Returns
77 -------
78 in_bounds : np.array of bools
80 Examples
81 --------
82 Create the DataArray, then use the keyword arguments to define along which
83 coordinate to check whether the points are within bounds.
85 >>> nrow, ncol = 3, 4
86 >>> data = np.arange(12.0).reshape(nrow, ncol)
87 >>> coords = {"x": [0.5, 1.5, 2.5, 3.5], "y": [2.5, 1.5, 0.5]}
88 >>> dims = ("y", "x")
89 >>> da = xr.DataArray(data, coords, dims)
90 >>> x = [0.4, 2.6]
91 >>> points_in_bounds(da, x=x)
93 This works for an arbitrary number of coordinates:
95 >>> y = [1.3, 2.7]
96 >>> points_in_bounds(da, x=x, y=y)
98 """
100 _check_points(points)
102 in_bounds = __arr_like_points(points, True)
104 if isinstance(da, xu.UgridDataArray):
105 index = get_unstructured_cell2d_from_xy(da, **points)
106 # xu.Ugrid2d.locate_points makes cells outside grid -1
107 in_bounds = index > -1
108 points.pop("x")
109 points.pop("y")
111 for key, x in points.items():
112 da_x = da.coords[key]
113 _, xmin, xmax = imod.util.spatial.coord_reference(da_x)
114 # Inplace bitwise operator
115 in_bounds &= (x >= xmin) & (x < xmax)
117 return in_bounds
120def check_points_in_bounds(da, points, out_of_bounds):
121 inside = points_in_bounds(da, **points)
122 # Error handling
123 msg = "Not all points are located within the bounds of the DataArray"
124 if not inside.all():
125 if out_of_bounds == "raise":
126 raise ValueError(msg)
127 elif out_of_bounds == "warn":
128 warnings.warn(msg)
129 elif out_of_bounds == "ignore":
130 points = {dim: x[inside] for dim, x in points.items()}
131 else:
132 raise ValueError(
133 f"Unrecognized option {out_of_bounds} for out_of_bounds, "
134 "should be one of: error, warn, ignore."
135 )
137 return points, inside
140def _get_indices_1d(da, coordname, x):
141 x = np.atleast_1d(x)
142 x_decreasing = da.indexes[coordname].is_monotonic_decreasing
143 dx, xmin, _ = imod.util.spatial.coord_reference(da.coords[coordname])
145 ncell = da[coordname].size
146 # Compute edges
147 xs = np.full(ncell + 1, xmin)
148 # Turn dx into array
149 if isinstance(dx, float):
150 dx = np.full(ncell, dx)
151 # Always increasing
152 if x_decreasing:
153 xs[1:] += np.abs(dx[::-1]).cumsum()
154 else:
155 xs[1:] += np.abs(dx).cumsum()
157 # From np.searchsorted docstring:
158 # Find the indices into a sorted array a such that, if the corresponding
159 # elements in v were inserted before the indices, the order of a would
160 # be preserved.
161 ixs = np.searchsorted(xs, x, side="right")
163 # Take care of decreasing coordinates
164 if x_decreasing:
165 ixs = ncell - ixs
166 else:
167 ixs = ixs - 1
169 return ixs
172def points_indices(da, out_of_bounds="raise", **points):
173 """
174 Get the indices for points as defined by the arrays x and y.
176 Not all points may be located in the bounds of the DataArray. By default,
177 this function raises an error. This behavior can be controlled with the
178 ``out_of_bounds`` argument. If ``out_of_bounds`` is set to "warn" or
179 "ignore", out of bounds point are removed. Which points have been removed
180 is visible in the ``index`` coordinate of the resulting selection.
182 Parameters
183 ----------
184 da : xr.DataArray
185 out_of_bounds : {"raise", "warn", "ignore"}, default: "raise"
186 What to do if the points are not located in the bounds of the
187 DataArray:
188 - "raise": raise an exception
189 - "warn": raise a warning, and ignore the missing points
190 - "ignore": ignore the missing points
191 points : keyword arguments of coordinates and values
193 Returns
194 -------
195 indices : dict of {coordinate: xr.DataArray with indices}
197 Examples
198 --------
200 To extract values:
202 >>> x = [1.0, 2.2, 3.0]
203 >>> y = [4.0, 5.6, 7.0]
204 >>> indices = imod.select.points_indices(da, x=x, y=y)
205 >>> ind_y = indices["y"]
206 >>> ind_x = indices["x"]
207 >>> selection = da.isel(x=ind_x, y=ind_y)
209 Or shorter, using dictionary unpacking:
211 >>> indices = imod.select.points_indices(da, x=x, y=y)
212 >>> selection = da.isel(**indices)
214 To set values (in a new array), the following will do the trick:
216 >>> empty = xr.full_like(da, np.nan)
217 >>> empty.data[indices["y"].values, indices["x"].values] = values_to_set
219 Unfortunately, at the time of writing, xarray's .sel method does not
220 support setting values yet. The method here works for both numpy and dask
221 arrays, but you'll have to manage dimensions yourself!
223 The ``imod.select.points_set_values()`` function will take care of the
224 dimensions.
225 """
226 points, inside = check_points_in_bounds(da, points, out_of_bounds)
228 indices = {}
229 index = np.arange(len(inside))[inside]
230 if isinstance(da, xu.UgridDataArray):
231 ind_arr = get_unstructured_cell2d_from_xy(da, **points)
232 ind_da = xr.DataArray(ind_arr, coords={"index": index}, dims=("index",))
233 face_dim = da.ugrid.grid.face_dimension
234 indices[face_dim] = ind_da
235 points.pop("x")
236 points.pop("y")
238 for coordname, value in points.items():
239 ind_arr = _get_indices_1d(da, coordname, value)
240 ind_da = xr.DataArray(ind_arr, coords={"index": index}, dims=("index",))
241 indices[coordname] = ind_da
243 return indices
246def points_values(da, out_of_bounds="raise", **points):
247 """
248 Get values from specified points.
250 This function will raise a ValueError if the points fall outside of the
251 bounds of the DataArray to avoid ambiguous behavior. Use the
252 ``imod.select.points_in_bounds`` function to detect these points.
254 Parameters
255 ----------
256 da : xr.DataArray
257 out_of_bounds : {"raise", "warn", "ignore"}, default: "raise"
258 What to do if the points are not located in the bounds of the
259 DataArray:
260 - "raise": raise an exception
261 - "warn": raise a warning, and ignore the missing points
262 - "ignore": ignore the missing points
263 points : keyword arguments of coordinate=values
264 keyword arguments specifying coordinate and values.
265 Returns
266 -------
267 selection : xr.DataArray
269 Examples
270 --------
272 >>> x = [1.0, 2.2, 3.0]
273 >>> y = [4.0, 5.6, 7.0]
274 >>> selection = imod.select.points_values(da, x=x, y=y)
276 """
277 iterable_points = {}
278 for coordname, value in points.items():
279 if not isinstance(da, xu.UgridDataArray) and (coordname not in da.coords):
280 raise ValueError(f'DataArray has no coordinate "{coordname}"')
281 # contents must be iterable
282 iterable_points[coordname] = np.atleast_1d(value)
284 indices = imod.select.points.points_indices(
285 da, out_of_bounds=out_of_bounds, **iterable_points
286 )
287 selection = da.isel(**indices)
289 return selection
292def points_set_values(da, values, out_of_bounds="raise", **points):
293 """
294 Set values at specified points.
296 This function will raise a ValueError if the points fall outside of the
297 bounds of the DataArray to avoid ambiguous behavior. Use the
298 ``imod.select.points_in_bounds`` function to detect these points.
300 Parameters
301 ----------
302 da : xr.DataArray
303 values : (int, float) or array of (int, float)
304 out_of_bounds : {"raise", "warn", "ignore"}, default: "raise"
305 What to do if the points are not located in the bounds of the
306 DataArray:
307 - "raise": raise an exception
308 - "warn": raise a warning, and ignore the missing points
309 - "ignore": ignore the missing points
310 points : keyword arguments of coordinate=values
311 keyword arguments specifying coordinate and values.
313 Returns
314 -------
315 da : xr.DataArray
316 DataArray with values set at the point locations.
318 Examples
319 --------
321 >>> x = [1.0, 2.2, 3.0]
322 >>> y = [4.0, 5.6, 7.0]
323 >>> values = [10.0, 11.0, 12.0]
324 >>> out = imod.select.points_set_values(da, values, x=x, y=y)
326 """
327 points, inside = check_points_in_bounds(da, points, out_of_bounds)
328 if not isinstance(values, (bool, float, int, str)): # then it might be an array
329 if len(values) != len(inside):
330 raise ValueError(
331 "Shape of ``values`` does not match shape of coordinates."
332 f"Shape of values: {values.shape}; shape of coordinates: {inside.shape}."
333 )
334 # Make sure a list or tuple is indexable by inside.
335 values = np.atleast_1d(values)[inside]
337 # Avoid side-effects just in case
338 # Load into memory, so values can be set efficiently via numpy indexing.
339 da = da.copy(deep=True).load()
341 sel_indices = {}
342 for coordname in points.keys():
343 if coordname not in da.coords:
344 raise ValueError(f'DataArray has no coordinate "{coordname}"')
345 underlying_dims = da.coords[coordname].dims
346 if len(underlying_dims) != 1:
347 raise ValueError(f"Coordinate {coordname} is not one-dimensional")
348 # Use the first and only element of underlying_dims
349 sel_indices[underlying_dims[0]] = _get_indices_1d(
350 da, coordname, points[coordname]
351 )
353 # Collect indices in the right order
354 indices = []
355 for dim in da.dims:
356 indices.append(sel_indices.get(dim, slice(None, None)))
358 # Set values in dask or numpy array
359 da.data[tuple(indices)] = values
360 return da