Mercurial > repos > ecology > xarray_coords_info

comparison xarray_mapplot.py @ 0:fea8a53f8099 draft

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