Coverage for test_util.py: 98%

357 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-03-01 15:36 +0100

1import datetime 

2import pathlib 

3import re 

4 

5import affine 

6import cftime 

7import dask 

8import numpy as np 

9import pandas as pd 

10import pytest 

11import xarray as xr 

12import xugrid as xu 

13 

14import imod 

15 

16 

17@pytest.fixture(scope="function") 

18def ugrid_ds(): 

19 vertices = np.array( 

20 [ 

21 [0.0, 0.0], # 0 

22 [1.0, 0.0], # 1 

23 [2.0, 0.0], # 2 

24 [0.0, 1.0], # 3 

25 [1.0, 1.0], # 4 

26 [2.0, 1.0], # 5 

27 [1.0, 2.0], # 6 

28 ] 

29 ) 

30 faces = np.array( 

31 [ 

32 [0, 1, 4, 3], 

33 [1, 2, 5, 4], 

34 [3, 4, 6, -1], 

35 [4, 5, 6, -1], 

36 ] 

37 ) 

38 grid = xu.Ugrid2d( 

39 node_x=vertices[:, 0], 

40 node_y=vertices[:, 1], 

41 fill_value=-1, 

42 face_node_connectivity=faces, 

43 ) 

44 darray = xr.DataArray( 

45 data=np.random.rand(5, 3, grid.n_face), 

46 coords={"time": pd.date_range("2000-01-01", "2000-01-05"), "layer": [1, 2, 3]}, 

47 dims=["time", "layer", grid.face_dimension], 

48 ) 

49 ds = grid.to_dataset() 

50 ds["a"] = darray 

51 return ds 

52 

53 

54def test_compose(): 

55 d = { 

56 "name": "head", 

57 "directory": pathlib.Path("path", "to"), 

58 "extension": ".idf", 

59 "layer": 5, 

60 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

61 "species": 6, 

62 } 

63 path = imod.util.path.compose(d) 

64 targetpath = pathlib.Path(d["directory"], "head_c6_20180222090657_l5.idf") 

65 assert path == targetpath 

66 

67 d.pop("species") 

68 path = imod.util.path.compose(d) 

69 targetpath = pathlib.Path(d["directory"], "head_20180222090657_l5.idf") 

70 assert path == targetpath 

71 

72 d.pop("layer") 

73 path = imod.util.path.compose(d) 

74 targetpath = pathlib.Path(d["directory"], "head_20180222090657.idf") 

75 assert path == targetpath 

76 

77 d.pop("time") 

78 d["layer"] = 1 

79 path = imod.util.path.compose(d) 

80 targetpath = pathlib.Path(d["directory"], "head_l1.idf") 

81 assert path == targetpath 

82 

83 d["species"] = 6 

84 path = imod.util.path.compose(d) 

85 targetpath = pathlib.Path(d["directory"], "head_c6_l1.idf") 

86 assert path == targetpath 

87 

88 

89def test_compose__pattern(): 

90 d = { 

91 "name": "head", 

92 "directory": pathlib.Path("path", "to"), 

93 "extension": ".foo", 

94 "layer": 5, 

95 } 

96 targetpath = pathlib.Path(d["directory"], "head_2018-02-22_l05.foo") 

97 

98 d["time"] = datetime.datetime(2018, 2, 22, 9, 6, 57) 

99 path = imod.util.path.compose(d, pattern="{name}_{time:%Y-%m-%d}_l{layer:02d}{extension}") 

100 assert path == targetpath 

101 

102 d["time"] = cftime.DatetimeProlepticGregorian(2018, 2, 22, 9, 6, 57) 

103 path = imod.util.path.compose(d, pattern="{name}_{time:%Y-%m-%d}_l{layer:02d}{extension}") 

104 assert path == targetpath 

105 

106 d["time"] = np.datetime64("2018-02-22 09:06:57") 

107 path = imod.util.path.compose(d, pattern="{name}_{time:%Y-%m-%d}_l{layer:02d}{extension}") 

108 assert path == targetpath 

109 

110 targetpath = pathlib.Path(d["directory"], ".foo_makes_head_no_layer5_sense_day22") 

111 path = imod.util.path.compose( 

112 d, pattern="{extension}_makes_{name}_no_layer{layer:d}_sense_day{time:%d}" 

113 ) 

114 assert path == targetpath 

115 

116 

117def test_decompose(): 

118 d = imod.util.path.decompose("path/to/head_20180222090657_l5.idf") 

119 refd = { 

120 "extension": ".idf", 

121 "directory": pathlib.Path("path", "to"), 

122 "name": "head", 

123 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

124 "layer": 5, 

125 "dims": ["time", "layer"], 

126 } 

127 assert isinstance(d, dict) 

128 assert d == refd 

129 

130 

131def test_decompose_species(): 

132 d = imod.util.path.decompose("path/to/conc_c3_20180222090657_l5.idf") 

133 refd = { 

134 "extension": ".idf", 

135 "species": 3, 

136 "directory": pathlib.Path("path", "to"), 

137 "name": "conc", 

138 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

139 "layer": 5, 

140 "dims": ["species", "time", "layer"], 

141 } 

142 assert isinstance(d, dict) 

143 assert d == refd 

144 

145 

146def test_decompose_short_date(): 

147 d = imod.util.path.decompose("path/to/head_20180222_l5.idf") 

148 refd = { 

149 "extension": ".idf", 

150 "directory": pathlib.Path("path", "to"), 

151 "name": "head", 

152 "time": datetime.datetime(2018, 2, 22), 

153 "layer": 5, 

154 "dims": ["time", "layer"], 

155 } 

156 assert isinstance(d, dict) 

157 assert d == refd 

158 

159 

160def test_decompose_nonstandard_date(): 

161 d = imod.util.path.decompose("path/to/head_2018-02-22_l5.idf") 

162 refd = { 

163 "extension": ".idf", 

164 "directory": pathlib.Path("path", "to"), 

165 "name": "head", 

166 "time": datetime.datetime(2018, 2, 22), 

167 "layer": 5, 

168 "dims": ["time", "layer"], 

169 } 

170 assert isinstance(d, dict) 

171 assert d == refd 

172 

173 

174def test_decompose_only_year(): 

175 d = imod.util.path.decompose("path/to/head_2018_l5.idf", pattern="{name}_{time}_l{layer}") 

176 refd = { 

177 "extension": ".idf", 

178 "directory": pathlib.Path("path", "to"), 

179 "name": "head", 

180 "time": datetime.datetime(2018, 1, 1), 

181 "layer": 5, 

182 "dims": ["time", "layer"], 

183 } 

184 assert isinstance(d, dict) 

185 assert d == refd 

186 

187 

188def test_decompose_underscore(): 

189 d = imod.util.path.decompose("path/to/starting_head_20180222090657_l5.idf") 

190 refd = { 

191 "extension": ".idf", 

192 "directory": pathlib.Path("path", "to"), 

193 "name": "starting_head", 

194 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

195 "layer": 5, 

196 "dims": ["time", "layer"], 

197 } 

198 assert isinstance(d, dict) 

199 assert d == refd 

200 

201 

202def test_decompose_dash(): 

203 d = imod.util.path.decompose("path/to/starting-head_20180222090657_l5.idf") 

204 refd = { 

205 "extension": ".idf", 

206 "directory": pathlib.Path("path", "to"), 

207 "name": "starting-head", 

208 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

209 "layer": 5, 

210 "dims": ["time", "layer"], 

211 } 

212 assert isinstance(d, dict) 

213 assert d == refd 

214 

215 

216def test_decompose_steady_state(): 

217 d = imod.util.path.decompose("path/to/head_steady-state_l64.idf") 

218 refd = { 

219 "extension": ".idf", 

220 "directory": pathlib.Path("path", "to"), 

221 "name": "head", 

222 "time": "steady-state", 

223 "layer": 64, 

224 "dims": ["layer", "time"], 

225 } 

226 assert isinstance(d, dict) 

227 assert d == refd 

228 

229 

230def test_decompose_underscore_in_name(): 

231 d = imod.util.path.decompose("path/to/some_name.idf") 

232 refd = { 

233 "extension": ".idf", 

234 "directory": pathlib.Path("path", "to"), 

235 "name": "some_name", 

236 "dims": [], 

237 } 

238 assert isinstance(d, dict) 

239 assert d == refd 

240 

241 

242def test_decompose_pattern_underscore(): 

243 d = imod.util.path.decompose( 

244 "path/to/starting_head_20180222090657_l5.idf", pattern="{name}_{time}_l{layer}" 

245 ) 

246 refd = { 

247 "extension": ".idf", 

248 "directory": pathlib.Path("path", "to"), 

249 "name": "starting_head", 

250 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

251 "layer": 5, 

252 "dims": ["time", "layer"], 

253 } 

254 assert isinstance(d, dict) 

255 assert d == refd 

256 

257 

258def test_decompose_pattern_dash(): 

259 d = imod.util.path.decompose( 

260 "path/to/starting-head_20180222090657_l5.idf", pattern="{name}_{time}_l{layer}" 

261 ) 

262 refd = { 

263 "extension": ".idf", 

264 "directory": pathlib.Path("path", "to"), 

265 "name": "starting-head", 

266 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

267 "layer": 5, 

268 "dims": ["time", "layer"], 

269 } 

270 assert isinstance(d, dict) 

271 assert d == refd 

272 

273 

274def test_decompose_regexpattern(): 

275 pattern = re.compile(r"(?P<name>[\w]+)L(?P<layer>[\d+]*)", re.IGNORECASE) 

276 d = imod.util.path.decompose("headL11.idf", pattern=pattern) 

277 refd = { 

278 "extension": ".idf", 

279 "directory": pathlib.Path("."), 

280 "name": "head", 

281 "layer": 11, 

282 "dims": ["layer"], 

283 } 

284 assert isinstance(d, dict) 

285 assert d == refd 

286 

287 

288def test_decompose_nodate(): 

289 d = imod.util.path.decompose("dem_10m.idf") 

290 refd = { 

291 "extension": ".idf", 

292 "directory": pathlib.Path("."), 

293 "name": "dem_10m", 

294 "dims": [], 

295 } 

296 assert isinstance(d, dict) 

297 assert d == refd 

298 

299 

300def test_decompose_dateonly(): 

301 d = imod.util.path.decompose("20180222090657.idf", pattern="{time}") 

302 refd = { 

303 "extension": ".idf", 

304 "directory": pathlib.Path("."), 

305 "name": "20180222090657", 

306 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

307 "dims": ["time"], 

308 } 

309 assert isinstance(d, dict) 

310 assert d == refd 

311 

312 

313def test_decompose_datelayeronly(): 

314 d = imod.util.path.decompose("20180222090657_l7.idf", pattern="{time}_l{layer}") 

315 refd = { 

316 "extension": ".idf", 

317 "directory": pathlib.Path("."), 

318 "name": "20180222090657_7", 

319 "time": datetime.datetime(2018, 2, 22, 9, 6, 57), 

320 "layer": 7, 

321 "dims": ["time", "layer"], 

322 } 

323 assert isinstance(d, dict) 

324 assert d == refd 

325 

326 

327def test_decompose_z_float(): 

328 d = imod.util.path.decompose("test_0.25.idf", pattern="{name}_{z}") 

329 refd = { 

330 "extension": ".idf", 

331 "directory": pathlib.Path("."), 

332 "name": "test", 

333 "z": "0.25", 

334 "dims": ["z"], 

335 } 

336 assert isinstance(d, dict) 

337 assert d == refd 

338 

339 

340def test_compose_year9999(): 

341 d = { 

342 "name": "head", 

343 "directory": pathlib.Path("path", "to"), 

344 "extension": ".idf", 

345 "layer": 5, 

346 "time": datetime.datetime(9999, 2, 22, 9, 6, 57), 

347 "dims": ["time"], 

348 } 

349 path = imod.util.path.compose(d) 

350 targetpath = pathlib.Path(d["directory"], "head_99990222090657_l5.idf") 

351 assert path == targetpath 

352 

353 

354def test_decompose_dateonly_year9999(): 

355 d = imod.util.path.decompose("99990222090657.idf", pattern="{time}") 

356 refd = { 

357 "extension": ".idf", 

358 "directory": pathlib.Path("."), 

359 "name": "99990222090657", 

360 "time": datetime.datetime(9999, 2, 22, 9, 6, 57), 

361 "dims": ["time"], 

362 } 

363 assert isinstance(d, dict) 

364 assert d == refd 

365 

366 

367def test_datetime_conversion__withinbounds(): 

368 times = [datetime.datetime(y, 1, 1) for y in range(2000, 2010)] 

369 converted, use_cftime = imod.util.time._convert_datetimes(times, use_cftime=False) 

370 assert use_cftime is False 

371 assert all(t.dtype == "<M8[ns]" for t in converted) 

372 assert converted[0] == np.datetime64("2000-01-01", "ns") 

373 assert converted[-1] == np.datetime64("2009-01-01", "ns") 

374 

375 

376def test_datetime_conversion__outofbounds(): 

377 times = [datetime.datetime(y, 1, 1) for y in range(1670, 1680)] 

378 with pytest.warns(UserWarning): 

379 converted, use_cftime = imod.util.time._convert_datetimes(times, use_cftime=False) 

380 assert use_cftime is True 

381 assert all(type(t) is cftime.DatetimeProlepticGregorian for t in converted) 

382 assert converted[0] == cftime.DatetimeProlepticGregorian(1670, 1, 1) 

383 assert converted[-1] == cftime.DatetimeProlepticGregorian(1679, 1, 1) 

384 

385 

386def test_datetime_conversion__withinbounds_cftime(): 

387 times = [datetime.datetime(y, 1, 1) for y in range(2000, 2010)] 

388 converted, use_cftime = imod.util.time._convert_datetimes(times, use_cftime=True) 

389 assert use_cftime is True 

390 assert all(type(t) is cftime.DatetimeProlepticGregorian for t in converted) 

391 assert converted[0] == cftime.DatetimeProlepticGregorian(2000, 1, 1) 

392 assert converted[-1] == cftime.DatetimeProlepticGregorian(2009, 1, 1) 

393 

394 

395def test_transform(): 

396 # implicit dx dy 

397 data = np.ones((2, 3)) 

398 coords = {"x": [0.5, 1.5, 2.5], "y": [1.5, 0.5]} 

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

400 da = xr.DataArray(data, coords, dims) 

401 actual = imod.util.spatial.transform(da) 

402 expected = affine.Affine(1.0, 0.0, 0.0, 0.0, -1.0, 2.0) 

403 assert actual == expected 

404 

405 # explicit dx dy, equidistant 

406 coords = { 

407 "x": [0.5, 1.5, 2.5], 

408 "y": [1.5, 0.5], 

409 "dx": ("x", [1.0, 1.0, 1.0]), 

410 "dy": ("y", [-1.0, -1.0]), 

411 } 

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

413 da = xr.DataArray(data, coords, dims) 

414 actual = imod.util.spatial.transform(da) 

415 assert actual == expected 

416 

417 # explicit dx dy, non-equidistant 

418 coords = { 

419 "x": [0.5, 1.5, 3.5], 

420 "y": [1.5, 0.5], 

421 "dx": ("x", [1.0, 1.0, 2.0]), 

422 "dy": ("y", [-1.0, -1.0]), 

423 } 

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

425 da = xr.DataArray(data, coords, dims) 

426 with pytest.raises(ValueError): 

427 imod.util.spatial.transform(da) 

428 

429 

430def test_is_divisor(): 

431 a = np.array([1.0, 0.5, 0.1]) 

432 b = 0.05 

433 assert imod.util.spatial.is_divisor(a, b) 

434 assert imod.util.spatial.is_divisor(-a, b) 

435 assert imod.util.spatial.is_divisor(a, -b) 

436 assert imod.util.spatial.is_divisor(-a, -b) 

437 b = 0.07 

438 assert not imod.util.spatial.is_divisor(a, b) 

439 assert not imod.util.spatial.is_divisor(-a, b) 

440 assert not imod.util.spatial.is_divisor(a, -b) 

441 assert not imod.util.spatial.is_divisor(-a, -b) 

442 

443 

444def test_empty(): 

445 da = imod.util.spatial.empty_2d(1.0, 0.0, 2.0, -1.0, 10.0, 12.0) 

446 assert da.isnull().all() 

447 assert np.allclose(da["x"], [0.5, 1.5]) 

448 assert np.allclose(da["y"], [11.5, 10.5]) 

449 assert da.dims == ("y", "x") 

450 # Sign on dx, dy should be ignored 

451 da = imod.util.spatial.empty_2d(-1.0, 0.0, 2.0, 1.0, 10.0, 12.0) 

452 assert np.allclose(da["x"], [0.5, 1.5]) 

453 assert np.allclose(da["y"], [11.5, 10.5]) 

454 

455 with pytest.raises(ValueError, match="layer must be 1d"): 

456 imod.util.spatial.empty_3d(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, [[1, 2]]) 

457 

458 da3d = imod.util.spatial.empty_3d(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, 1) 

459 assert da3d.ndim == 3 

460 da3d = imod.util.spatial.empty_3d(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, [1, 3]) 

461 assert np.array_equal(da3d["layer"], [1, 3]) 

462 assert da3d.dims == ("layer", "y", "x") 

463 

464 times = ["2000-01-01", "2001-01-01"] 

465 with pytest.raises(ValueError, match="time must be 1d"): 

466 imod.util.spatial.empty_2d_transient(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, [times]) 

467 

468 da2dt = imod.util.spatial.empty_2d_transient(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, times[0]) 

469 assert da2dt.ndim == 3 

470 da2dt = imod.util.spatial.empty_2d_transient(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, times) 

471 assert isinstance(da2dt["time"].values[0], np.datetime64) 

472 assert da2dt.dims == ("time", "y", "x") 

473 

474 da3dt = imod.util.spatial.empty_3d_transient(1.0, 0.0, 2.0, -1.0, 10.0, 12.0, [0, 1], times) 

475 assert da3dt.ndim == 4 

476 assert da3dt.dims == ("time", "layer", "y", "x") 

477 

478 

479def test_where(): 

480 a = xr.DataArray( 

481 [[0.0, 1.0], [2.0, np.nan]], 

482 {"y": [1.5, 0.5], "x": [0.5, 1.5]}, 

483 ["y", "x"], 

484 ) 

485 cond = a <= 1 

486 actual = imod.util.structured.where(cond, if_true=a, if_false=1.0) 

487 assert np.allclose(actual.values, [[0.0, 1.0], [1.0, np.nan]], equal_nan=True) 

488 

489 actual = imod.util.structured.where(cond, if_true=0.0, if_false=1.0) 

490 assert np.allclose(actual.values, [[0.0, 0.0], [1.0, 1.0]], equal_nan=True) 

491 

492 actual = imod.util.structured.where(cond, if_true=a, if_false=1.0, keep_nan=False) 

493 assert np.allclose(actual.values, [[0.0, 1.0], [1.0, 1.0]]) 

494 

495 with pytest.raises(ValueError, match="at least one of"): 

496 imod.util.structured.where(False, 1, 0) 

497 

498 

499def test_compliant_ugrid2d(ugrid_ds, write=False): 

500 uds = imod.util.spatial.mdal_compliant_ugrid2d(ugrid_ds) 

501 

502 assert isinstance(uds, xr.Dataset) 

503 for i in range(1, 4): 

504 assert f"a_layer_{i}" in uds 

505 

506 assert "mesh2d" in uds 

507 assert "mesh2d_face_nodes" in uds 

508 assert "mesh2d_node_x" in uds 

509 assert "mesh2d_node_y" in uds 

510 assert "mesh2d_nFaces" in uds.dims 

511 assert "mesh2d_nNodes" in uds.dims 

512 assert "mesh2d_nMax_face_nodes" in uds.dims 

513 attrs = uds["mesh2d"].attrs 

514 assert "face_coordinates" not in attrs 

515 

516 assert uds["time"].encoding["dtype"] == np.float64 

517 

518 if write: 

519 uds.to_netcdf("ugrid-mixed.nc") 

520 

521 

522def test_mdal_compliant_roundtrip(ugrid_ds): 

523 uds = xu.UgridDataset(imod.util.spatial.mdal_compliant_ugrid2d(ugrid_ds)) 

524 uds["b"] = (("time", "layer"), np.random.rand(5, 3)) 

525 uds["c"] = (("layer", "mesh2d_nFaces"), np.random.rand(3, 4)) 

526 back = imod.util.spatial.from_mdal_compliant_ugrid2d(uds) 

527 

528 assert isinstance(back, xu.UgridDataset) 

529 assert back["a"].dims == ("time", "layer", "mesh2d_nFaces") 

530 assert back["b"].dims == ("time", "layer") 

531 assert back["c"].dims == ("layer", "mesh2d_nFaces") 

532 assert np.array_equal(back["layer"], [1, 2, 3]) 

533 

534 

535def test_to_ugrid2d(write=False): 

536 a2d = imod.util.spatial.empty_2d( 

537 dx=1.0, 

538 xmin=0.0, 

539 xmax=2.0, 

540 dy=1.0, 

541 ymin=0.0, 

542 ymax=2.0, 

543 ) 

544 

545 with pytest.raises(TypeError, match="data must be xarray"): 

546 imod.util.spatial.to_ugrid2d(a2d.values) 

547 with pytest.raises(ValueError, match="A name is required"): 

548 imod.util.spatial.to_ugrid2d(a2d) 

549 

550 a2d.name = "a" 

551 with pytest.raises(ValueError, match="Last two dimensions of da"): 

552 imod.util.spatial.to_ugrid2d(a2d.transpose()) 

553 

554 uds = imod.util.spatial.to_ugrid2d(a2d) 

555 assert isinstance(uds, xr.Dataset) 

556 assert "a" in uds 

557 

558 assert "mesh2d" in uds 

559 assert "mesh2d_face_nodes" in uds 

560 assert "mesh2d_node_x" in uds 

561 assert "mesh2d_node_y" in uds 

562 assert "mesh2d_nFaces" in uds.dims 

563 assert "mesh2d_nNodes" in uds.dims 

564 assert "mesh2d_nMax_face_nodes" in uds.dims 

565 attrs = uds["mesh2d"].attrs 

566 assert "face_coordinates" not in attrs 

567 

568 if write: 

569 uds.to_netcdf("ugrid-a2d.nc") 

570 

571 # 2d Dataset 

572 ds = xr.Dataset() 

573 ds["a"] = a2d 

574 ds["b"] = a2d.copy() 

575 uds = imod.util.spatial.to_ugrid2d(ds) 

576 assert "a" in uds 

577 assert "b" in uds 

578 assert isinstance(uds, xr.Dataset) 

579 

580 if write: 

581 uds.to_netcdf("ugrid-a2d-ds.nc") 

582 

583 # transient 2d 

584 a2dt = imod.util.spatial.empty_2d_transient( 

585 dx=1.0, 

586 xmin=0.0, 

587 xmax=2.0, 

588 dy=1.0, 

589 ymin=0.0, 

590 ymax=2.0, 

591 time=pd.date_range("2000-01-01", "2000-01-05"), 

592 ) 

593 a2dt.name = "a" 

594 uds = imod.util.spatial.to_ugrid2d(a2dt) 

595 assert "a" in uds 

596 assert uds["time"].encoding["dtype"] == np.float64 

597 

598 if write: 

599 uds.to_netcdf("ugrid-a2dt.nc") 

600 

601 # 3d 

602 a3d = imod.util.spatial.empty_3d( 

603 dx=1.0, 

604 xmin=0.0, 

605 xmax=2.0, 

606 dy=1.0, 

607 ymin=0.0, 

608 ymax=2.0, 

609 layer=[1, 2, 3], 

610 ) 

611 a3d.name = "a" 

612 uds = imod.util.spatial.to_ugrid2d(a3d) 

613 assert isinstance(uds, xr.Dataset) 

614 for i in range(1, 4): 

615 assert f"a_layer_{i}" in uds 

616 

617 if write: 

618 uds.to_netcdf("ugrid-a3d.nc") 

619 

620 # transient 3d 

621 a3dt = imod.util.spatial.empty_3d_transient( 

622 dx=1.0, 

623 xmin=0.0, 

624 xmax=2.0, 

625 dy=1.0, 

626 ymin=0.0, 

627 ymax=2.0, 

628 layer=[1, 2, 3], 

629 time=pd.date_range("2000-01-01", "2000-01-05"), 

630 ) 

631 a3dt.name = "a" 

632 uds = imod.util.spatial.to_ugrid2d(a3dt) 

633 for i in range(1, 4): 

634 assert f"a_layer_{i}" in uds 

635 assert uds["time"].encoding["dtype"] == np.float64 

636 

637 if write: 

638 uds.to_netcdf("ugrid-a3dt.nc") 

639 

640 

641def test_replace(): 

642 # replace scalar 

643 da = xr.DataArray([0, 1, 2]) 

644 out = imod.util.structured.replace(da, 1, 10) 

645 assert out.equals(xr.DataArray([0, 10, 2])) 

646 

647 # Replace NaN by scalar 

648 da = xr.DataArray([np.nan, 1.0, 2.0]) 

649 out = imod.util.structured.replace(da, np.nan, 10.0) 

650 assert out.equals(xr.DataArray([10.0, 1.0, 2.0])) 

651 

652 # replace two 

653 da = xr.DataArray([0, 1, 2]) 

654 out = imod.util.structured.replace(da, [1, 2], [10, 20]) 

655 assert out.equals(xr.DataArray([0, 10, 20])) 

656 

657 # With a NaN in the data 

658 da = xr.DataArray([np.nan, 1.0, 2.0]) 

659 out = imod.util.structured.replace(da, [1, 2], [10, 20]) 

660 assert out.equals(xr.DataArray([np.nan, 10.0, 20.0])) 

661 

662 # Replace a NaN value 

663 da = xr.DataArray([np.nan, 1.0, 2.0]) 

664 out = imod.util.structured.replace(da, [np.nan, 2], [10, 20]) 

665 assert out.equals(xr.DataArray([10.0, 1.0, 20.0])) 

666 

667 # With non-present values in to_replace 

668 da = xr.DataArray([np.nan, 1.0, 1.0, 2.0]) 

669 out = imod.util.structured.replace(da, [1.0, 2.0, 30.0], [10.0, 20.0, 30.0]) 

670 assert out.equals(xr.DataArray([np.nan, 10.0, 10.0, 20.0])) 

671 

672 # With a nan and non-present values 

673 da = xr.DataArray([np.nan, 1.0, 1.0, 2.0]) 

674 out = imod.util.structured.replace(da, [np.nan, 1.0, 2.0, 30.0], 10.0) 

675 assert out.equals(xr.DataArray([10.0, 10.0, 10.0, 10.0])) 

676 

677 # With a dask array 

678 da = xr.DataArray(dask.array.full(3, 1.0)) 

679 out = imod.util.structured.replace(da, [1.0, 2.0], [10.0, 20.0]) 

680 assert isinstance(out.data, dask.array.Array) 

681 assert out.equals(xr.DataArray([10.0, 10.0, 10.0])) 

682 

683 # scalar to_replace, non-scalar value 

684 with pytest.raises(TypeError): 

685 imod.util.structured.replace(da, 1.0, [10.0, 20.0]) 

686 

687 # 2D arrays 

688 with pytest.raises(ValueError): 

689 imod.util.structured.replace(da, [[1.0, 2.0]], [[10.0, 20.0]]) 

690 

691 # 1D to_replace, 2D value 

692 with pytest.raises(ValueError): 

693 imod.util.structured.replace(da, [1.0, 2.0], [[10.0, 20.0]]) 

694 

695 # 1D, different size 

696 with pytest.raises(ValueError): 

697 imod.util.structured.replace(da, [1.0, 2.0], [10.0, 20.0, 30.0]) 

698 

699 

700def test_public_api(): 

701 """ 

702 Test if functions previously in imod.util.py are still available under same 

703 namespace 

704 """ 

705 from imod.util import ( 

706 cd, # noqa: F401 

707 empty_2d, # noqa: F401 

708 empty_2d_transient, # noqa: F401 

709 empty_3d, # noqa: F401 

710 empty_3d_transient, # noqa: F401 

711 from_mdal_compliant_ugrid2d, # noqa: F401 

712 ignore_warnings, # noqa: F401 

713 mdal_compliant_ugrid2d, # noqa: F401 

714 replace, # noqa: F401 

715 round_extent, # noqa: F401 

716 spatial_reference, # noqa: F401 

717 temporary_directory, # noqa: F401 

718 to_datetime, # noqa: F401 

719 to_ugrid2d, # noqa: F401 

720 transform, # noqa: F401 

721 ugrid2d_data, # noqa: F401 

722 values_within_range, # noqa: F401 

723 where, # noqa: F401 

724 )