Mercurial > repos > ecology > timeseries_extraction
comparison xarray_mapplot.py @ 0:810820a0d45c draft default tip
planemo upload for repository https://github.com/galaxyecology/tools-ecology/tree/master/tools/data_manipulation/xarray/ commit fd8ad4d97db7b1fd3876ff63e14280474e06fdf7
author | ecology |
---|---|
date | Sun, 31 Jul 2022 21:23:21 +0000 |
parents | |
children |
comparison
equal
deleted
inserted
replaced
-1:000000000000 | 0:810820a0d45c |
---|---|
1 #!/usr/bin/env python3 | |
2 # | |
3 # | |
4 # usage: xarray_mapplot.py [-h] [--proj PROJ] | |
5 # [--cmap CMAP] | |
6 # [--output OUTPUT] | |
7 # [--time TIMES] | |
8 # [--nrow NROW] | |
9 # [--ncol NCOL] | |
10 # [--title title] | |
11 # [--latitude LATITUDE] | |
12 # [--longitude LONGITUDE] | |
13 # [--land ALPHA-LAND] | |
14 # [--ocean ALPHA-OCEAN] | |
15 # [--coastline ALPHA-COASTLINE] | |
16 # [--borders ALPHA-BORDERS] | |
17 # [--xlim "x1,x2"] | |
18 # [--ylim "y1,y2"] | |
19 # [--range "valmin,valmax"] | |
20 # [--threshold VAL] | |
21 # [--label label-colorbar] | |
22 # [--config config-file] | |
23 # [--shift] | |
24 # [-v] | |
25 # input varname | |
26 # | |
27 # positional arguments: | |
28 # input input filename with geographical coordinates (netCDF | |
29 # format) | |
30 # varname Specify which variable to plot (case sensitive) | |
31 # | |
32 # optional arguments: | |
33 # -h, --help show this help message and exit | |
34 # --proj PROJ Specify the projection on which we draw | |
35 # --cmap CMAP Specify which colormap to use for plotting | |
36 # --output OUTPUT output filename to store resulting image (png format) | |
37 # --time TIMES time index from the file for multiple plots ("0 1 2 3") | |
38 # --title plot or subplot title | |
39 # --latitude variable name for latitude | |
40 # --longitude variable name for longitude | |
41 # --land add land on plot with alpha value [0-1] | |
42 # --ocean add oceans on plot with alpha value [0-1] | |
43 # --coastline add coastline with alpha value [0-1] | |
44 # --borders add country borders with alpha value [0-1] | |
45 # --xlim limited geographical area longitudes "x1,x2" | |
46 # --ylim limited geographical area latitudes "y1,y2" | |
47 # --range "valmin,valmax" for plotting | |
48 # --threshold do not plot values below threshold | |
49 # --label set a label for colormap | |
50 # --config plotting parameters are passed via a config file | |
51 # (overwrite other plotting options) | |
52 # --shift shift longitudes if specified | |
53 # -v, --verbose switch on verbose mode | |
54 # | |
55 | |
56 import argparse | |
57 import ast | |
58 import warnings | |
59 from pathlib import Path | |
60 | |
61 import cartopy.crs as ccrs | |
62 import cartopy.feature as feature | |
63 | |
64 from cmcrameri import cm | |
65 | |
66 import matplotlib as mpl | |
67 mpl.use('Agg') | |
68 from matplotlib import pyplot # noqa: I202,E402 | |
69 | |
70 import xarray as xr # noqa: E402 | |
71 | |
72 | |
73 class MapPlotXr (): | |
74 def __init__(self, input, varname, output, verbose=False, | |
75 config_file="", proj="", shift=False): | |
76 | |
77 li = list(input.split(",")) | |
78 if len(li) > 1: | |
79 self.input = li | |
80 else: | |
81 self.input = input | |
82 | |
83 if proj != "" and proj is not None and Path(proj).exists(): | |
84 f = open(proj) | |
85 sdict = ''.join( | |
86 f.read().replace("\n", "").split('{')[1].split('}')[0] | |
87 ) | |
88 self.proj = '{' + sdict.strip() + '}' | |
89 else: | |
90 self.proj = None | |
91 self.varname = varname | |
92 self.shift = shift | |
93 self.xylim_supported = False | |
94 self.colorbar = True | |
95 if output is None: | |
96 if type(self.input) is list: | |
97 self.output = Path(self.input[0]).stem + '.png' | |
98 else: | |
99 self.output = Path(self.input).stem + '.png' | |
100 else: | |
101 self.output = output | |
102 self.verbose = verbose | |
103 self.label = {} | |
104 self.time = [] | |
105 self.xlim = [] | |
106 self.ylim = [] | |
107 self.range = [] | |
108 self.latitude = "latitude" | |
109 self.longitude = "longitude" | |
110 self.land = 0 | |
111 self.ocean = 0 | |
112 self.coastline = 0 | |
113 self.borders = 0 | |
114 self.cmap = "coolwarm" | |
115 self.threshold = "" | |
116 self.title = "" | |
117 | |
118 if config_file != "" and config_file is not None: | |
119 with open(config_file) as f: | |
120 sdict = ''.join( | |
121 f.read().replace("\n", "").split('{')[1].split('}')[0] | |
122 ) | |
123 tmp = ast.literal_eval('{' + sdict.strip() + '}') | |
124 for key in tmp: | |
125 if key == 'time': | |
126 time = tmp[key] | |
127 self.time = list(map(int, time.split(","))) | |
128 if key == 'cmap': | |
129 self.get_cmap(tmp[key]) | |
130 if key == 'latitude': | |
131 self.latitude = tmp[key] | |
132 if key == 'longitude': | |
133 self.longitude = tmp[key] | |
134 if key == 'land': | |
135 self.land = float(tmp[key]) | |
136 if key == 'ocean': | |
137 self.ocean = float(tmp[key]) | |
138 if key == 'coastline': | |
139 self.coastline = float(tmp[key]) | |
140 if key == 'borders': | |
141 self.borders = float(tmp[key]) | |
142 if key == 'xlim': | |
143 xlim = tmp[key] | |
144 self.xlim = list(map(float, xlim.split(","))) | |
145 if key == 'ylim': | |
146 ylim = tmp[key] | |
147 self.ylim = list(map(float, ylim.split(","))) | |
148 if key == 'range': | |
149 range_values = tmp[key] | |
150 self.range = list(map(float, range_values.split(","))) | |
151 if key == 'threshold': | |
152 self.threshold = float(tmp[key]) | |
153 if key == 'label': | |
154 self.label['label'] = tmp[key] | |
155 if key == 'title': | |
156 self.title = tmp[key] | |
157 | |
158 if type(self.input) is list: | |
159 self.dset = xr.open_mfdataset(self.input, use_cftime=True) | |
160 else: | |
161 self.dset = xr.open_dataset(self.input, use_cftime=True) | |
162 | |
163 if verbose: | |
164 print("input: ", self.input) | |
165 print("proj: ", self.proj) | |
166 print("varname: ", self.varname) | |
167 print("time: ", self.time) | |
168 print("minval, maxval: ", self.range) | |
169 print("title: ", self.title) | |
170 print("output: ", self.output) | |
171 print("label: ", self.label) | |
172 print("shift: ", self.shift) | |
173 print("ocean: ", self.ocean) | |
174 print("land: ", self.land) | |
175 print("coastline: ", self.coastline) | |
176 print("borders: ", self.borders) | |
177 print("latitude: ", self.latitude) | |
178 print("longitude: ", self.longitude) | |
179 print("xlim: ", self.xlim) | |
180 print("ylim: ", self.ylim) | |
181 | |
182 def get_cmap(self, cmap): | |
183 if cmap[0:3] == 'cm.': | |
184 self.cmap = cm.__dict__[cmap[3:]] | |
185 else: | |
186 self.cmap = cmap | |
187 | |
188 def projection(self): | |
189 if self.proj is None: | |
190 return ccrs.PlateCarree() | |
191 | |
192 proj_dict = ast.literal_eval(self.proj) | |
193 user_proj = proj_dict.pop("proj") | |
194 if user_proj == 'PlateCarree': | |
195 self.xylim_supported = True | |
196 return ccrs.PlateCarree(**proj_dict) | |
197 elif user_proj == 'AlbersEqualArea': | |
198 return ccrs.AlbersEqualArea(**proj_dict) | |
199 elif user_proj == 'AzimuthalEquidistant': | |
200 return ccrs.AzimuthalEquidistant(**proj_dict) | |
201 elif user_proj == 'EquidistantConic': | |
202 return ccrs.EquidistantConic(**proj_dict) | |
203 elif user_proj == 'LambertConformal': | |
204 return ccrs.LambertConformal(**proj_dict) | |
205 elif user_proj == 'LambertCylindrical': | |
206 return ccrs.LambertCylindrical(**proj_dict) | |
207 elif user_proj == 'Mercator': | |
208 return ccrs.Mercator(**proj_dict) | |
209 elif user_proj == 'Miller': | |
210 return ccrs.Miller(**proj_dict) | |
211 elif user_proj == 'Mollweide': | |
212 return ccrs.Mollweide(**proj_dict) | |
213 elif user_proj == 'Orthographic': | |
214 return ccrs.Orthographic(**proj_dict) | |
215 elif user_proj == 'Robinson': | |
216 return ccrs.Robinson(**proj_dict) | |
217 elif user_proj == 'Sinusoidal': | |
218 return ccrs.Sinusoidal(**proj_dict) | |
219 elif user_proj == 'Stereographic': | |
220 return ccrs.Stereographic(**proj_dict) | |
221 elif user_proj == 'TransverseMercator': | |
222 return ccrs.TransverseMercator(**proj_dict) | |
223 elif user_proj == 'UTM': | |
224 return ccrs.UTM(**proj_dict) | |
225 elif user_proj == 'InterruptedGoodeHomolosine': | |
226 return ccrs.InterruptedGoodeHomolosine(**proj_dict) | |
227 elif user_proj == 'RotatedPole': | |
228 return ccrs.RotatedPole(**proj_dict) | |
229 elif user_proj == 'OSGB': | |
230 self.xylim_supported = False | |
231 return ccrs.OSGB(**proj_dict) | |
232 elif user_proj == 'EuroPP': | |
233 self.xylim_supported = False | |
234 return ccrs.EuroPP(**proj_dict) | |
235 elif user_proj == 'Geostationary': | |
236 return ccrs.Geostationary(**proj_dict) | |
237 elif user_proj == 'NearsidePerspective': | |
238 return ccrs.NearsidePerspective(**proj_dict) | |
239 elif user_proj == 'EckertI': | |
240 return ccrs.EckertI(**proj_dict) | |
241 elif user_proj == 'EckertII': | |
242 return ccrs.EckertII(**proj_dict) | |
243 elif user_proj == 'EckertIII': | |
244 return ccrs.EckertIII(**proj_dict) | |
245 elif user_proj == 'EckertIV': | |
246 return ccrs.EckertIV(**proj_dict) | |
247 elif user_proj == 'EckertV': | |
248 return ccrs.EckertV(**proj_dict) | |
249 elif user_proj == 'EckertVI': | |
250 return ccrs.EckertVI(**proj_dict) | |
251 elif user_proj == 'EqualEarth': | |
252 return ccrs.EqualEarth(**proj_dict) | |
253 elif user_proj == 'Gnomonic': | |
254 return ccrs.Gnomonic(**proj_dict) | |
255 elif user_proj == 'LambertAzimuthalEqualArea': | |
256 return ccrs.LambertAzimuthalEqualArea(**proj_dict) | |
257 elif user_proj == 'NorthPolarStereo': | |
258 return ccrs.NorthPolarStereo(**proj_dict) | |
259 elif user_proj == 'OSNI': | |
260 return ccrs.OSNI(**proj_dict) | |
261 elif user_proj == 'SouthPolarStereo': | |
262 return ccrs.SouthPolarStereo(**proj_dict) | |
263 | |
264 def plot(self, ts=None): | |
265 if self.shift: | |
266 if self.longitude == 'longitude': | |
267 self.dset = self.dset.assign_coords( | |
268 longitude=((( | |
269 self.dset[self.longitude] | |
270 + 180) % 360) - 180)) | |
271 elif self.longitude == 'lon': | |
272 self.dset = self.dset.assign_coords( | |
273 lon=(((self.dset[self.longitude] | |
274 + 180) % 360) - 180)) | |
275 | |
276 pyplot.figure(1, figsize=[20, 10]) | |
277 | |
278 # Set the projection to use for plotting | |
279 ax = pyplot.subplot(1, 1, 1, projection=self.projection()) | |
280 if self.land: | |
281 ax.add_feature(feature.LAND, alpha=self.land) | |
282 | |
283 if self.ocean: | |
284 ax.add_feature(feature.OCEAN, alpha=self.ocean) | |
285 if self.coastline: | |
286 ax.coastlines(resolution='10m', alpha=self.coastline) | |
287 if self.borders: | |
288 ax.add_feature(feature.BORDERS, linestyle=':', alpha=self.borders) | |
289 | |
290 if self.xlim: | |
291 min_lon = min(self.xlim[0], self.xlim[1]) | |
292 max_lon = max(self.xlim[0], self.xlim[1]) | |
293 else: | |
294 min_lon = self.dset[self.longitude].min() | |
295 max_lon = self.dset[self.longitude].max() | |
296 | |
297 if self.ylim: | |
298 min_lat = min(self.ylim[0], self.ylim[1]) | |
299 max_lat = max(self.ylim[0], self.ylim[1]) | |
300 else: | |
301 min_lat = self.dset[self.latitude].min() | |
302 max_lat = self.dset[self.latitude].max() | |
303 | |
304 if self.xylim_supported: | |
305 pyplot.xlim(min_lon, max_lon) | |
306 pyplot.ylim(min_lat, max_lat) | |
307 | |
308 # Fix extent | |
309 if self.threshold == "" or self.threshold is None: | |
310 threshold = self.dset[self.varname].min() | |
311 else: | |
312 threshold = float(self.threshold) | |
313 | |
314 if self.range == []: | |
315 minval = self.dset[self.varname].min() | |
316 maxval = self.dset[self.varname].max() | |
317 else: | |
318 minval = self.range[0] | |
319 maxval = self.range[1] | |
320 | |
321 if self.verbose: | |
322 print("minval: ", minval) | |
323 print("maxval: ", maxval) | |
324 | |
325 # pass extent with vmin and vmax parameters | |
326 proj_t = ccrs.PlateCarree() | |
327 if ts is None: | |
328 self.dset.where( | |
329 self.dset[self.varname] > threshold | |
330 )[self.varname].plot(ax=ax, | |
331 vmin=minval, | |
332 vmax=maxval, | |
333 transform=proj_t, | |
334 cmap=self.cmap, | |
335 cbar_kwargs=self.label | |
336 ) | |
337 if self.title != "" and self.title is not None: | |
338 pyplot.title(self.title) | |
339 pyplot.savefig(self.output) | |
340 else: | |
341 if self.colorbar: | |
342 self.dset.where( | |
343 self.dset[self.varname] > threshold | |
344 )[self.varname].isel(time=ts).plot(ax=ax, | |
345 vmin=minval, | |
346 vmax=maxval, | |
347 transform=proj_t, | |
348 cmap=self.cmap, | |
349 cbar_kwargs=self.label | |
350 ) | |
351 else: | |
352 self.dset.where( | |
353 self.dset[self.varname] > minval | |
354 )[self.varname].isel(time=ts).plot(ax=ax, | |
355 vmin=minval, | |
356 vmax=maxval, | |
357 transform=proj_t, | |
358 cmap=self.cmap, | |
359 add_colorbar=False) | |
360 if self.title != "" and self.title is not None: | |
361 pyplot.title(self.title + "(time = " + str(ts) + ')') | |
362 pyplot.savefig(self.output[:-4] + "_time" + str(ts) + | |
363 self.output[-4:]) # assume png format | |
364 | |
365 | |
366 if __name__ == '__main__': | |
367 warnings.filterwarnings("ignore") | |
368 parser = argparse.ArgumentParser() | |
369 parser.add_argument( | |
370 'input', | |
371 help='input filename with geographical coordinates (netCDF format)' | |
372 ) | |
373 parser.add_argument( | |
374 '--proj', | |
375 help='Config file with the projection on which we draw' | |
376 ) | |
377 parser.add_argument( | |
378 'varname', | |
379 help='Specify which variable to plot (case sensitive)' | |
380 ) | |
381 parser.add_argument( | |
382 '--output', | |
383 help='output filename to store resulting image (png format)' | |
384 ) | |
385 parser.add_argument( | |
386 '--config', | |
387 help='pass plotting parameters via a config file' | |
388 ) | |
389 parser.add_argument( | |
390 '--shift', | |
391 help='shift longitudes if specified', | |
392 action="store_true" | |
393 ) | |
394 parser.add_argument( | |
395 "-v", "--verbose", | |
396 help="switch on verbose mode", | |
397 action="store_true") | |
398 args = parser.parse_args() | |
399 | |
400 dset = MapPlotXr(input=args.input, varname=args.varname, | |
401 output=args.output, verbose=args.verbose, | |
402 config_file=args.config, proj=args.proj, | |
403 shift=args.shift) | |
404 | |
405 if dset.time == []: | |
406 dset.plot() | |
407 else: | |
408 for t in dset.time: | |
409 dset.plot(t) | |
410 dset.shift = False # only shift once | |
411 dset.colorbar = True |