Coverage for test_util.py: 98%

1 import datetime

2 import pathlib

3 import re

4

5 import affine

6 import cftime

7 import dask

8 import numpy as np

9 import pandas as pd

10 import pytest

11 import xarray as xr

12 import xugrid as xu

13

14 import imod

15

16

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

18 def 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

54 def 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

89 def 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

117 def 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

131 def 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

146 def 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

160 def 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

174 def 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

188 def 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

202 def 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

216 def 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

230 def 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

242 def 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

258 def 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

274 def 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

288 def 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

300 def 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

313 def 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

327 def 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

340 def 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

354 def 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

367 def 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

376 def 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

386 def 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

395 def 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

430 def 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

444 def 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

479 def 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

499 def 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

522 def 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

535 def 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

641 def 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

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