Coverage for test_reproject.py: 100%

158 statements  

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

1import affine 

2import numpy as np 

3import pytest 

4import rasterio 

5import rasterio.warp 

6import rioxarray 

7import xarray as xr 

8 

9import imod 

10 

11 

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

13def write_tif(): 

14 def _write_tif(path, epsg, dtype=np.float64, rotation_angle=None): 

15 nrow = 5 

16 ncol = 8 

17 values = (10 * np.random.rand(nrow, ncol)).astype(dtype) 

18 

19 profile = {} 

20 profile["crs"] = rasterio.crs.CRS.from_epsg(epsg) 

21 profile["transform"] = affine.Affine(1.0, 0.0, 155_000.0, 0.0, -1.0, 463_000.0) 

22 if rotation_angle: 

23 profile["transform"] *= affine.Affine.rotation(rotation_angle) 

24 profile["driver"] = "GTiff" 

25 profile["height"] = nrow 

26 profile["width"] = ncol 

27 profile["count"] = 1 

28 profile["dtype"] = dtype 

29 

30 with rasterio.Env(): 

31 with rasterio.open(path, "w", **profile) as ds: 

32 ds.write(values, 1) 

33 

34 return _write_tif 

35 

36 

37def test_basic_resample__nearest(write_tif, tmp_path): 

38 """Test nearest neighbour upsampling to halved cellsize""" 

39 write_tif(tmp_path / "basic.tif", epsg=28992) 

40 

41 with rioxarray.open_rasterio(tmp_path / "basic.tif").squeeze("band") as da: 

42 dx, xmin, xmax, dy, ymin, ymax = imod.util.spatial.spatial_reference(da) 

43 data = np.empty((10, 16)) 

44 coords = { 

45 "y": np.linspace(ymax + 0.25 * dy, ymin - 0.25 * dy, 10), 

46 "x": np.linspace(xmin + 0.25 * dx, xmax - 0.25 * dx, 16), 

47 } 

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

49 like = xr.DataArray(data, coords, dims) 

50 newda = imod.prepare.reproject(da, like, method="nearest") 

51 

52 newarr = np.empty((10, 16)) 

53 with rasterio.open(tmp_path / "basic.tif") as src: 

54 arr = src.read() 

55 aff = src.transform 

56 crs = src.crs 

57 newaff = affine.Affine(aff.a / 2.0, aff.b, aff.c, aff.d, aff.e / 2.0, aff.f) 

58 

59 rasterio.warp.reproject( 

60 arr, 

61 newarr, 

62 src_transform=aff, 

63 dst_transform=newaff, 

64 src_crs=crs, 

65 dst_crs=crs, 

66 resampling=rasterio.enums.Resampling.nearest, 

67 src_nodata=np.nan, 

68 ) 

69 

70 # Double values in both axes 

71 repeated = np.repeat(da.values, 2, axis=0) 

72 repeated = np.repeat(repeated, 2, axis=1) 

73 assert np.allclose(newda.values, repeated) 

74 assert np.allclose(newda.values, newarr) 

75 

76 

77def test_basic_resample__bilinear(write_tif, tmp_path): 

78 """Test bilinear upsampling to halved cellsize""" 

79 write_tif(tmp_path / "basic.tif", epsg=28992) 

80 

81 with rioxarray.open_rasterio(tmp_path / "basic.tif").squeeze("band") as da: 

82 dx, xmin, xmax, dy, ymin, ymax = imod.util.spatial.spatial_reference(da) 

83 data = np.empty((10, 16)) 

84 coords = { 

85 "y": np.linspace(ymax + 0.25 * dy, ymin - 0.25 * dy, 10), 

86 "x": np.linspace(xmin + 0.25 * dx, xmax - 0.25 * dx, 16), 

87 } 

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

89 like = xr.DataArray(data, coords, dims) 

90 newda = imod.prepare.reproject(da, like, method="bilinear") 

91 

92 newarr = np.empty((10, 16)) 

93 with rasterio.open(tmp_path / "basic.tif") as src: 

94 arr = src.read() 

95 aff = src.transform 

96 crs = src.crs 

97 newaff = affine.Affine(aff.a / 2.0, aff.b, aff.c, aff.d, aff.e / 2.0, aff.f) 

98 

99 rasterio.warp.reproject( 

100 arr, 

101 newarr, 

102 src_transform=aff, 

103 dst_transform=newaff, 

104 src_crs=crs, 

105 dst_crs=crs, 

106 resampling=rasterio.enums.Resampling.bilinear, 

107 src_nodata=np.nan, 

108 ) 

109 

110 assert np.allclose(newda.values, newarr) 

111 

112 

113def test_basic_reproject(write_tif, tmp_path): 

114 """Basic reprojection from EPSG:28992 to EPSG:32631""" 

115 write_tif(tmp_path / "basic.tif", epsg=28992) 

116 dst_crs = {"init": "EPSG:32631"} 

117 with rioxarray.open_rasterio(tmp_path / "basic.tif").squeeze("band") as da: 

118 newda = imod.prepare.reproject(da, src_crs="EPSG:28992", dst_crs=dst_crs) 

119 

120 with rasterio.open(tmp_path / "basic.tif") as src: 

121 arr = src.read() 

122 src_transform = src.transform 

123 src_crs = src.crs 

124 src_height = src.height 

125 src_width = src.width 

126 

127 src_crs = rasterio.crs.CRS(src_crs) 

128 bounds = rasterio.transform.array_bounds(src_height, src_width, src_transform) 

129 dst_transform, dst_width, dst_height = rasterio.warp.calculate_default_transform( 

130 src_crs, dst_crs, src_width, src_height, *bounds 

131 ) 

132 

133 newarr = np.empty((dst_height, dst_width)) 

134 

135 rasterio.warp.reproject( 

136 arr, 

137 newarr, 

138 src_transform=src_transform, 

139 dst_transform=dst_transform, 

140 src_crs=src_crs, 

141 dst_crs=dst_crs, 

142 resampling=rasterio.enums.Resampling.nearest, 

143 src_nodata=np.nan, 

144 ) 

145 

146 assert np.allclose(newda.values, newarr) 

147 

148 

149def test_reproject__use_src_attrs(write_tif, tmp_path): 

150 """Reprojection from EPSG:28992 to EPSG:32631, using on attrs generated by xarray.""" 

151 write_tif(tmp_path / "basic.tif", epsg=28992) 

152 dst_crs = "EPSG:32631" 

153 with rioxarray.open_rasterio(tmp_path / "basic.tif").squeeze("band") as da: 

154 newda = imod.prepare.reproject(da, dst_crs=dst_crs, use_src_attrs=True) 

155 

156 with rasterio.open(tmp_path / "basic.tif") as src: 

157 arr = src.read() 

158 src_transform = src.transform 

159 src_crs = src.crs 

160 src_height = src.height 

161 src_width = src.width 

162 

163 src_crs = rasterio.crs.CRS(src_crs) 

164 bounds = rasterio.transform.array_bounds(src_height, src_width, src_transform) 

165 dst_transform, dst_width, dst_height = rasterio.warp.calculate_default_transform( 

166 src_crs, dst_crs, src_width, src_height, *bounds 

167 ) 

168 

169 newarr = np.empty((dst_height, dst_width)) 

170 

171 rasterio.warp.reproject( 

172 arr, 

173 newarr, 

174 src_transform=src_transform, 

175 dst_transform=dst_transform, 

176 src_crs=src_crs, 

177 dst_crs=dst_crs, 

178 resampling=rasterio.enums.Resampling.nearest, 

179 src_nodata=np.nan, 

180 ) 

181 

182 assert np.allclose(newda.values, newarr) 

183 

184 

185def test_reproject_resample(write_tif, tmp_path): 

186 """ 

187 Reprojection from EPSG:28992 to EPSG:32631, using on attrs generated by xarray, 

188 then resample to a like DataArray 

189 """ 

190 write_tif(tmp_path / "basic.tif", epsg=28992) 

191 # ESPG:32631 coords 

192 xmin = 663_304.0 

193 xmax = 663_312.0 

194 ymin = 5_780_979.5 

195 ymax = 5_780_984.5 

196 dx = 0.5 

197 dy = -0.5 

198 with rioxarray.open_rasterio(tmp_path / "basic.tif").squeeze("band") as da: 

199 dst_crs = "EPSG:32631" 

200 data = np.empty((10, 16)) 

201 coords = { 

202 "y": np.linspace(ymax + 0.5 * dy, ymin - 0.5 * dy, 10), 

203 "x": np.linspace(xmin + 0.5 * dx, xmax - 0.5 * dx, 16), 

204 } 

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

206 like = xr.DataArray(data, coords, dims) 

207 newda = imod.prepare.reproject(da, like, dst_crs=dst_crs, use_src_attrs=True) 

208 

209 with rasterio.open(tmp_path / "basic.tif") as src: 

210 arr = src.read() 

211 src_transform = src.transform 

212 src_crs = src.crs 

213 

214 dst_transform = affine.Affine(dx, 0.0, xmin, 0.0, dy, ymax) 

215 newarr = np.empty((10, 16)) 

216 rasterio.warp.reproject( 

217 arr, 

218 newarr, 

219 src_transform=src_transform, 

220 dst_transform=dst_transform, 

221 src_crs=src_crs, 

222 dst_crs=dst_crs, 

223 resampling=rasterio.enums.Resampling.nearest, 

224 src_nodata=np.nan, 

225 ) 

226 

227 assert ~np.isnan(newda.values).all() 

228 assert np.allclose(newda.values, newarr) 

229 

230 

231def test_reproject_rotation__via_kwargs(write_tif, tmp_path): 

232 """Reprojection from EPSG:28992 to EPSG:32631, by specifying kwarg""" 

233 # Arrange 

234 write_tif(tmp_path / "rotated.tif", epsg=28992, rotation_angle=45.0) 

235 dst_crs = "EPSG:32631" 

236 da = rioxarray.open_rasterio(tmp_path / "rotated.tif").squeeze("band") 

237 src_transform = ( 

238 affine.Affine.translation(155000.0, 463000.0) 

239 * affine.Affine.scale(1.0, -1.0) 

240 * affine.Affine.rotation(45.0) 

241 ) 

242 newda = imod.prepare.reproject( 

243 da, src_crs="EPSG:28992", dst_crs=dst_crs, src_transform=src_transform 

244 ) 

245 

246 with rasterio.open(tmp_path / "rotated.tif") as src: 

247 arr = src.read() 

248 src_transform = src.transform 

249 src_crs = src.crs 

250 src_height = src.height 

251 src_width = src.width 

252 

253 src_crs = rasterio.crs.CRS(src_crs) 

254 bounds = rasterio.transform.array_bounds(src_height, src_width, src_transform) 

255 dst_transform, dst_width, dst_height = rasterio.warp.calculate_default_transform( 

256 src_crs, dst_crs, src_width, src_height, *bounds 

257 ) 

258 

259 # Act 

260 newarr = np.empty((dst_height, dst_width)) 

261 

262 rasterio.warp.reproject( 

263 arr, 

264 newarr, 

265 src_transform=src_transform, 

266 dst_transform=dst_transform, 

267 src_crs=src_crs, 

268 dst_crs=dst_crs, 

269 resampling=rasterio.enums.Resampling.nearest, 

270 src_nodata=np.nan, 

271 ) 

272 

273 # Assert 

274 np.testing.assert_allclose(newda.values, newarr, equal_nan=True) 

275 

276 

277def test_reproject_rotation__use_src_attrs(write_tif, tmp_path): 

278 """Reprojection from EPSG:28992 to EPSG:32631, using attrs generated by xarray.""" 

279 write_tif(tmp_path / "rotated.tif", epsg=28992, rotation_angle=45.0) 

280 dst_crs = "EPSG:32631" 

281 da = rioxarray.open_rasterio(tmp_path / "rotated.tif").squeeze("band") 

282 newda = imod.prepare.reproject(da, dst_crs=dst_crs, use_src_attrs=True) 

283 

284 with rasterio.open(tmp_path / "rotated.tif") as src: 

285 arr = src.read() 

286 src_transform = src.transform 

287 src_crs = src.crs 

288 src_height = src.height 

289 src_width = src.width 

290 

291 src_crs = rasterio.crs.CRS(src_crs) 

292 bounds = rasterio.transform.array_bounds(src_height, src_width, src_transform) 

293 dst_transform, dst_width, dst_height = rasterio.warp.calculate_default_transform( 

294 src_crs, dst_crs, src_width, src_height, *bounds 

295 ) 

296 

297 newarr = np.empty((dst_height, dst_width)) 

298 

299 rasterio.warp.reproject( 

300 arr, 

301 newarr, 

302 src_transform=src_transform, 

303 dst_transform=dst_transform, 

304 src_crs=src_crs, 

305 dst_crs=dst_crs, 

306 resampling=rasterio.enums.Resampling.nearest, 

307 src_nodata=np.nan, 

308 ) 

309 assert np.allclose(newda.values, newarr, equal_nan=True)