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

1 import affine

2 import numpy as np

3 import xarray as xr

4

5 import imod

6 from 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

10 try:

11 import rasterio

12 import rasterio.warp

13 except ImportError:

14 rasterio = MissingOptionalModule("rasterio")

15

16

17 def _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

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