Coverage for test_spatial.py: 100%

221 statements  

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

1import pathlib 

2 

3import geopandas as gpd 

4import numpy as np 

5import pandas as pd 

6import pytest 

7import shapely.geometry as sg 

8import xarray as xr 

9from pytest_cases import parametrize_with_cases 

10 

11import imod 

12from imod.testing import assert_frame_equal 

13 

14 

15@pytest.fixture(scope="module") 

16def geom(): 

17 return sg.Polygon([(0.0, 0.0), (1.1, 0.0), (1.1, 1.1), (0.0, 1.1)]) 

18 

19 

20class ShapefileCases: 

21 def case_integer(self, tmp_path_factory, geom): 

22 tmp_dir = tmp_path_factory.mktemp("testvector") 

23 gdf = gpd.GeoDataFrame() 

24 gdf.geometry = [geom] 

25 value = 2 

26 gdf["values"] = value 

27 gdf.to_file(tmp_dir / "shape_int.shp") 

28 return tmp_dir / "shape_int.shp", value 

29 

30 def case_float(self, tmp_path_factory, geom): 

31 tmp_dir = tmp_path_factory.mktemp("testvector") 

32 gdf = gpd.GeoDataFrame() 

33 gdf.geometry = [geom] 

34 value = 2.2 

35 gdf["values"] = value 

36 gdf.to_file(tmp_dir / "shape_int.shp") 

37 return tmp_dir / "shape_int.shp", value 

38 

39 

40def test_round_extent(): 

41 extent = (2.5, 2.5, 52.5, 52.5) 

42 cellsize = 5.0 

43 expected = (0.0, 0.0, 55.0, 55.0) 

44 assert imod.prepare.spatial.round_extent(extent, cellsize) == expected 

45 # Test if previous namespace available. 

46 assert imod.util.round_extent(extent, cellsize) == expected 

47 

48 

49def test_fill(): 

50 da = xr.DataArray([np.nan, 1.0, 2.0], {"x": [1, 2, 3]}, ("x",)) 

51 expected = xr.DataArray([1.0, 1.0, 2.0], {"x": [1, 2, 3]}, ("x",)) 

52 actual = imod.prepare.spatial.fill(da) 

53 assert actual.identical(expected) 

54 

55 # by dimension 

56 da = xr.DataArray( 

57 [[np.nan, np.nan, 1.0], [2.0, 2.0, 2.0]], 

58 {"x": [1, 2, 3], "y": [1, 2]}, 

59 ("y", "x"), 

60 ) 

61 expected_y = xr.DataArray( 

62 [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]], {"x": [1, 2, 3], "y": [1, 2]}, ("y", "x") 

63 ) 

64 expected_x = xr.DataArray( 

65 [[2.0, 2.0, 1.0], [2.0, 2.0, 2.0]], {"x": [1, 2, 3], "y": [1, 2]}, ("y", "x") 

66 ) 

67 actual_x = imod.prepare.spatial.fill(da, by="x") 

68 actual_y = imod.prepare.spatial.fill(da, by="y") 

69 assert actual_x.identical(expected_x) 

70 assert actual_y.identical(expected_y) 

71 

72 

73def test_laplace_interpolate(): 

74 nrow, ncol = 3, 4 

75 dx, dy = 1.0, -1.0 

76 xmin, xmax = 0.0, 4.0 

77 ymin, ymax = 0.0, 3.0 

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

79 kwargs = {"name": "test", "coords": coords, "dims": ("y", "x")} 

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

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

82 # remove some values 

83 da.values[:, 1] = np.nan 

84 actual = imod.prepare.laplace_interpolate(da, mxiter=5, iter1=30) 

85 assert np.allclose(actual.values, 1.0) 

86 

87 ibound = xr.full_like(da, True) 

88 ibound.values[1, 1] = False 

89 actual = imod.prepare.laplace_interpolate(da, ibound=ibound, mxiter=5, iter1=30) 

90 assert np.isnan(actual.values[1, 1]) 

91 actual.values[1, 1] = 1.0 

92 assert np.allclose(actual.values, 1.0) 

93 

94 da.values[:] = 1.0 

95 da.values[1, :] = np.nan 

96 actual = imod.prepare.laplace_interpolate(da, ibound=ibound, mxiter=5, iter1=30) 

97 assert np.isnan(actual.values[1, 1]) 

98 actual.values[1, 1] = 1.0 

99 assert np.allclose(actual.values, 1.0) 

100 

101 

102def test_rasterize(): 

103 geom = sg.Polygon([(0.0, 0.0), (1.1, 0.0), (1.1, 1.1), (0.0, 1.1)]) 

104 gdf = gpd.GeoDataFrame() 

105 gdf.geometry = [geom] 

106 gdf["values"] = [2.0] 

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

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

109 like = xr.DataArray(np.full((2, 2), np.nan), coords, dims) 

110 # No column given, default to 1.0 where polygon 

111 expected = xr.DataArray([[np.nan, np.nan], [1.0, np.nan]], coords, dims) 

112 actual = imod.prepare.spatial.rasterize(gdf, like) 

113 assert actual.identical(expected) 

114 

115 # Column given, use actual value 

116 expected = xr.DataArray([[np.nan, np.nan], [2.0, np.nan]], coords, dims) 

117 actual = imod.prepare.spatial.rasterize(gdf, like, column="values") 

118 assert actual.identical(expected) 

119 

120 

121def test_polygonize(): 

122 nrow, ncol = 2, 2 

123 dx, dy = 1.0, -1.0 

124 xmin, xmax = 0.0, 2.0 

125 ymin, ymax = 0.0, 2.0 

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

127 kwargs = {"name": "test", "coords": coords, "dims": ("y", "x")} 

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

129 data[0, 1] = 2.0 

130 data[1, 1] = 3.0 

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

132 gdf = imod.prepare.polygonize(da) 

133 assert len(gdf) == 3 

134 assert sorted(gdf["value"]) == [1.0, 2.0, 3.0] 

135 

136 

137def test_handle_dtype(): 

138 assert imod.prepare.spatial._handle_dtype(np.uint8, None) == (1, 0) 

139 assert imod.prepare.spatial._handle_dtype(np.uint16, None) == (2, 0) 

140 assert imod.prepare.spatial._handle_dtype(np.int16, None) == (3, -32768) 

141 assert imod.prepare.spatial._handle_dtype(np.uint32, None) == (4, 0) 

142 assert imod.prepare.spatial._handle_dtype(np.int32, None) == (5, -2147483648) 

143 assert imod.prepare.spatial._handle_dtype(np.float32, None) == (6, np.nan) 

144 assert imod.prepare.spatial._handle_dtype(np.float64, None) == (7, np.nan) 

145 

146 with pytest.raises(ValueError): # out of bounds 

147 imod.prepare.spatial._handle_dtype(np.uint32, -1) 

148 with pytest.raises(ValueError): # invalid dtype 

149 imod.prepare.spatial._handle_dtype(np.int64, -1) 

150 

151 

152@parametrize_with_cases("shapefile, expected_value", cases=ShapefileCases) 

153def test_gdal_rasterize(shapefile, expected_value): 

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

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

156 like = xr.DataArray(np.full((2, 2), np.nan), coords, dims) 

157 spatial_reference = {"bounds": (0.0, 2.0, 0.0, 2.0), "cellsizes": (1.0, -1.0)} 

158 expected = xr.DataArray([[np.nan, np.nan], [expected_value, np.nan]], coords, dims) 

159 

160 # Test with like 

161 actual = imod.prepare.spatial.gdal_rasterize(shapefile, "values", like) 

162 assert actual.identical(expected) 

163 

164 # Test with all_touched=True 

165 actual = imod.prepare.spatial.gdal_rasterize( 

166 shapefile, "values", like, all_touched=True 

167 ) 

168 assert (actual == expected_value).all() 

169 

170 # Test whether GDAL error results in a RuntimeError 

171 with pytest.raises(RuntimeError): # misnamed column 

172 imod.prepare.spatial.gdal_rasterize(shapefile, "value", like) 

173 

174 # Can't determine dtype without like, raise ValueError 

175 with pytest.raises(ValueError): 

176 imod.prepare.spatial.gdal_rasterize( 

177 shapefile, "values", spatial_reference=spatial_reference 

178 ) 

179 

180 # Test without like 

181 actual = imod.prepare.spatial.gdal_rasterize( 

182 shapefile, "values", dtype=np.float64, spatial_reference=spatial_reference 

183 ) 

184 coords = {"y": [1.5, 0.5], "x": [0.5, 1.5], "dx": 1.0, "dy": -1.0} 

185 expected = xr.DataArray([[np.nan, np.nan], [expected_value, np.nan]], coords, dims) 

186 assert actual.identical(expected) 

187 

188 # Test integer dtype, and nodata default (0 for uint8) 

189 expected = xr.DataArray([[0, 0], [2, 0]], coords, dims) 

190 actual = imod.prepare.spatial.gdal_rasterize( 

191 shapefile, "values", dtype=np.uint8, spatial_reference=spatial_reference 

192 ) 

193 assert actual.identical(expected) 

194 

195 # test with pathlib 

196 expected = xr.DataArray([[0, 0], [2, 0]], coords, dims) 

197 actual = imod.prepare.spatial.gdal_rasterize( 

198 pathlib.Path(shapefile), 

199 "values", 

200 dtype=np.uint8, 

201 spatial_reference=spatial_reference, 

202 ) 

203 assert actual.identical(expected) 

204 

205 

206@parametrize_with_cases("shapefile, expected_value", cases=ShapefileCases) 

207def test_private_celltable(shapefile, expected_value): 

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

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

210 like = xr.DataArray(np.full((2, 2), np.nan), coords, dims) 

211 

212 expected = pd.DataFrame() 

213 expected["row_index"] = [1] 

214 expected["col_index"] = [0] 

215 expected["values"] = [expected_value] 

216 expected["area"] = [1.0] 

217 

218 actual = imod.prepare.spatial._celltable( 

219 shapefile, "values", 1.0, like, type(expected_value) 

220 ) 

221 assert_frame_equal(actual, expected, check_dtype=True) 

222 

223 

224@parametrize_with_cases("shapefile, expected_value", cases=ShapefileCases) 

225def test_celltable(shapefile, expected_value): 

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

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

228 like = xr.DataArray(np.full((2, 2), np.nan), coords, dims) 

229 

230 expected = pd.DataFrame() 

231 expected["row_index"] = [1] 

232 expected["col_index"] = [0] 

233 expected["values"] = [expected_value] 

234 expected["area"] = [1.0] 

235 

236 actual = imod.prepare.spatial.celltable( 

237 shapefile, "values", 1.0, like, dtype=type(expected_value) 

238 ) 

239 assert_frame_equal(actual, expected, check_dtype=True) 

240 

241 # test resolution error: 

242 with pytest.raises(ValueError): 

243 actual = imod.prepare.spatial.celltable(shapefile, "values", 0.17, like) 

244 

245 

246def test_rasterize_table(): 

247 table = pd.DataFrame() 

248 table["row_index"] = [1] 

249 table["col_index"] = [0] 

250 table["values"] = [2] 

251 table["area"] = [1.0] 

252 

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

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

255 like = xr.DataArray(np.full((2, 2), np.nan), coords, dims) 

256 expected = xr.DataArray([[np.nan, np.nan], [1.0, np.nan]], coords, dims) 

257 

258 actual = imod.prepare.spatial.rasterize_celltable(table, "area", like) 

259 assert actual.identical(expected) 

260 

261 

262def test_zonal_aggregate_raster(tmp_path): 

263 raster = xr.DataArray( 

264 data=[[0.0, 0.0], [1.0, 1.0]], 

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

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

267 name="my-raster", 

268 ) 

269 geometries = [ 

270 sg.Polygon( 

271 [ 

272 (0.0, 1.0), 

273 (2.0, 1.0), 

274 (2.0, 2.0), 

275 (0.0, 2.0), 

276 ] 

277 ), 

278 sg.Polygon( 

279 [ 

280 (0.0, 0.0), 

281 (2.0, 0.0), 

282 (2.0, 1.0), 

283 (0.0, 1.0), 

284 ] 

285 ), 

286 ] 

287 gdf = gpd.GeoDataFrame(geometry=geometries) 

288 gdf["id"] = [1, 2] 

289 path = tmp_path / "two-zones1.shp" 

290 gdf.to_file(path) 

291 

292 actual = imod.prepare.spatial.zonal_aggregate_raster( 

293 path=path, 

294 column="id", 

295 raster=raster, 

296 resolution=1.0, 

297 method="mean", 

298 ) 

299 expected = pd.DataFrame.from_dict( 

300 { 

301 "id": [1, 2], 

302 "area": [2.0, 2.0], 

303 "my-raster": [0.0, 1.0], 

304 } 

305 ) 

306 assert_frame_equal(actual, expected) 

307 

308 raster.name = None 

309 geometries = [ 

310 sg.Polygon( 

311 [ 

312 (0.0, 1.5), 

313 (2.0, 1.5), 

314 (2.0, 2.0), 

315 (0.0, 2.0), 

316 ] 

317 ), 

318 sg.Polygon( 

319 [ 

320 (0.0, 0.0), 

321 (2.0, 0.0), 

322 (2.0, 1.5), 

323 (0.0, 1.5), 

324 ] 

325 ), 

326 ] 

327 gdf.geometry = geometries 

328 path = tmp_path / "two-zones2.shp" 

329 gdf.to_file(path) 

330 

331 actual = imod.prepare.spatial.zonal_aggregate_raster( 

332 path=path, 

333 column="id", 

334 raster=raster, 

335 resolution=0.5, 

336 method="mean", 

337 ) 

338 expected = pd.DataFrame.from_dict( 

339 { 

340 "id": [1, 2], 

341 "area": [1.0, 3.0], 

342 "aggregated": [0.0, 2.0 / 3.0], 

343 } 

344 ) 

345 assert_frame_equal(actual, expected) 

346 

347 actual = imod.prepare.spatial.zonal_aggregate_raster( 

348 path=path, 

349 column="id", 

350 raster=raster, 

351 resolution=0.5, 

352 method=pd.Series.mode, 

353 ) 

354 expected = pd.DataFrame.from_dict( 

355 { 

356 "id": [1, 2], 

357 "area": [1.0, 3.0], 

358 "aggregated": [0.0, 1.0], 

359 } 

360 ) 

361 assert_frame_equal(actual, expected) 

362 

363 # Now test no overlap 

364 # The functions internally explicitly return a zero row dataframe 

365 # Since a method like pd.Series.mode cannot deal with 0 rows. 

366 raster = xr.DataArray( 

367 data=[[0.0, 0.0], [1.0, 1.0]], 

368 coords={"x": [10.5, 11.5], "y": [11.5, 10.5]}, 

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

370 name="my-raster", 

371 ) 

372 actual = imod.prepare.spatial.zonal_aggregate_raster( 

373 path=path, 

374 column="id", 

375 raster=raster, 

376 resolution=0.5, 

377 method=pd.Series.mode, 

378 ) 

379 assert len(actual) == 0 

380 

381 

382def test_zonal_aggregate_polygons(tmp_path): 

383 raster = xr.DataArray( 

384 data=[[np.nan, np.nan], [np.nan, np.nan]], 

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

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

387 ) 

388 geometries_a = [ 

389 sg.Polygon( 

390 [ 

391 (0.0, 1.0), 

392 (2.0, 1.0), 

393 (2.0, 2.0), 

394 (0.0, 2.0), 

395 ] 

396 ), 

397 sg.Polygon( 

398 [ 

399 (0.0, 0.0), 

400 (2.0, 0.0), 

401 (2.0, 1.0), 

402 (0.0, 1.0), 

403 ] 

404 ), 

405 ] 

406 gdf_a = gpd.GeoDataFrame(geometry=geometries_a) 

407 gdf_a["id_a"] = [1, 2] 

408 gdf_a["data_a"] = [3, 4] 

409 path_a = tmp_path / "two-zones_a.shp" 

410 gdf_a.to_file(path_a) 

411 

412 actual = imod.prepare.spatial.zonal_aggregate_polygons( 

413 path_a=path_a, 

414 column_a="id_a", 

415 path_b=path_a, 

416 column_b="data_a", 

417 like=raster, 

418 resolution=1.0, 

419 method="mean", 

420 ) 

421 expected = pd.DataFrame.from_dict( 

422 { 

423 "id_a": [1, 2], 

424 "area": [2.0, 2.0], 

425 "data_a": [3.0, 4.0], 

426 } 

427 ) 

428 assert_frame_equal(actual, expected) 

429 

430 geometries_b = [ 

431 sg.Polygon( 

432 [ 

433 (0.0, 2.0), 

434 (0.0, 0.0), 

435 (1.0, 0.0), 

436 (1.0, 2.0), 

437 ] 

438 ), 

439 sg.Polygon( 

440 [ 

441 (1.0, 2.0), 

442 (1.0, 0.0), 

443 (2.0, 0.0), 

444 (2.0, 2.0), 

445 ] 

446 ), 

447 ] 

448 gdf_b = gpd.GeoDataFrame(geometry=geometries_b) 

449 gdf_b["data_b"] = [3, 4] 

450 path_b = tmp_path / "two-zones_b.shp" 

451 gdf_b.to_file(path_b) 

452 

453 actual = imod.prepare.spatial.zonal_aggregate_polygons( 

454 path_a=path_a, 

455 column_a="id_a", 

456 path_b=path_b, 

457 column_b="data_b", 

458 like=raster, 

459 resolution=1.0, 

460 method="mean", 

461 ) 

462 expected = pd.DataFrame.from_dict( 

463 { 

464 "id_a": [1, 2], 

465 "area": [2.0, 2.0], 

466 "data_b": [3.5, 3.5], 

467 } 

468 ) 

469 assert_frame_equal(actual, expected) 

470 

471 # Now test no overlap 

472 # The functions internally explicitly return a zero row dataframe 

473 # Since a method like pd.Series.mode cannot deal with 0 rows. 

474 raster = xr.DataArray( 

475 data=[[0.0, 0.0], [1.0, 1.0]], 

476 coords={"x": [10.5, 11.5], "y": [11.5, 10.5]}, 

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

478 name="my-raster", 

479 ) 

480 actual = imod.prepare.spatial.zonal_aggregate_polygons( 

481 path_a=path_a, 

482 column_a="id_a", 

483 path_b=path_b, 

484 column_b="data_b", 

485 like=raster, 

486 resolution=1.0, 

487 method="mean", 

488 ) 

489 assert len(actual) == 0