Coverage for C:\src\imod-python\imod\select\points.py: 92%

118 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import warnings 

2 

3import numpy as np 

4import xarray as xr 

5import xugrid as xu 

6 

7import imod 

8 

9 

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) 

22 

23 

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 

37 

38 

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}") 

43 

44 

45def _check_points(points): 

46 """ 

47 Check whether the array with points has the right and consistent shape. 

48 """ 

49 

50 shapes = __check_and_get_points_shape(points) 

51 __check_point_shapes_consistency(shapes) 

52 

53 

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 

60 

61 return np.full(shape, fill_value) 

62 

63 

64def points_in_bounds(da, **points): 

65 """ 

66 Returns whether points specified by keyword arguments fall within the bounds 

67 of ``da``. 

68 

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. 

75 

76 Returns 

77 ------- 

78 in_bounds : np.array of bools 

79 

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. 

84 

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) 

92 

93 This works for an arbitrary number of coordinates: 

94 

95 >>> y = [1.3, 2.7] 

96 >>> points_in_bounds(da, x=x, y=y) 

97 

98 """ 

99 

100 _check_points(points) 

101 

102 in_bounds = __arr_like_points(points, True) 

103 

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") 

110 

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) 

116 

117 return in_bounds 

118 

119 

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 ) 

136 

137 return points, inside 

138 

139 

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]) 

144 

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() 

156 

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") 

162 

163 # Take care of decreasing coordinates 

164 if x_decreasing: 

165 ixs = ncell - ixs 

166 else: 

167 ixs = ixs - 1 

168 

169 return ixs 

170 

171 

172def points_indices(da, out_of_bounds="raise", **points): 

173 """ 

174 Get the indices for points as defined by the arrays x and y. 

175 

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. 

181 

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 

192 

193 Returns 

194 ------- 

195 indices : dict of {coordinate: xr.DataArray with indices} 

196 

197 Examples 

198 -------- 

199 

200 To extract values: 

201 

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) 

208 

209 Or shorter, using dictionary unpacking: 

210 

211 >>> indices = imod.select.points_indices(da, x=x, y=y) 

212 >>> selection = da.isel(**indices) 

213 

214 To set values (in a new array), the following will do the trick: 

215 

216 >>> empty = xr.full_like(da, np.nan) 

217 >>> empty.data[indices["y"].values, indices["x"].values] = values_to_set 

218 

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! 

222 

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) 

227 

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") 

237 

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 

242 

243 return indices 

244 

245 

246def points_values(da, out_of_bounds="raise", **points): 

247 """ 

248 Get values from specified points. 

249 

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. 

253 

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 

268 

269 Examples 

270 -------- 

271 

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) 

275 

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) 

283 

284 indices = imod.select.points.points_indices( 

285 da, out_of_bounds=out_of_bounds, **iterable_points 

286 ) 

287 selection = da.isel(**indices) 

288 

289 return selection 

290 

291 

292def points_set_values(da, values, out_of_bounds="raise", **points): 

293 """ 

294 Set values at specified points. 

295 

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. 

299 

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. 

312 

313 Returns 

314 ------- 

315 da : xr.DataArray 

316 DataArray with values set at the point locations. 

317 

318 Examples 

319 -------- 

320 

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) 

325 

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] 

336 

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() 

340 

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 ) 

352 

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

357 

358 # Set values in dask or numpy array 

359 da.data[tuple(indices)] = values 

360 return da