Coverage for C:\src\imod-python\imod\prepare\reproject.py: 75%
88 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import affine
2import numpy as np
3import xarray as xr
5import imod
6from imod.util.imports import MissingOptionalModule
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")
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
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.
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``.
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.
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:
71 * string: e.g. ``"EPSG:4326"``
72 * rasterio.crs.CRS
73 dst_crs: string, dict, rasterio.crs.CRS
74 Coordinate system of result. Options:
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:
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()``.
101 Returns
102 -------
103 xarray.DataArray
104 Resampled/reprojected DataArray.
106 Examples
107 --------
108 Resample a DataArray ``a`` to a new cellsize, using an existing DataArray ``b``:
110 >>> c = imod.rasterio.reproject(source=a, like=b)
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.)
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)
120 Reproject a DataArray from one coordinate system (WGS84, EPSG:4326) to another (UTM30N, EPSG:32630):
122 >>> c = imod.rasterio.reproject(source=a, src_crs="EPSG:4326", dst_crs="EPSG:32630")
124 Get the reprojected DataArray in the desired shape and coordinates by providing ``like``:
126 >>> c = imod.rasterio.reproject(source=a, like=b, src_crs="EPSG:4326", dst_crs="EPSG:32630")
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.
131 >>> a = rioxarray.open_rasterio("example.tif").squeeze("band")
132 >>> c = imod.rasterio.reproject(source=a, use_src_attrs=True, dst_crs="EPSG:28992")
134 In case of a rotated ``source``, provide ``src_transform`` directly or ``use_src_attrs=True`` to rely on generated attributes:
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
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
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 )
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
180 resampling_methods = {e.name: e for e in rasterio.enums.Resampling}
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")
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)
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)
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 )
243 else:
244 raise ValueError(
245 "At least ``like``, or crs information for source and destination must be provided."
246 )
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")
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 )
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