Coverage for C:\src\imod-python\imod\visualize\spatial.py: 10%
179 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
1import pathlib
3import matplotlib
4import matplotlib.pyplot as plt
5import numpy as np
6import pandas as pd
7from mpl_toolkits.axes_grid1 import make_axes_locatable
9import imod
10from imod.util.imports import MissingOptionalModule
11from imod.visualize import common
13try:
14 import contextily as ctx
15except ImportError:
16 ctx = MissingOptionalModule("contextily")
19def read_imod_legend(path):
20 """
21 Parameters
22 ----------
23 path : str
24 Path to iMOD .leg file.
26 Returns
27 -------
28 colors : List of hex colors of length N.
29 levels : List of floats of length N-1. These are the boundaries between
30 the legend colors/classes.
31 """
33 # Read file. Do not rely the headers in the leg file.
34 def _read(delim_whitespace):
35 return pd.read_csv(
36 path,
37 header=1,
38 delim_whitespace=delim_whitespace,
39 index_col=False,
40 usecols=[0, 1, 2, 3, 4],
41 names=["upper", "lower", "red", "green", "blue"],
42 )
44 # Try both comma and whitespace separated
45 try:
46 legend = _read(delim_whitespace=False)
47 except ValueError:
48 legend = _read(delim_whitespace=True)
50 # The colors in iMOD are formatted in RGB. Format to hexadecimal.
51 red = legend["red"]
52 blue = legend["blue"]
53 green = legend["green"]
54 colors = [f"#{r:02x}{g:02x}{b:02x}" for r, g, b in zip(red, green, blue)][::-1]
55 levels = list(legend["lower"])[::-1][1:]
57 return colors, levels
60def _crs2string(crs):
61 if isinstance(crs, str):
62 if "epsg" in crs.lower():
63 return crs
64 try:
65 return crs.to_string()
66 except AttributeError:
67 try:
68 return crs["init"]
69 except KeyError:
70 return crs
73def plot_map(
74 raster,
75 colors,
76 levels,
77 overlays=[],
78 basemap=None,
79 kwargs_raster=None,
80 kwargs_colorbar=None,
81 kwargs_basemap={},
82 figsize=None,
83 return_cbar=False,
84 fig=None,
85 ax=None,
86):
87 """
88 Plot raster on a map with optional vector overlays and basemap.
90 Parameters
91 ----------
92 raster : xr.DataArray
93 2D grid to plot.
94 colors : list of str, list of RGBA/RGBA tuples, colormap name (str), or
95 LinearSegmentedColormap.
97 If list, it should be a Matplotlib acceptable list of colors. Length N.
98 Accepts both tuples of (R, G, B) and hexidecimal (e.g. `#7ec0ee`).
100 If str, use an existing Matplotlib colormap. This function will
101 autmatically add distinctive colors for pixels lower or high than the
102 given min respectivly max level.
104 If LinearSegmentedColormap, you can use something like
105 `matplotlib.cm.get_cmap('jet')` as input. This function will not alter
106 the colormap, so add under- and over-colors yourself.
108 Looking for good colormaps? Try: http://colorbrewer2.org/ Choose a
109 colormap, and use the HEX JS array.
110 levels : listlike of floats or integers
111 Boundaries between the legend colors/classes. Length: N - 1.
112 overlays : list of dicts, optional
113 Dicts contain geodataframe (key is "gdf"), and the keyword arguments for
114 plotting the geodataframe.
115 basemap : bool or contextily._providers.TileProvider, optional
116 When `True` or a `contextily._providers.TileProvider` object: plot a
117 basemap as a background for the plot and make the raster translucent. If
118 `basemap=True`, then `CartoDB.Positron` is used as the default provider.
119 If not set explicitly through kwargs_basemap, plot_map() will try and
120 infer the crs from the raster or overlays, or fall back to EPSG:28992
121 (Amersfoort/RDnew).
123 *Requires contextily*
125 kwargs_raster : dict of keyword arguments, optional
126 These arguments are forwarded to ax.imshow()
127 kwargs_colorbar : dict of keyword arguments, optional
128 These arguments are forwarded to fig.colorbar(). The key label can be
129 used to label the colorbar. Key whiten_triangles can be set to False to
130 alter the default behavior of coloring the min / max triangles of the
131 colorbar white if the value is not present in the map.
132 kwargs_basemap : dict of keyword arguments, optional
133 Except for "alpha", these arguments are forwarded to
134 contextily.add_basemap(). Parameter "alpha" controls the transparency of
135 raster.
136 figsize : tuple of two floats or integers, optional
137 This is used in plt.subplots(figsize)
138 return_cbar : boolean, optional
139 Return the matplotlib.Colorbar instance. Defaults to False.
140 fig : matplotlib.figure, optional
141 If provided, figure to which to add the map
142 ax : matplot.ax, optional
143 If provided, axis to which to add the map
145 Returns
146 -------
147 fig : matplotlib.figure ax : matplotlib.ax if return_cbar == True: cbar :
148 matplotlib.Colorbar
150 Examples
151 --------
152 Plot with an overlay:
154 >>> overlays = [{"gdf": geodataframe, "edgecolor": "black", "facecolor": "None"}]
155 >>> imod.visualize.plot_map(raster, colors, levels, overlays)
157 Label the colorbar:
159 >>> imod.visualize.plot_map(raster, colors, levels, kwargs_colorbar={"label":"Head aquifer (m)"})
161 Plot with a basemap:
163 >>> import contextily as ctx
164 >>> src = ctx.providers.Stamen.TonerLite
165 >>> imod.visualize.plot_map(raster, colors, levels, basemap=src, kwargs_basemap={"alpha":0.6})
167 """
168 # Account for both None or False to skip adding a basemap
169 if basemap is None or (isinstance(basemap, bool) and not basemap):
170 add_basemap = False
171 else:
172 add_basemap = True
174 # Read legend settings
175 cmap, norm = common._cmapnorm_from_colorslevels(colors, levels)
177 # Get extent
178 _, xmin, xmax, _, ymin, ymax = imod.util.spatial.spatial_reference(raster)
180 # raster kwargs
181 settings_raster = {"interpolation": "nearest", "extent": [xmin, xmax, ymin, ymax]}
182 # if a basemap is added: set alpha of raster
183 if add_basemap:
184 settings_raster["alpha"] = kwargs_basemap.pop("alpha", 0.7)
185 if kwargs_raster is not None:
186 settings_raster.update(kwargs_raster)
188 # cbar kwargs
189 settings_cbar = {"ticks": levels, "extend": "both"}
191 # Find a unit in the raster to use in the colorbar label
192 try:
193 settings_cbar["label"] = raster.attrs["units"]
194 except (KeyError, AttributeError):
195 try:
196 settings_cbar["label"] = raster.attrs["unit"]
197 except (KeyError, AttributeError):
198 pass
200 whiten_triangles = True
201 if kwargs_colorbar is not None:
202 whiten_triangles = kwargs_colorbar.pop("whiten_triangles", True)
203 settings_cbar.update(kwargs_colorbar)
205 # If not provided, make figure and axes
206 # Catch case first where no figure provided, but ax was provided
207 if fig is None and ax is not None:
208 raise ValueError(
209 "Axes provided, yet no figure is provided. "
210 "Please provide a figure as well."
211 )
212 if fig is None:
213 fig = plt.figure(figsize=figsize)
214 if ax is None:
215 ax = plt.axes()
217 # Make sure x is increasing, y is decreasing
218 raster = raster.copy(deep=False)
219 flip = slice(None, None, -1)
220 if not raster.indexes["x"].is_monotonic_increasing:
221 raster = raster.isel(x=flip)
222 if not raster.indexes["y"].is_monotonic_decreasing:
223 raster = raster.isel(y=flip)
225 # Plot raster
226 ax1 = ax.imshow(raster, cmap=cmap, norm=norm, zorder=1, **settings_raster)
228 # Set ax imits
229 ax.set_xlim(xmin, xmax)
230 ax.set_ylim(ymin, ymax)
232 # Make triangles white if data is not larger/smaller than legend_levels-range
233 if whiten_triangles:
234 if float(raster.max().compute()) < levels[-1]:
235 ax1.cmap.set_over("#FFFFFF")
236 if float(raster.min().compute()) > levels[0]:
237 ax1.cmap.set_under("#FFFFFF")
239 # Add colorbar
240 divider = make_axes_locatable(ax)
241 cbar_ax = divider.append_axes("right", size="5%", pad="5%")
242 settings_cbar.pop("ticklabels", None)
243 cbar = fig.colorbar(ax1, cax=cbar_ax, **settings_cbar)
245 # Add overlays
246 for i, overlay in enumerate(overlays):
247 tmp = overlay.copy()
248 gdf = tmp.pop("gdf")
249 gdf.plot(ax=ax, zorder=2 + i, **tmp)
251 # Add basemap, if basemap is neither None nor False
252 if add_basemap:
253 crs = "EPSG:28992" # default Amersfoort/RDnew
254 try:
255 crs = _crs2string(kwargs_basemap.pop("crs"))
256 except (KeyError, AttributeError):
257 try:
258 crs = _crs2string(raster.attrs["crs"])
259 except (KeyError, AttributeError):
260 for overlay in overlays:
261 if "crs" in overlay["gdf"]:
262 crs = _crs2string(overlay["gdf"].crs)
263 break
265 if isinstance(basemap, bool):
266 source = ctx.providers["CartoDB"]["Positron"]
267 else:
268 source = basemap
270 ctx.add_basemap(ax=ax, source=source, crs=crs, **kwargs_basemap)
272 # Return
273 if return_cbar:
274 return fig, ax, cbar
275 else:
276 return fig, ax
279def _colorscale(a_yx, levels, cmap, quantile_colorscale):
280 """
281 This is an attempt to automatically create somewhat robust color scales.
283 Parameters
284 ----------
285 a_yx : xr.DataArray
286 2D DataArray with only dimensions ("y", "x")
287 levels : integer, np.ndarray or None
288 Number of levels (if integer), or level boundaries (if ndarray)
289 quantile_colorscale : boolean
290 Whether to create a colorscale based on quantile classification
292 Returns
293 -------
294 norm : matplotlib.colors.BoundaryNorm
295 cmap : matplotlib.colors.ListedColormap
296 """
297 # This is all an attempt at a somewhat robust colorscale handling
298 if levels is None: # Nothing given, default to 25 colors
299 levels = 25
300 if isinstance(levels, int):
301 nlevels = levels
302 if quantile_colorscale:
303 levels = np.unique(np.nanpercentile(a_yx.values, np.linspace(0, 100, 101)))
304 if len(levels) > nlevels:
305 # Decrease the number of levels
306 # Pretty rough approach, but should be sufficient
307 x = np.linspace(0.0, 100.0, nlevels)
308 xp = np.linspace(0.0, 100.0, len(levels))
309 yp = levels
310 levels = np.interp(x, xp, yp)
311 else: # Can't make more levels out of only a few quantiles
312 nlevels = len(levels)
313 else: # Go start to end
314 vmin = float(a_yx.min())
315 vmax = float(a_yx.max())
316 levels = np.linspace(vmin, vmax, nlevels)
317 elif isinstance(levels, (np.ndarray, list, tuple)): # Pre-defined by user
318 nlevels = len(levels)
319 else:
320 raise ValueError("levels argument should be None, an integer, or an array.")
322 if nlevels < 3: # let matplotlib take care of it
323 norm = None
324 else:
325 norm = matplotlib.colors.BoundaryNorm(boundaries=levels, ncolors=nlevels)
327 # Interpolate colormap to nlevels
328 if isinstance(cmap, str):
329 cmap = matplotlib.cm.get_cmap(cmap)
330 # cmap is a callable object
331 cmap = matplotlib.colors.ListedColormap(cmap(np.linspace(0.0, 1.0, nlevels)))
333 return norm, cmap
336def _imshow_xy(
337 a_yx, fname, title, cmap, overlays, quantile_colorscale, figsize, settings, levels
338):
339 fig, ax = plt.subplots(figsize=figsize)
340 norm, cmap = _colorscale(a_yx, levels, cmap, quantile_colorscale)
341 ax1 = ax.imshow(a_yx, cmap=cmap, norm=norm, **settings)
342 for overlay in overlays:
343 tmp = overlay.copy()
344 gdf = tmp.pop("gdf")
345 gdf.plot(ax=ax, **tmp)
347 divider = make_axes_locatable(ax)
348 cbar_ax = divider.append_axes("right", size="5%", pad="5%")
349 fig.colorbar(ax1, cax=cbar_ax)
351 ax.set_title(title)
352 plt.savefig(fname, dpi=200, bbox_inches="tight")
353 plt.close()
356def format_time(time):
357 if isinstance(time, np.datetime64):
358 # The following line is because numpy.datetime64[ns] does not
359 # support converting to datetime, but returns an integer instead.
360 # This solution is 20 times faster than using pd.to_datetime()
361 return time.astype("datetime64[us]").item().strftime("%Y%m%d%H%M%S")
362 else:
363 return time.strftime("%Y%m%d%H%M%S")
366def imshow_topview(
367 da,
368 name,
369 directory=".",
370 cmap="viridis",
371 overlays=[],
372 quantile_colorscale=True,
373 figsize=(8, 8),
374 levels=None,
375):
376 """
377 Automatically colors by quantile.
379 Dumps PNGs into directory of choice.
380 """
381 directory = pathlib.Path(directory)
382 directory.mkdir(parents=True, exist_ok=True)
383 if "x" not in da.dims or "y" not in da.dims:
384 raise ValueError("DataArray must have dims x and y.")
385 directory = pathlib.Path(directory)
386 _, xmin, xmax, _, ymin, ymax = imod.util.spatial.spatial_reference(da)
387 settings = {"interpolation": "nearest", "extent": [xmin, xmax, ymin, ymax]}
388 extradims = [dim for dim in da.dims if dim not in ("x", "y")]
390 if len(extradims) == 0:
391 fname = directory / f"{name}.png"
392 _imshow_xy(
393 da,
394 fname,
395 name,
396 cmap,
397 overlays,
398 quantile_colorscale,
399 figsize,
400 settings,
401 levels,
402 )
403 else:
404 stacked = da.stack(idf=extradims)
405 for coordvals, a_yx in list(stacked.groupby("idf")):
406 if a_yx.isnull().all():
407 continue
409 fname_parts = []
410 title_parts = []
411 for key, coordval in zip(extradims, coordvals):
412 title_parts.append(f"{key}: {coordval}")
413 if key == "time":
414 coordval = format_time(coordval)
415 fname_parts.append(f"{key}{coordval}")
417 fname_parts = "_".join(fname_parts)
418 title_parts = ", ".join(title_parts)
419 fname = directory / f"{name}_{fname_parts}.png"
420 title = f"{name}, {title_parts}"
421 _imshow_xy(
422 a_yx,
423 fname,
424 title,
425 cmap,
426 overlays,
427 quantile_colorscale,
428 figsize,
429 settings,
430 levels,
431 )