Mercurial > repos > ecology > timeseries_extraction

comparison xarray_mapplot.py @ 0:dd19259f7da5 draft default tip

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 fd8ad4d97db7b1fd3876ff63e14280474e06fdf7
author ecology
date Sun, 31 Jul 2022 21:18:44 +0000
parents
children
comparison
equal deleted inserted replaced
-1:000000000000 0:dd19259f7da5
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