Coverage for test_regrid.py: 98%

385 statements  

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

1import os 

2 

3import numba 

4import numpy as np 

5import pandas as pd 

6import pytest 

7import xarray as xr 

8 

9import imod 

10 

11 

12def first(values, weights): 

13 return values[0] 

14 

15 

16def mean(values, weights): 

17 vsum = 0.0 

18 wsum = 0.0 

19 for i in range(values.size): 

20 v = values[i] 

21 vsum += v 

22 wsum += 1.0 

23 return vsum / wsum 

24 

25 

26def weightedmean(values, weights): 

27 vsum = 0.0 

28 wsum = 0.0 

29 for i in range(values.size): 

30 v = values[i] 

31 w = weights[i] 

32 vsum += w * v 

33 wsum += w 

34 return vsum / wsum 

35 

36 

37def conductance(values, weights): 

38 v_agg = 0.0 

39 w_sum = 0.0 

40 for i in range(values.size): 

41 v = values[i] 

42 w = weights[i] 

43 if np.isnan(v): 

44 continue 

45 v_agg += v * w 

46 w_sum += w 

47 if w_sum == 0: 

48 return np.nan 

49 else: 

50 return v_agg 

51 

52 

53def test_make_regrid(): 

54 if "NUMBA_DISABLE_JIT" in os.environ: 

55 pass 

56 else: 

57 # Cannot really test functionality, since it's compiled by numba at runtime 

58 # This just checks whether it's ingested okay 

59 func = imod.prepare.regrid._jit_regrid(mean, 1) 

60 assert isinstance(func, numba.core.registry.CPUDispatcher) 

61 

62 func = imod.prepare.regrid._make_regrid(mean, 1) 

63 assert isinstance(func, numba.core.registry.CPUDispatcher) 

64 

65 

66def test_regrid_1d(): 

67 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

68 dst_x = np.array([0.0, 2.5, 5.0]) 

69 alloc_len, i_w = imod.prepare.common._weights_1d(src_x, dst_x, False) 

70 inds_weights = [tuple(elem) for elem in i_w] 

71 values = np.zeros(alloc_len) 

72 weights = np.zeros(alloc_len) 

73 src = np.array([10.0, 20.0, 30.0, 40.0, 50.0]) 

74 dst = np.array([0.0, 0.0]) 

75 

76 # Regrid method 1 

77 first_regrid = imod.prepare.regrid._jit_regrid(numba.njit(first), 1) 

78 dst = first_regrid(src, dst, values, weights, *inds_weights) 

79 assert np.allclose(dst, np.array([10.0, 30.0])) 

80 

81 # Regrid method 2 

82 mean_regrid = imod.prepare.regrid._jit_regrid(numba.njit(mean), 1) 

83 dst = mean_regrid(src, dst, values, weights, *inds_weights) 

84 assert np.allclose( 

85 dst, np.array([(10.0 + 20.0 + 30.0) / 3.0, (30.0 + 40.0 + 50.0) / 3.0]) 

86 ) 

87 

88 # Regrid method 3 

89 wmean_regrid = imod.prepare.regrid._jit_regrid(numba.njit(weightedmean), 1) 

90 dst = wmean_regrid(src, dst, values, weights, *inds_weights) 

91 assert np.allclose( 

92 dst, 

93 np.array([(10.0 + 20.0 + 0.5 * 30.0) / 2.5, (30.0 * 0.5 + 40.0 + 50.0) / 2.5]), 

94 ) 

95 

96 

97def test_iter_regrid__1d(): 

98 ndim_regrid = 1 

99 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

100 dst_x = np.array([0.0, 2.5, 5.0]) 

101 alloc_len, i_w = imod.prepare.common._weights_1d(src_x, dst_x, False) 

102 inds_weights = [tuple(elem) for elem in i_w] 

103 

104 # 1D regrid over 1D array 

105 src = np.array([10.0, 20.0, 30.0, 40.0, 50.0]) 

106 dst = np.zeros(2) 

107 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

108 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid) 

109 iter_dst = iter_regrid(iter_src, iter_dst, alloc_len, *inds_weights) 

110 assert np.allclose(dst, np.array([10.0, 30.0])) 

111 

112 # 1D regrid over 2D array 

113 src = np.array([[10.0, 20.0, 30.0, 40.0, 50.0] for _ in range(3)]) 

114 dst = np.zeros((3, 2)) 

115 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

116 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid) 

117 iter_dst = iter_regrid(iter_src, iter_dst, alloc_len, *inds_weights) 

118 assert np.allclose(dst, np.array([[10.0, 30.0], [10.0, 30.0], [10.0, 30.0]])) 

119 

120 # 1D regrid over 3D array 

121 src = np.zeros((4, 3, 5)) 

122 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

123 dst = np.zeros((4, 3, 2)) 

124 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

125 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid) 

126 iter_dst = iter_regrid(iter_src, iter_dst, alloc_len, *inds_weights) 

127 compare = np.zeros((4, 3, 2)) 

128 compare[..., :] = [10.0, 30.0] 

129 assert np.allclose(dst, compare) 

130 

131 

132def test_nd_regrid__1d(): 

133 # 1D regrid over 3D array 

134 src_coords = (np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),) 

135 dst_coords = (np.array([0.0, 2.5, 5.0]),) 

136 ndim_regrid = len(src_coords) 

137 src = np.zeros((4, 3, 5)) 

138 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

139 dst = np.zeros((4, 3, 2)) 

140 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

141 

142 dst = imod.prepare.regrid._nd_regrid( 

143 src, dst, src_coords, dst_coords, iter_regrid, False 

144 ) 

145 compare = np.zeros((4, 3, 2)) 

146 compare[..., :] = [10.0, 30.0] 

147 assert np.allclose(dst, compare) 

148 

149 

150def test_nd_regrid__2d__first(): 

151 # 2D regrid over 3D array 

152 src_coords = ( 

153 np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), 

154 np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), 

155 ) 

156 dst_coords = (np.array([0.0, 2.5, 5.0]), np.array([0.0, 2.5, 5.0])) 

157 ndim_regrid = len(src_coords) 

158 src = np.zeros((4, 5, 5)) 

159 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

160 dst = np.zeros((4, 2, 2)) 

161 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

162 

163 dst = imod.prepare.regrid._nd_regrid( 

164 src, dst, src_coords, dst_coords, iter_regrid, False 

165 ) 

166 compare = np.zeros((4, 2, 2)) 

167 compare[..., :] = [10.0, 30.0] 

168 assert np.allclose(dst, compare) 

169 

170 

171def test_nd_regrid__2d__mean(): 

172 # 2D regrid over 3D array 

173 src_coords = ( 

174 np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), 

175 np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), 

176 ) 

177 dst_coords = (np.array([0.0, 2.5, 5.0]), np.array([0.0, 2.5, 5.0])) 

178 ndim_regrid = len(src_coords) 

179 src = np.zeros((4, 5, 5)) 

180 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

181 dst = np.zeros((4, 2, 2)) 

182 iter_regrid = imod.prepare.regrid._make_regrid(mean, ndim_regrid) 

183 

184 dst = imod.prepare.regrid._nd_regrid( 

185 src, dst, src_coords, dst_coords, iter_regrid, False 

186 ) 

187 compare = np.zeros((4, 2, 2)) 

188 compare[..., :] = [20.0, 40.0] 

189 assert np.allclose(dst, compare) 

190 

191 

192def test_nd_regrid__3d__first(): 

193 # 3D regrid over 3D array 

194 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

195 dst_x = np.array([0.0, 2.5, 5.0]) 

196 src_coords = [src_x for _ in range(3)] 

197 dst_coords = [dst_x for _ in range(3)] 

198 ndim_regrid = len(src_coords) 

199 src = np.zeros((5, 5, 5)) 

200 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

201 dst = np.zeros((2, 2, 2)) 

202 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

203 

204 dst = imod.prepare.regrid._nd_regrid( 

205 src, dst, src_coords, dst_coords, iter_regrid, False 

206 ) 

207 compare = np.zeros((2, 2, 2)) 

208 compare[..., :] = [10.0, 30.0] 

209 assert np.allclose(dst, compare) 

210 

211 

212def test_nd_regrid__4d3d__first(): 

213 # 3D regrid over 4D array 

214 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

215 dst_x = np.array([0.0, 2.5, 5.0]) 

216 src_coords = [src_x for _ in range(3)] 

217 dst_coords = [dst_x for _ in range(3)] 

218 ndim_regrid = len(src_coords) 

219 src = np.zeros((3, 5, 5, 5)) 

220 src[..., :] = [10.0, 20.0, 30.0, 40.0, 50.0] 

221 dst = np.zeros((3, 2, 2, 2)) 

222 iter_regrid = imod.prepare.regrid._make_regrid(first, ndim_regrid) 

223 

224 dst = imod.prepare.regrid._nd_regrid( 

225 src, dst, src_coords, dst_coords, iter_regrid, False 

226 ) 

227 compare = np.zeros((3, 2, 2, 2)) 

228 compare[..., :] = [10.0, 30.0] 

229 assert np.allclose(dst, compare) 

230 

231 

232def test_regrid_coord(): 

233 # Regular 

234 da = xr.DataArray((np.zeros(4)), {"x": np.arange(4.0) + 0.5}, ("x",)) 

235 regridx = imod.prepare.common._coord(da, "x") 

236 assert np.allclose(regridx, np.arange(5.0)) 

237 

238 # Negative x 

239 da = xr.DataArray((np.zeros(4)), {"x": np.arange(-4.0, 0.0, 1.0) + 0.5}, ("x",)) 

240 regridx = imod.prepare.common._coord(da, "x") 

241 assert np.allclose(regridx, np.arange(-4.0, 1.0, 1.0)) 

242 

243 # Non-equidistant, postive dx, negative dy 

244 nrow, ncol = 3, 4 

245 dx = np.array([0.9, 1.1, 0.8, 1.2]) 

246 dy = np.array([-1.3, -0.7, -1.0]) 

247 xmin, xmax = 0.0, 4.0 

248 ymin, ymax = 0.0, 3.0 

249 coords = imod.util.spatial._xycoords((xmin, xmax, ymin, ymax), (dx, dy)) 

250 # Flip around the y to be increasing 

251 # this is what happens in the regridding fuctions with the DataArrays 

252 coords["y"] = coords["y"][::-1] 

253 coords["dy"] = (coords["dy"][0], coords["dy"][1][::-1]) 

254 kwargs = {"name": "nonequidistant", "coords": coords, "dims": ("y", "x")} 

255 data = np.ones((nrow, ncol), dtype=np.float32) 

256 da = xr.DataArray(data, **kwargs) 

257 

258 regridx = imod.prepare.common._coord(da, "x") 

259 regridy = imod.prepare.common._coord(da, "y") 

260 assert float(regridx.min()) == xmin 

261 assert float(regridx.max()) == xmax 

262 assert float(regridy.min()) == ymin 

263 assert float(regridy.max()) == ymax 

264 assert np.allclose(np.diff(regridx), dx) 

265 assert np.allclose(np.diff(regridy), dy[::-1] * -1.0) 

266 

267 # Now test it if dy doesn't have the right sign 

268 # it should automatically infer it based on y instead. 

269 da["dy"].values *= -1.0 

270 regridy2 = imod.prepare.common._coord(da, "y") 

271 assert np.allclose(regridy, regridy2) 

272 

273 

274def test_regrid_mean1d(): 

275 values = np.array([1.0, 2.0, 3.0]) 

276 src_x = np.array([0.5, 1.5, 2.5]) 

277 dst_x = np.array([0.5, 2.0]) 

278 coords = {"x": src_x, "dx": ("x", np.array([1.0, 1.0, 1.0]))} 

279 like_coords = {"x": dst_x, "dx": ("x", np.array([1.0, 2.0]))} 

280 dims = ("x",) 

281 source = xr.DataArray(values, coords, dims) 

282 like = xr.DataArray(np.empty(2), like_coords, dims) 

283 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

284 compare = np.array([1.0, 2.5]) 

285 assert np.allclose(out.values, compare) 

286 

287 

288def test_regrid_mean1d__dx_negative(): 

289 values = np.array([1.0, 2.0, 3.0]) 

290 src_x = np.array([2.5, 1.5, 0.5]) 

291 dst_x = np.array([2.0, 0.5]) 

292 coords = {"x": src_x, "dx": ("x", np.array([-1.0, -1.0, -1.0]))} 

293 like_coords = {"x": dst_x, "dx": ("x", np.array([-2.0, -1.0]))} 

294 dims = ("x",) 

295 source = xr.DataArray(values, coords, dims) 

296 like = xr.DataArray(np.empty(2), like_coords, dims) 

297 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

298 compare = np.array([1.5, 3.0]) 

299 assert np.allclose(out.values, compare) 

300 

301 

302def test_regrid_mean1d__with_time(): 

303 values = np.full((4, 3), [1.0, 2.0, 3.0]) 

304 src_x = np.array([2.5, 1.5, 0.5]) 

305 dst_x = np.array([2.0, 0.5]) 

306 time = pd.date_range("2001-01-01", "2001-01-04") 

307 coords = {"time": time, "x": src_x, "dx": ("x", np.array([-1.0, -1.0, -1.0]))} 

308 like_coords = {"x": dst_x, "dx": ("x", np.array([-2.0, -1.0]))} 

309 source = xr.DataArray(values, coords, ["time", "x"]) 

310 like = xr.DataArray(np.empty(2), like_coords, ["x"]) 

311 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

312 assert isinstance(out, xr.DataArray) 

313 

314 like_coords = {"time": time, "x": dst_x, "dx": ("x", np.array([-2.0, -1.0]))} 

315 like = xr.DataArray(np.empty((4, 2)), like_coords, ["time", "x"]) 

316 with pytest.raises(RuntimeError): 

317 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

318 

319 

320@pytest.mark.parametrize("chunksize", [1, 2, 3]) 

321def test_regrid_mean1d__opposite_dx(chunksize): 

322 values = np.array([1.0, 2.0, 3.0]) 

323 src_x = np.array([2.5, 1.5, 0.5]) 

324 dst_x = np.array([0.5, 2.0]) 

325 coords = {"x": src_x, "dx": ("x", np.array([-1.0, -1.0, -1.0]))} 

326 like_coords = {"x": dst_x, "dx": ("x", np.array([1.0, 2.0]))} 

327 dims = ("x",) 

328 source = xr.DataArray(values, coords, dims) 

329 like = xr.DataArray(np.empty(2), like_coords, dims) 

330 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

331 compare = np.array([3.0, 1.5]) 

332 assert np.allclose(out.values, compare) 

333 

334 # Check sensitivity for dx, it shouldn't matter 

335 like_coords = {"x": dst_x, "dx": ("x", np.array([-1.0, -2.0]))} 

336 like = xr.DataArray(np.empty(2), like_coords, dims) 

337 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

338 compare = np.array([3.0, 1.5]) 

339 assert np.allclose(out.values, compare) 

340 

341 # Now with chunking 

342 source = source.chunk({"x": chunksize}) 

343 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

344 assert np.allclose(out.values, compare) 

345 

346 

347def test_regrid__coordinate_errors(): 

348 values = np.array([1.0, 2.0, 3.0]) 

349 # Non-equidistant, dx not provided 

350 src_x = np.array([0.5, 1.5, 3.5]) 

351 dst_x = np.array([0.5, 2.0]) 

352 coords = {"x": src_x} 

353 like_coords = {"x": dst_x} 

354 dims = ("x",) 

355 source = xr.DataArray(values, coords, dims) 

356 like = xr.DataArray(np.empty(2), like_coords, dims) 

357 with pytest.raises(ValueError): 

358 imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

359 

360 # Single cell src, dx not provided 

361 # dst stays the same 

362 values = np.array([1.0]) 

363 src_x = np.array([0.5]) 

364 coords = {"x": src_x} 

365 source = xr.DataArray(values, coords, dims) 

366 like = xr.DataArray(np.empty(2), like_coords, dims) 

367 with pytest.raises(ValueError): 

368 imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

369 

370 # Now for dst 

371 values = np.array([1.0, 2.0, 3.0]) 

372 src_x = np.array([2.5, 1.5, 0.5]) 

373 dst_x = np.array([0.5]) 

374 coords = {"x": src_x, "dx": ("x", np.array([-1.0, -1.0, -1.0]))} 

375 like_coords = {"x": dst_x} 

376 dims = ("x",) 

377 source = xr.DataArray(values, coords, dims) 

378 like = xr.DataArray(np.empty(1), like_coords, dims) 

379 with pytest.raises(ValueError): 

380 imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

381 

382 

383@pytest.mark.parametrize("chunksize", [1, 2, 3]) 

384def test_regrid_mean2d(chunksize): 

385 values = np.array([[0.6, 0.2, 3.4], [1.4, 1.6, 1.0], [4.0, 2.8, 3.0]]) 

386 src_x = np.arange(3.0) + 0.5 

387 coords = {"y": src_x, "x": src_x} 

388 dims = ("y", "x") 

389 source = xr.DataArray(values, coords, dims) 

390 dst_x = np.arange(0.0, 3.0, 1.5) + 0.75 

391 likecoords = {"y": dst_x, "x": dst_x} 

392 like = xr.DataArray(np.empty((2, 2)), likecoords, dims) 

393 

394 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

395 compare = np.array( 

396 [ 

397 [ 

398 (0.6 + 0.5 * 0.2 + 0.5 * 1.4 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

399 (3.4 + 0.5 * 0.2 + 0.5 * 1.0 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

400 ], 

401 [ 

402 (4.0 + 0.5 * 1.4 + 0.5 * 2.8 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

403 (3.0 + 0.5 * 1.0 + 0.5 * 2.8 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

404 ], 

405 ] 

406 ) 

407 assert np.allclose(out.values, compare) 

408 

409 # Now with chunking 

410 source = source.chunk({"x": chunksize, "y": chunksize}) 

411 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

412 assert np.allclose(out.values, compare) 

413 

414 # Now with orthogonal chunks 

415 source = source.chunk({"x": 2}) 

416 likecoords = {"y": dst_x} 

417 like = xr.DataArray(np.empty(2), likecoords, ["y"]) 

418 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

419 

420 

421@pytest.mark.parametrize("chunksize", [1, 2, 3]) 

422def test_regrid_mean2d_over3darray(chunksize): 

423 values = np.array([[0.6, 0.2, 3.4], [1.4, 1.6, 1.0], [4.0, 2.8, 3.0]]) 

424 values = np.stack([values for _ in range(5)]) 

425 src_x = np.arange(3.0) + 0.5 

426 src_z = np.arange(5.0) 

427 coords = {"z": src_z, "y": src_x, "x": src_x} 

428 dims = ("z", "y", "x") 

429 source = xr.DataArray(values, coords, dims) 

430 dst_x = np.arange(0.0, 3.0, 1.5) + 0.75 

431 likecoords = {"z": src_z, "y": dst_x, "x": dst_x} 

432 like = xr.DataArray(np.empty((5, 2, 2)), likecoords, dims) 

433 

434 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

435 compare_values = np.array( 

436 [ 

437 [ 

438 (0.6 + 0.5 * 0.2 + 0.5 * 1.4 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

439 (3.4 + 0.5 * 0.2 + 0.5 * 1.0 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

440 ], 

441 [ 

442 (4.0 + 0.5 * 1.4 + 0.5 * 2.8 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

443 (3.0 + 0.5 * 1.0 + 0.5 * 2.8 + 0.25 * 1.6) / (1.0 + 0.5 + 0.5 + 0.25), 

444 ], 

445 ] 

446 ) 

447 compare = np.empty((5, 2, 2)) 

448 compare[:, ...] = compare_values 

449 

450 assert np.allclose(out.values, compare) 

451 

452 # Now with chunking 

453 source = source.chunk({"x": chunksize, "y": chunksize}) 

454 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

455 assert np.allclose(out.values, compare) 

456 

457 

458def test_regrid_conductance2d(): 

459 # First case, same domain, smaller cellsizes 

460 y = np.arange(10.0, 0.0, -2.5) - 1.25 

461 x = np.arange(0.0, 10.0, 2.5) + 1.25 

462 coords = {"y": y, "x": x} 

463 dims = ("y", "x") 

464 like_da = xr.DataArray(np.empty((4, 4)), coords, dims) 

465 src_da = xr.DataArray( 

466 [[10.0, 10.0], [10.0, 10.0]], {"y": [7.5, 2.5], "x": [2.5, 7.5]}, dims 

467 ) 

468 

469 regridder = imod.prepare.Regridder(method=conductance, use_relative_weights=True) 

470 dst_da = regridder.regrid(src_da, like_da) 

471 assert float(src_da.sum()) == float(dst_da.sum()) 

472 

473 # Second case, different domain, smaller cellsizes 

474 dx = np.array([2.5, 2.5, 2.5, 3.5]) 

475 x = np.cumsum(dx) - dx * 0.5 

476 coords["x"] = x 

477 coords["dx"] = ("x", dx) 

478 like_da = xr.DataArray(np.empty((4, 4)), coords, dims) 

479 dst_da = regridder.regrid(src_da, like_da) 

480 assert float(src_da.sum()) == float(dst_da.sum()) 

481 

482 # Third case, same domain, small to large cellsizes 

483 y = np.arange(10.0, 0.0, -2.5) - 1.25 

484 x = np.arange(0.0, 10.0, 2.5) + 1.25 

485 coords = {"y": y, "x": x} 

486 dims = ("y", "x") 

487 src_da = xr.DataArray(np.full((4, 4), 10.0), coords, dims) 

488 like_da = xr.DataArray( 

489 [[10.0, 10.0], [10.0, 10.0]], {"y": [7.5, 2.5], "x": [2.5, 7.5]}, dims 

490 ) 

491 

492 dst_da = regridder.regrid(src_da, like_da) 

493 assert float(src_da.sum()) == float(dst_da.sum()) 

494 

495 # Fourth case, larger domain, small to large cellsizes 

496 like_da = xr.DataArray( 

497 [[10.0, 10.0], [10.0, 10.0]], {"y": [15.0, 5.0], "x": [5.0, 15.0]}, dims 

498 ) 

499 

500 dst_da = regridder.regrid(src_da, like_da) 

501 assert float(src_da.sum()) == float(dst_da.sum()) 

502 

503 

504def test_regrid_errors(): 

505 values = np.array([[0.6, 0.2, 3.4], [1.4, 1.6, 1.0], [4.0, 2.8, 3.0]]) 

506 values = np.stack([values for _ in range(5)]) 

507 src_x = np.arange(3.0) + 0.5 

508 src_z = np.arange(5.0) 

509 coords = {"z": src_z, "y": src_x, "x": src_x} 

510 dims = ("z", "y", "x") 

511 source = xr.DataArray(values, coords, dims) 

512 dst_x = np.arange(0.0, 3.0, 1.5) + 0.75 

513 dst_z = np.arange(0.0, 2.5, 0.5) 

514 likecoords = {"z": dst_z, "y": dst_x, "x": dst_x} 

515 like = xr.DataArray(np.empty((5, 2, 2)), likecoords, dims) 

516 

517 with pytest.raises(ValueError, match="The conductance method"): 

518 _ = imod.prepare.Regridder(method="conductance").regrid(source, like) 

519 

520 like = source.isel(z=slice(1, 3)) 

521 like["z"] = [1.5, 2.5] 

522 with pytest.raises(ValueError, match="Number of dimensions to regrid"): 

523 _ = imod.prepare.Regridder(method="mean", ndim_regrid=3).regrid(source, like) 

524 

525 

526def test_str_method(): 

527 values = np.array([1.0, 2.0, 3.0]) 

528 src_x = np.array([2.5, 1.5, 0.5]) 

529 dst_x = np.array([2.0, 0.5]) 

530 coords = {"x": src_x, "dx": ("x", np.array([-1.0, -1.0, -1.0]))} 

531 like_coords = {"x": dst_x, "dx": ("x", np.array([-2.0, -1.0]))} 

532 dims = ("x",) 

533 source = xr.DataArray(values, coords, dims) 

534 like = xr.DataArray(np.empty(2), like_coords, dims) 

535 # Test function method 

536 out = imod.prepare.Regridder(method=mean).regrid(source, like) 

537 compare = np.array([1.5, 3.0]) 

538 assert np.allclose(out.values, compare) 

539 

540 # Now test str method 

541 out = imod.prepare.Regridder(method="mean").regrid(source, like) 

542 assert np.allclose(out.values, compare) 

543 

544 out = imod.prepare.Regridder(method="nearest").regrid(source, like) 

545 

546 

547def test_no_overlap(): 

548 values = np.array([1.0, 2.0, 3.0]) 

549 src_x = np.array([0.5, 1.5, 2.5]) 

550 dst_x = np.array([4.5, 5.5]) 

551 coords = {"x": src_x, "dx": ("x", np.array([1.0, 1.0, 1.0]))} 

552 like_coords = {"x": dst_x, "dx": ("x", np.array([1.0, 1.0]))} 

553 dims = ("x",) 

554 source = xr.DataArray(values, coords, dims) 

555 like = xr.DataArray(np.empty(2), like_coords, dims) 

556 out = imod.prepare.Regridder(method=weightedmean).regrid(source, like) 

557 assert out.shape == (2,) 

558 assert out.isnull().all()