Coverage for C:\src\imod-python\imod\prepare\reproject.py: 75%

88 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import affine 

2import numpy as np 

3import xarray as xr 

4 

5import imod 

6from imod.util.imports import MissingOptionalModule 

7 

8# since rasterio is a big dependency that is sometimes hard to install 

9# and not always required, we made this an optional dependency 

10try: 

11 import rasterio 

12 import rasterio.warp 

13except ImportError: 

14 rasterio = MissingOptionalModule("rasterio") 

15 

16 

17def _reproject_dst(source, src_crs, dst_crs, src_transform): 

18 """ 

19 Prepares destination transform Affine and DataArray for projection. 

20 """ 

21 src_height, src_width = source.y.size, source.x.size 

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

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

24 src_crs, dst_crs, src_width, src_height, *bounds 

25 ) 

26 # from: http://xarray.pydata.org/en/stable/generated/xarray.open_rasterio.html 

27 x, y = dst_transform * np.meshgrid( 

28 np.arange(dst_width) + 0.5, np.arange(dst_height) + 0.5 

29 ) 

30 dst = xr.DataArray( 

31 data=np.zeros((dst_height, dst_width), source.dtype), 

32 coords={"y": y[:, 0], "x": x[0, :]}, 

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

34 ) 

35 return dst_transform, dst 

36 

37 

38def reproject( 

39 source, 

40 like=None, 

41 src_crs=None, 

42 dst_crs=None, 

43 method="nearest", 

44 use_src_attrs=False, 

45 src_nodata=np.nan, 

46 **reproject_kwargs, 

47): 

48 """ 

49 Reprojects and/or resamples a 2D xarray DataArray to a 

50 different cellsize or coordinate system. 

51 

52 * To resample to a new cellsize in the same projection: provide only ``like``. 

53 * To only reproject: provide only ``src_crs`` and ``src_crs``. 

54 * To reproject and resample to a specific domain: provide ``src_crs``, ``src_crs``, and ``like``. 

55 

56 Note: when only ``like`` is provided, Cartesian (projected) coordinates are a 

57 ssumed for resampling. In case of non-Cartesian coordinates, specify 

58 ``src_crs`` and ``dst_crs`` for correct resampling. 

59 

60 Parameters 

61 ---------- 

62 source: xarray DataArray 

63 The DataArray to be resampled and/or reprojected. Must contain dimensions 

64 ``y`` and ``x``. 

65 like: xarray DataArray 

66 Example DataArray that shows what the resampled result should look like 

67 in terms of coordinates. Must contain dimensions ``y`` and ``x``. 

68 src_crs: string, dict, rasterio.crs.CRS 

69 Coordinate system of ``source``. Options: 

70 

71 * string: e.g. ``"EPSG:4326"`` 

72 * rasterio.crs.CRS 

73 dst_crs: string, dict, rasterio.crs.CRS 

74 Coordinate system of result. Options: 

75 

76 * string: e.g. ``"EPSG:4326"`` 

77 * rasterio.crs.CRS 

78 use_src_attrs: boolean 

79 If True: Use metadata in ``source.attrs``, as generated by ``xarray.open_rasterio()``, to do 

80 reprojection. 

81 method: string 

82 The method to use for resampling/reprojection. 

83 Defaults to "nearest". GDAL methods are available: 

84 

85 * nearest 

86 * bilinear 

87 * cubic 

88 * cubic_spline 

89 * lanczos 

90 * average 

91 * mode 

92 * gauss 

93 * max 

94 * min 

95 * med (50th percentile) 

96 * q1 (25th percentile) 

97 * q3 (75th percentile) 

98 reproject_kwargs: dict, optional 

99 keyword arguments for ``rasterio.warp.reproject()``. 

100 

101 Returns 

102 ------- 

103 xarray.DataArray 

104 Resampled/reprojected DataArray. 

105 

106 Examples 

107 -------- 

108 Resample a DataArray ``a`` to a new cellsize, using an existing DataArray ``b``: 

109 

110 >>> c = imod.rasterio.reproject(source=a, like=b) 

111 

112 Resample a DataArray to a new cellsize of 100.0, by creating a ``like`` DataArray first: 

113 (Note that dy must be negative, as is usual for geospatial grids.) 

114 

115 >>> dims = ("y", "x") 

116 >>> coords = {"y": np.arange(200_000.0, 100_000.0, -100.0), "x": np.arange(0.0, 100_000.0, 100.0)} 

117 >>> b = xr.DataArray(data=np.empty((200, 100)), coords=coords, dims=dims) 

118 >>> c = imod.rasterio.reproject(source=a, like=b) 

119 

120 Reproject a DataArray from one coordinate system (WGS84, EPSG:4326) to another (UTM30N, EPSG:32630): 

121 

122 >>> c = imod.rasterio.reproject(source=a, src_crs="EPSG:4326", dst_crs="EPSG:32630") 

123 

124 Get the reprojected DataArray in the desired shape and coordinates by providing ``like``: 

125 

126 >>> c = imod.rasterio.reproject(source=a, like=b, src_crs="EPSG:4326", dst_crs="EPSG:32630") 

127 

128 Open a single band raster, and reproject to RD new coordinate system (EPSG:28992), without explicitly specifying ``src_crs``. 

129 ``src_crs`` is taken from ``a.attrs``, so the raster file has to include coordinate system metadata for this to work. 

130 

131 >>> a = rioxarray.open_rasterio("example.tif").squeeze("band") 

132 >>> c = imod.rasterio.reproject(source=a, use_src_attrs=True, dst_crs="EPSG:28992") 

133 

134 In case of a rotated ``source``, provide ``src_transform`` directly or ``use_src_attrs=True`` to rely on generated attributes: 

135 

136 >>> rotated = rioxarray.open_rasterio("rotated_example.tif").squeeze("band") 

137 >>> c = imod.rasterio.reproject(source=rotated, dst_crs="EPSG:28992", reproject_kwargs={"src_transform":affine.Affine(...)}) 

138 >>> c = imod.rasterio.reproject(source=rotated, dst_crs="EPSG:28992", use_src_attrs=True) 

139 """ 

140 # Make sure the rio accessor is avaible. 

141 import rioxarray # noqa pylint: F401 

142 

143 if not source.dims == ("y", "x"): 

144 raise ValueError( 

145 "reproject does not support dimensions other than ``x`` and ``y`` for ``source``." 

146 ) 

147 if like is not None: 

148 if not like.dims == ("y", "x"): 

149 raise ValueError( 

150 "reproject does not support dimensions other than ``x`` and ``y`` for ``like``." 

151 ) 

152 if use_src_attrs: # only provided when reproject is necessary 

153 try: 

154 src_crs = source.attrs["crs"] 

155 except KeyError: 

156 src_crs = source.rio.crs 

157 

158 if isinstance(src_crs, str): 

159 if "epsg:" in src_crs.lower(): 

160 # Workaround for rioxarray.open_rasterio generation proj4 strings 

161 # https://github.com/mapbox/rasterio/issues/1809 

162 epsg_code = src_crs.lower().split("epsg:")[-1] 

163 src_crs = rasterio.crs.CRS.from_epsg(epsg_code) 

164 else: 

165 src_crs = rasterio.crs.CRS.from_string(src_crs) 

166 elif isinstance(src_crs, rasterio.crs.CRS): 

167 pass 

168 else: 

169 raise ValueError( 

170 f"Invalid src_crs: {src_crs}. Must be either str or rasterio.crs.CRS object" 

171 ) 

172 

173 if "nodatavals" in source.attrs: 

174 src_nodata = source.attrs["nodatavals"][0] 

175 else: 

176 rio_nodata = source.rio.nodata 

177 if rio_nodata is not None: 

178 src_nodata = rio_nodata 

179 

180 resampling_methods = {e.name: e for e in rasterio.enums.Resampling} 

181 

182 if isinstance(method, str): 

183 try: 

184 resampling_method = resampling_methods[method] 

185 except KeyError as e: 

186 raise ValueError( 

187 "Invalid resampling method. Available methods are: {}".format( 

188 resampling_methods.keys() 

189 ) 

190 ) from e 

191 elif isinstance(method, rasterio.enums.Resampling): 

192 resampling_method = method 

193 else: 

194 raise TypeError("method must be a string or rasterio.enums.Resampling") 

195 

196 # Givens: source, like, method. No reprojection necessary. 

197 if src_crs is None and dst_crs is None: 

198 if like is None: 

199 raise ValueError( 

200 "If crs information is not provided, ``like`` must be provided." 

201 ) 

202 if resampling_method == rasterio.enums.Resampling.nearest: 

203 # this can be handled with xarray 

204 # xarray 0.10.9 needs .compute() 

205 # see https://github.com/pydata/xarray/issues/2454 

206 return source.compute().reindex_like(like, method="nearest") 

207 else: 

208 # if no crs is defined, assume it should remain the same 

209 # in this case use UTM30, ESPG:32630, as a dummy value for GDAL 

210 # (Any projected coordinate system should suffice, Cartesian plane == Cartesian plane) 

211 dst = xr.DataArray( 

212 data=np.zeros(like.shape, source.dtype), 

213 coords={"y": like.y, "x": like.x}, 

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

215 ) 

216 src_crs = dst_crs = rasterio.crs.CRS.from_epsg(32630) 

217 src_transform = imod.util.spatial.transform(source) 

218 dst_transform = imod.util.spatial.transform(like) 

219 

220 elif src_crs and dst_crs: 

221 if use_src_attrs: 

222 # TODO: modify if/when xarray uses affine by default for transform 

223 try: 

224 src_transform = affine.Affine(*source.attrs["transform"][:6]) 

225 except KeyError: 

226 src_transform = source.rio.transform() 

227 elif "src_transform" in reproject_kwargs.keys(): 

228 src_transform = reproject_kwargs.pop("src_transform") 

229 else: 

230 src_transform = imod.util.spatial.transform(source) 

231 

232 # If no like is provided, just reproject to different coordinate system 

233 if like is None: 

234 dst_transform, dst = _reproject_dst(source, src_crs, dst_crs, src_transform) 

235 else: 

236 dst_transform = imod.util.spatial.transform(like) 

237 dst = xr.DataArray( 

238 data=np.zeros(like.shape, source.dtype), 

239 coords={"y": like.y, "x": like.x}, 

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

241 ) 

242 

243 else: 

244 raise ValueError( 

245 "At least ``like``, or crs information for source and destination must be provided." 

246 ) 

247 

248 if not src_transform[0] > 0: 

249 raise ValueError("dx of 'source' must be positive") 

250 if not src_transform[4] < 0: 

251 raise ValueError("dy of 'source' must be negative") 

252 if not dst_transform[0] > 0: 

253 raise ValueError("dx of 'like' must be positive") 

254 if not dst_transform[4] < 0: 

255 raise ValueError("dy of 'like' must be negative") 

256 

257 rasterio.warp.reproject( 

258 source.values, 

259 dst.values, 

260 src_transform=src_transform, 

261 dst_transform=dst_transform, 

262 src_crs=src_crs, 

263 dst_crs=dst_crs, 

264 resampling=resampling_method, 

265 src_nodata=src_nodata, 

266 **reproject_kwargs, 

267 ) 

268 

269 dst.attrs = source.attrs 

270 dst.attrs["transform"] = dst_transform 

271 dst.attrs["res"] = (abs(dst_transform[0]), abs(dst_transform[4])) 

272 dst.attrs["crs"] = dst_crs 

273 # TODO: what should be the type of "crs" field in attrs? 

274 # Doesn't actually matter for functionality 

275 # rasterio accepts string, dict, and CRS object 

276 return dst