Coverage for C:\src\imod-python\imod\visualize\spatial.py: 47%

179 statements  

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

1import pathlib 

2 

3import matplotlib 

4import matplotlib.pyplot as plt 

5import numpy as np 

6import pandas as pd 

7from mpl_toolkits.axes_grid1 import make_axes_locatable 

8 

9import imod 

10from imod.util.imports import MissingOptionalModule 

11from imod.visualize import common 

12 

13try: 

14 import contextily as ctx 

15except ImportError: 

16 ctx = MissingOptionalModule("contextily") 

17 

18 

19def read_imod_legend(path): 

20 """ 

21 Parameters 

22 ---------- 

23 path : str 

24 Path to iMOD .leg file. 

25 

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 """ 

32 

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 ) 

43 

44 # Try both comma and whitespace separated 

45 try: 

46 legend = _read(delim_whitespace=False) 

47 except ValueError: 

48 legend = _read(delim_whitespace=True) 

49 

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:] 

56 

57 return colors, levels 

58 

59 

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 

71 

72 

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. 

89 

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. 

96 

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`). 

99 

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. 

103 

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. 

107 

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). 

122 

123 *Requires contextily* 

124 

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 

144 

145 Returns 

146 ------- 

147 fig : matplotlib.figure ax : matplotlib.ax if return_cbar == True: cbar : 

148 matplotlib.Colorbar 

149 

150 Examples 

151 -------- 

152 Plot with an overlay: 

153 

154 >>> overlays = [{"gdf": geodataframe, "edgecolor": "black", "facecolor": "None"}] 

155 >>> imod.visualize.plot_map(raster, colors, levels, overlays) 

156 

157 Label the colorbar: 

158 

159 >>> imod.visualize.plot_map(raster, colors, levels, kwargs_colorbar={"label":"Head aquifer (m)"}) 

160 

161 Plot with a basemap: 

162 

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}) 

166 

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 

173 

174 # Read legend settings 

175 cmap, norm = common._cmapnorm_from_colorslevels(colors, levels) 

176 

177 # Get extent 

178 _, xmin, xmax, _, ymin, ymax = imod.util.spatial.spatial_reference(raster) 

179 

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) 

187 

188 # cbar kwargs 

189 settings_cbar = {"ticks": levels, "extend": "both"} 

190 

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 

199 

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) 

204 

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() 

216 

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) 

224 

225 # Plot raster 

226 ax1 = ax.imshow(raster, cmap=cmap, norm=norm, zorder=1, **settings_raster) 

227 

228 # Set ax imits 

229 ax.set_xlim(xmin, xmax) 

230 ax.set_ylim(ymin, ymax) 

231 

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") 

238 

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) 

244 

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) 

250 

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 

264 

265 if isinstance(basemap, bool): 

266 source = ctx.providers["CartoDB"]["Positron"] 

267 else: 

268 source = basemap 

269 

270 ctx.add_basemap(ax=ax, source=source, crs=crs, **kwargs_basemap) 

271 

272 # Return 

273 if return_cbar: 

274 return fig, ax, cbar 

275 else: 

276 return fig, ax 

277 

278 

279def _colorscale(a_yx, levels, cmap, quantile_colorscale): 

280 """ 

281 This is an attempt to automatically create somewhat robust color scales. 

282 

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 

291 

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.") 

321 

322 if nlevels < 3: # let matplotlib take care of it 

323 norm = None 

324 else: 

325 norm = matplotlib.colors.BoundaryNorm(boundaries=levels, ncolors=nlevels) 

326 

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))) 

332 

333 return norm, cmap 

334 

335 

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) 

346 

347 divider = make_axes_locatable(ax) 

348 cbar_ax = divider.append_axes("right", size="5%", pad="5%") 

349 fig.colorbar(ax1, cax=cbar_ax) 

350 

351 ax.set_title(title) 

352 plt.savefig(fname, dpi=200, bbox_inches="tight") 

353 plt.close() 

354 

355 

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") 

364 

365 

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. 

378 

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")] 

389 

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 

408 

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}") 

416 

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 )