source: flex_extract.git/python/plot_retrieved.py @ 70a0bec

ctbtodev
Last change on this file since 70a0bec was 54a8a01, checked in by Anne Philipp <anne.philipp@…>, 6 years ago

restructuring, documentations and bug fixes

  • Property mode set to 100755
File size: 22.9 KB
Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3#*******************************************************************************
4# @Author: Leopold Haimberger (University of Vienna)
5#
6# @Date: November 2015
7#
8# @Change History:
9#
10#    February 2018 - Anne Philipp (University of Vienna):
11#        - applied PEP8 style guide
12#        - added documentation
13#        - created function main and moved the two function calls for
14#          arguments and plotting into it
15#        - added function get_basics to extract the boundary conditions
16#          of the data fields from the first grib file it gets.
17#
18# @License:
19#    (C) Copyright 2015-2018.
20#
21#    This software is licensed under the terms of the Apache Licence Version 2.0
22#    which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
23#
24# @Program Functionality:
25#    Simple tool for creating maps and time series of retrieved fields.
26#
27# @Program Content:
28#    - main
29#    - get_basics
30#    - get_files_per_date
31#    - plot_retrieved
32#    - plot_timeseries
33#    - plot_map
34#    - get_plot_args
35#
36#*******************************************************************************
37
38# ------------------------------------------------------------------------------
39# MODULES
40# ------------------------------------------------------------------------------
41import time
42import datetime
43import os
44import inspect
45import sys
46from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
47import matplotlib
48import matplotlib.pyplot as plt
49from mpl_toolkits.basemap import Basemap
50from eccodes import codes_grib_new_from_file, codes_get, codes_release, \
51                    codes_get_values
52import numpy as np
53
54# software specific classes and modules from flex_extract
55from ControlFile import ControlFile
56from UioFiles import UioFiles
57
58# add path to pythonpath so that python finds its buddies
59LOCAL_PYTHON_PATH = os.path.dirname(os.path.abspath(
60    inspect.getfile(inspect.currentframe())))
61if LOCAL_PYTHON_PATH not in sys.path:
62    sys.path.append(LOCAL_PYTHON_PATH)
63
64font = {'family': 'monospace', 'size': 12}
65matplotlib.rcParams['xtick.major.pad'] = '20'
66matplotlib.rc('font', **font)
67# ------------------------------------------------------------------------------
68# FUNCTION
69# ------------------------------------------------------------------------------
70def main():
71    '''
72    @Description:
73        If plot_retrieved is called from command line, this function controls
74        the program flow and calls the argumentparser function and
75        the plot_retrieved function for plotting the retrieved GRIB data.
76
77    @Input:
78        <nothing>
79
80    @Return:
81        <nothing>
82    '''
83    args, c = get_plot_args()
84    plot_retrieved(c)
85
86    return
87
88def get_basics(ifile, verb=False):
89    """
90    @Description:
91        An example grib file will be opened and basic information will
92        be extracted. These information are important for later use and the
93        initialization of numpy arrays for data storing.
94
95    @Input:
96        ifile: string
97            Contains the full absolute path to the ECMWF grib file.
98
99        verb (opt): bool
100            Is True if there should be extra output in verbose mode.
101            Default value is False.
102
103    @Return:
104        data: dict
105            Contains basic informations of the ECMWF grib files, e.g.
106            'Ni', 'Nj', 'latitudeOfFirstGridPointInDegrees',
107            'longitudeOfFirstGridPointInDegrees',
108            'latitudeOfLastGridPointInDegrees',
109            'longitudeOfLastGridPointInDegrees',
110            'jDirectionIncrementInDegrees',
111            'iDirectionIncrementInDegrees'
112    """
113
114    data = {}
115
116    # --- open file ---
117    print("Opening file for getting information data --- %s" %
118          os.path.basename(ifile))
119
120    with open(ifile) as f:
121
122        # load first message from file
123        gid = codes_grib_new_from_file(f)
124
125        # information needed from grib message
126        keys = ['Ni',
127                'Nj',
128                'latitudeOfFirstGridPointInDegrees',
129                'longitudeOfFirstGridPointInDegrees',
130                'latitudeOfLastGridPointInDegrees',
131                'longitudeOfLastGridPointInDegrees',
132                'jDirectionIncrementInDegrees',
133                'iDirectionIncrementInDegrees']
134
135        if verb:
136            print '\nInformations are: '
137        for key in keys:
138            # Get the value of the key in a grib message.
139            data[key] = codes_get(gid, key)
140            if verb:
141                print "%s = %s" % (key, data[key])
142        if verb:
143            print '\n'
144
145        # Free the memory for the message referred as gribid.
146        codes_release(gid)
147
148    return data
149
150def get_files_per_date(files, datelist):
151    '''
152    @Description:
153        The filenames contain dates which are used to select a list
154        of files for a specific time period specified in datelist.
155
156    @Input:
157        files: instance of UioFiles
158            For description see class documentation.
159            It contains the attribute "files" which is a list of pathes
160            to filenames.
161
162        datelist: list of datetimes
163            Contains the list of dates which should be processed for plotting.
164
165    @Return:
166        filelist: list of strings
167            Contains the selected files for the time period.
168    '''
169
170    filelist = []
171    for filename in files:
172        filedate = filename[-8:]
173        ddate = datetime.datetime.strptime(filedate, '%y%m%d%H')
174        if ddate in datelist:
175            filelist.append(filename)
176
177    return filelist
178
179def plot_retrieved(c):
180    '''
181    @Description:
182        Reads GRIB data from a specified time period, a list of levels
183        and a specified list of parameter.
184
185    @Input:
186        c: instance of class ControlFile
187            Contains all necessary information of a CONTROL file. The parameters
188            are: DAY1, DAY2, DTIME, MAXSTEP, TYPE, TIME, STEP, CLASS, STREAM,
189            NUMBER, EXPVER, GRID, LEFT, LOWER, UPPER, RIGHT, LEVEL, LEVELIST,
190            RESOL, GAUSS, ACCURACY, OMEGA, OMEGADIFF, ETA, ETADIFF, DPDETA,
191            SMOOTH, FORMAT, ADDPAR, WRF, CWC, PREFIX, ECSTORAGE, ECTRANS,
192            ECFSDIR, MAILOPS, MAILFAIL, GRIB2FLEXPART, DEBUG, INPUTDIR,
193            OUTPUTDIR, FLEXPART_ROOT_SCRIPTS
194            For more information about format and content of the parameter see
195            documentation.
196
197    @Return:
198        <nothing>
199    '''
200    start = datetime.datetime.strptime(c.start_date, '%Y%m%d%H')
201    end = datetime.datetime.strptime(c.end_date, '%Y%m%d%H')
202
203    # create datelist between start and end date
204    datelist = [start] # initialise datelist with first date
205    run_date = start
206    while run_date < end:
207        run_date += datetime.timedelta(hours=int(c.dtime))
208        datelist.append(run_date)
209
210    print 'datelist: ', datelist
211
212    c.paramIds = np.asarray(c.paramIds, dtype='int')
213    c.levels = np.asarray(c.levels, dtype='int')
214    c.area = np.asarray(c.area)
215
216    files = UioFiles(c.inputdir, c.prefix+'*')
217    ifiles = get_files_per_date(files.files, datelist)
218    ifiles.sort()
219
220    gdict = get_basics(ifiles[0], verb=False)
221
222    fdict = dict()
223    fmeta = dict()
224    fstamp = dict()
225    for p in c.paramIds:
226        for l in c.levels:
227            key = '{:0>3}_{:0>3}'.format(p, l)
228            fdict[key] = []
229            fmeta[key] = []
230            fstamp[key] = []
231
232    for filename in ifiles:
233        f = open(filename)
234        print "Opening file for reading data --- %s" % filename
235        fdate = datetime.datetime.strptime(filename[-8:], "%y%m%d%H")
236
237        # Load in memory a grib message from a file.
238        gid = codes_grib_new_from_file(f)
239        while gid is not None:
240            gtype = codes_get(gid, 'type')
241            paramId = codes_get(gid, 'paramId')
242            parameterName = codes_get(gid, 'parameterName')
243            level = codes_get(gid, 'level')
244
245            if paramId in c.paramIds and level in c.levels:
246                key = '{:0>3}_{:0>3}'.format(paramId, level)
247                print 'key: ', key
248                if fstamp[key]:
249                    for i in range(len(fstamp[key])):
250                        if fdate < fstamp[key][i]:
251                            fstamp[key].insert(i, fdate)
252                            fmeta[key].insert(i, [paramId, parameterName, gtype,
253                                                  fdate, level])
254                            fdict[key].insert(i, np.flipud(np.reshape(
255                                codes_get_values(gid),
256                                [gdict['Nj'], gdict['Ni']])))
257                            break
258                        elif fdate > fstamp[key][i] and i == len(fstamp[key])-1:
259                            fstamp[key].append(fdate)
260                            fmeta[key].append([paramId, parameterName, gtype,
261                                               fdate, level])
262                            fdict[key].append(np.flipud(np.reshape(
263                                codes_get_values(gid),
264                                [gdict['Nj'], gdict['Ni']])))
265                            break
266                        elif fdate > fstamp[key][i] and i != len(fstamp[key])-1 \
267                             and fdate < fstamp[key][i+1]:
268                            fstamp[key].insert(i, fdate)
269                            fmeta[key].insert(i, [paramId, parameterName, gtype,
270                                                  fdate, level])
271                            fdict[key].insert(i, np.flipud(np.reshape(
272                                codes_get_values(gid),
273                                [gdict['Nj'], gdict['Ni']])))
274                            break
275                        else:
276                            pass
277                else:
278                    fstamp[key].append(fdate)
279                    fmeta[key].append((paramId, parameterName, gtype,
280                                       fdate, level))
281                    fdict[key].append(np.flipud(np.reshape(
282                        codes_get_values(gid), [gdict['Nj'], gdict['Ni']])))
283
284            codes_release(gid)
285
286            # Load in memory a grib message from a file.
287            gid = codes_grib_new_from_file(f)
288
289        f.close()
290
291    for k in fdict.iterkeys():
292        print 'fmeta: ', len(fmeta), fmeta
293        fml = fmeta[k]
294        fdl = fdict[k]
295        print 'fm1: ', len(fml), fml
296        for fd, fm in zip(fdl, fml):
297            print fm
298            ftitle = fm[1] + ' {} '.format(fm[-1]) + \
299                datetime.datetime.strftime(fm[3], '%Y%m%d%H')
300            pname = '_'.join(fm[1].split()) + '_{}_'.format(fm[-1]) + \
301                datetime.datetime.strftime(fm[3], '%Y%m%d%H')
302            plot_map(c, fd, fm, gdict, ftitle, pname, 'png')
303
304    for k in fdict.iterkeys():
305        fml = fmeta[k]
306        fdl = fdict[k]
307        fsl = fstamp[k]
308        if fdl:
309            fm = fml[0]
310            fd = fdl[0]
311            ftitle = fm[1] + ' {} '.format(fm[-1]) + \
312                datetime.datetime.strftime(fm[3], '%Y%m%d%H')
313            pname = '_'.join(fm[1].split()) + '_{}_'.format(fm[-1]) + \
314                datetime.datetime.strftime(fm[3], '%Y%m%d%H')
315            lat = -20.
316            lon = 20.
317            plot_timeseries(c, fdl, fml, fsl, lat, lon, gdict,
318                            ftitle, pname, 'png')
319
320    return
321
322def plot_timeseries(c, flist, fmetalist, ftimestamps, lat, lon,
323                    gdict, ftitle, filename, fending, show=False):
324    '''
325    @Description:
326        Creates a timeseries plot for a given lat/lon position.
327
328    @Input:
329        c: instance of class ControlFile
330            Contains all necessary information of a CONTROL file. The parameters
331            are: DAY1, DAY2, DTIME, MAXSTEP, TYPE, TIME, STEP, CLASS, STREAM,
332            NUMBER, EXPVER, GRID, LEFT, LOWER, UPPER, RIGHT, LEVEL, LEVELIST,
333            RESOL, GAUSS, ACCURACY, OMEGA, OMEGADIFF, ETA, ETADIFF, DPDETA,
334            SMOOTH, FORMAT, ADDPAR, WRF, CWC, PREFIX, ECSTORAGE, ECTRANS,
335            ECFSDIR, MAILOPS, MAILFAIL, GRIB2FLEXPART, DEBUG, INPUTDIR,
336            OUTPUTDIR, FLEXPART_ROOT_SCRIPTS
337            For more information about format and content of the parameter see
338            documentation.
339
340        flist: numpy array, 2d
341            The actual data values to be plotted from the grib messages.
342
343        fmetalist: list of strings
344            Contains some meta date for the data field to be plotted:
345            parameter id, parameter Name, grid type, datetime, level
346
347        ftimestamps: list of datetime
348            Contains the time stamps.
349
350        lat: float
351            The latitude for which the timeseries should be plotted.
352
353        lon: float
354            The longitude for which the timeseries should be plotted.
355
356        gdict: dict
357            Contains basic informations of the ECMWF grib files, e.g.
358            'Ni', 'Nj', 'latitudeOfFirstGridPointInDegrees',
359            'longitudeOfFirstGridPointInDegrees',
360            'latitudeOfLastGridPointInDegrees',
361            'longitudeOfLastGridPointInDegrees',
362            'jDirectionIncrementInDegrees',
363            'iDirectionIncrementInDegrees'
364
365        ftitle: string
366            The title of the timeseries.
367
368        filename: string
369            The time series is stored in a file with this name.
370
371        fending: string
372            Contains the type of plot, e.g. pdf or png
373
374        show: boolean
375            Decides if the plot is shown after plotting or hidden.
376
377    @Return:
378        <nothing>
379    '''
380    print 'plotting timeseries'
381
382    t1 = time.time()
383
384    #llx = gdict['longitudeOfFirstGridPointInDegrees']
385    #if llx > 180. :
386    #    llx -= 360.
387    #lly = gdict['latitudeOfLastGridPointInDegrees']
388    #dxout = gdict['iDirectionIncrementInDegrees']
389    #dyout = gdict['jDirectionIncrementInDegrees']
390    #urx = gdict['longitudeOfLastGridPointInDegrees']
391    #ury = gdict['latitudeOfFirstGridPointInDegrees']
392    #numxgrid = gdict['Ni']
393    #numygrid = gdict['Nj']
394
395    farr = np.asarray(flist)
396    #(time, lat, lon)
397
398    #lonindex = linspace(llx, urx, numxgrid)
399    #latindex = linspace(lly, ury, numygrid)
400
401    ts = farr[:, 0, 0]
402
403    fig = plt.figure(figsize=(12, 6.7))
404
405    plt.plot(ftimestamps, ts)
406    plt.title(ftitle)
407
408    plt.savefig(c.outputdir + '/' + filename + '_TS.' + fending,
409                facecolor=fig.get_facecolor(),
410                edgecolor='none',
411                format=fending)
412    print 'created ', c.outputdir + '/' + filename
413    if show:
414        plt.show()
415    fig.clf()
416    plt.close(fig)
417
418    print time.time() - t1, 's'
419
420    return
421
422def plot_map(c, flist, fmetalist, gdict, ftitle, filename, fending, show=False):
423    '''
424    @Description:
425        Creates a basemap plot with imshow for a given data field.
426
427    @Input:
428        c: instance of class ControlFile
429            Contains all necessary information of a CONTROL file. The parameters
430            are: DAY1, DAY2, DTIME, MAXSTEP, TYPE, TIME, STEP, CLASS, STREAM,
431            NUMBER, EXPVER, GRID, LEFT, LOWER, UPPER, RIGHT, LEVEL, LEVELIST,
432            RESOL, GAUSS, ACCURACY, OMEGA, OMEGADIFF, ETA, ETADIFF, DPDETA,
433            SMOOTH, FORMAT, ADDPAR, WRF, CWC, PREFIX, ECSTORAGE, ECTRANS,
434            ECFSDIR, MAILOPS, MAILFAIL, GRIB2FLEXPART, DEBUG, INPUTDIR,
435            OUTPUTDIR, FLEXPART_ROOT_SCRIPTS
436            For more information about format and content of the parameter see
437            documentation.
438
439        flist: numpy array, 2d
440            The actual data values to be plotted from the grib messages.
441
442        fmetalist: list of strings
443            Contains some meta date for the data field to be plotted:
444            parameter id, parameter Name, grid type, datetime, level
445
446        gdict: dict
447            Contains basic informations of the ECMWF grib files, e.g.
448            'Ni', 'Nj', 'latitudeOfFirstGridPointInDegrees',
449            'longitudeOfFirstGridPointInDegrees',
450            'latitudeOfLastGridPointInDegrees',
451            'longitudeOfLastGridPointInDegrees',
452            'jDirectionIncrementInDegrees',
453            'iDirectionIncrementInDegrees'
454
455        ftitle: string
456            The titel of the plot.
457
458        filename: string
459            The plot is stored in a file with this name.
460
461        fending: string
462            Contains the type of plot, e.g. pdf or png
463
464        show: boolean
465            Decides if the plot is shown after plotting or hidden.
466
467    @Return:
468        <nothing>
469    '''
470    print 'plotting map'
471
472    t1 = time.time()
473
474    fig = plt.figure(figsize=(12, 6.7))
475    #mbaxes = fig.add_axes([0.05, 0.15, 0.8, 0.7])
476
477    llx = gdict['longitudeOfFirstGridPointInDegrees'] #- 360
478    if llx > 180.:
479        llx -= 360.
480    lly = gdict['latitudeOfLastGridPointInDegrees']
481    #dxout = gdict['iDirectionIncrementInDegrees']
482    #dyout = gdict['jDirectionIncrementInDegrees']
483    urx = gdict['longitudeOfLastGridPointInDegrees']
484    ury = gdict['latitudeOfFirstGridPointInDegrees']
485    #numxgrid = gdict['Ni']
486    #numygrid = gdict['Nj']
487
488    m = Basemap(projection='cyl', llcrnrlon=llx, llcrnrlat=lly,
489                urcrnrlon=urx, urcrnrlat=ury, resolution='i')
490
491    #lw = 0.5
492    m.drawmapboundary()
493    #x = linspace(llx, urx, numxgrid)
494    #y = linspace(lly, ury, numygrid)
495
496    #xx, yy = m(*meshgrid(x, y))
497
498    #s = m.contourf(xx, yy, flist)
499
500    s = plt.imshow(flist.T,
501                   extent=(llx, urx, lly, ury),
502                   alpha=1.0,
503                   interpolation='nearest'
504                   #vmin=vn,
505                   #vmax=vx,
506                   #cmap=my_cmap,
507                   #levels=levels,
508                   #cmap=my_cmap,
509                   #norm=LogNorm(vn,vx)
510                  )
511
512    plt.title(ftitle, y=1.08)
513    cb = m.colorbar(s, location="right", pad="10%")
514    cb.set_label('label', size=14)
515
516    thickline = np.arange(lly, ury+1, 10.)
517    thinline = np.arange(lly, ury+1, 5.)
518    m.drawparallels(thickline,
519                    color='gray',
520                    dashes=[1, 1],
521                    linewidth=0.5,
522                    labels=[1, 1, 1, 1],
523                    xoffset=1.)
524    m.drawparallels(np.setdiff1d(thinline, thickline),
525                    color='lightgray',
526                    dashes=[1, 1],
527                    linewidth=0.5,
528                    labels=[0, 0, 0, 0])
529
530    thickline = np.arange(llx, urx+1, 10.)
531    thinline = np.arange(llx, urx+1, 5.)
532    m.drawmeridians(thickline,
533                    color='gray',
534                    dashes=[1, 1],
535                    linewidth=0.5,
536                    labels=[1, 1, 1, 1],
537                    yoffset=1.)
538    m.drawmeridians(np.setdiff1d(thinline, thickline),
539                    color='lightgray',
540                    dashes=[1, 1],
541                    linewidth=0.5,
542                    labels=[0, 0, 0, 0])
543
544    m.drawcoastlines()
545    m.drawcountries()
546
547    plt.savefig(c.outputdir + '/' + filename + '_MAP.' + fending,
548                facecolor=fig.get_facecolor(),
549                edgecolor='none',
550                format=fending)
551    print 'created ', c.outputdir + '/' + filename
552    if show:
553        plt.show()
554    fig.clf()
555    plt.close(fig)
556
557    print time.time() - t1, 's'
558
559    return
560
561def get_plot_args():
562    '''
563    @Description:
564        Assigns the command line arguments and reads CONTROL file
565        content. Apply default values for non mentioned arguments.
566
567    @Input:
568        <nothing>
569
570    @Return:
571        args: instance of ArgumentParser
572            Contains the commandline arguments from script/program call.
573
574        c: instance of class ControlFile
575            Contains all necessary information of a CONTROL file. The parameters
576            are: DAY1, DAY2, DTIME, MAXSTEP, TYPE, TIME, STEP, CLASS, STREAM,
577            NUMBER, EXPVER, GRID, LEFT, LOWER, UPPER, RIGHT, LEVEL, LEVELIST,
578            RESOL, GAUSS, ACCURACY, OMEGA, OMEGADIFF, ETA, ETADIFF, DPDETA,
579            SMOOTH, FORMAT, ADDPAR, WRF, CWC, PREFIX, ECSTORAGE, ECTRANS,
580            ECFSDIR, MAILOPS, MAILFAIL, GRIB2FLEXPART, DEBUG, INPUTDIR,
581            OUTPUTDIR, FLEXPART_ROOT_SCRIPTS
582            For more information about format and content of the parameter see
583            documentation.
584    '''
585    parser = ArgumentParser(description='Plot retrieved GRIB data from ' + \
586                            'ECMWF MARS archive',
587                            formatter_class=ArgumentDefaultsHelpFormatter)
588
589# the most important arguments
590    parser.add_argument("--start_date", dest="start_date",
591                        help="start date YYYYMMDD")
592    parser.add_argument("--end_date", dest="end_date",
593                        help="end_date YYYYMMDD")
594
595    parser.add_argument("--start_step", dest="start_step",
596                        help="start step in hours")
597    parser.add_argument("--end_step", dest="end_step",
598                        help="end step in hours")
599
600# some arguments that override the default in the CONTROL file
601    parser.add_argument("--levelist", dest="levelist",
602                        help="vertical levels to be retrieved, e.g. 30/to/60")
603    parser.add_argument("--area", dest="area",
604                        help="area defined as north/west/south/east")
605    parser.add_argument("--paramIds", dest="paramIds",
606                        help="parameter IDs")
607    parser.add_argument("--prefix", dest="prefix", default='EN',
608                        help="output file name prefix")
609
610# set the working directories
611    parser.add_argument("--inputdir", dest="inputdir", default=None,
612                        help="root directory for storing intermediate files")
613    parser.add_argument("--outputdir", dest="outputdir", default=None,
614                        help="root directory for storing output files")
615    parser.add_argument("--flexpart_root_scripts", dest="flexpart_root_scripts",
616                        help="FLEXPART root directory (to find \
617                        'grib2flexpart and COMMAND file)\n \
618                        Normally flex_extract resides in the scripts directory \
619                        of the FLEXPART distribution")
620
621    parser.add_argument("--controlfile", dest="controlfile",
622                        default='CONTROL.temp',
623                        help="file with CONTROL parameters")
624    args = parser.parse_args()
625
626    try:
627        c = ControlFile(args.controlfile)
628    except IOError:
629        try:
630            c = ControlFile(LOCAL_PYTHON_PATH + args.controlfile)
631        except IOError:
632            print 'Could not read CONTROL file "' + args.controlfile + '"'
633            print 'Either it does not exist or its syntax is wrong.'
634            print 'Try "' + sys.argv[0].split('/')[-1] + \
635                  ' -h" to print usage information'
636            exit(1)
637
638    if args.levelist:
639        c.levels = args.levelist.split('/')
640    else:
641        c.levels = [0]
642
643    if args.area:
644        c.area = args.area.split('/')
645    else:
646        c.area = '[0,0]'
647
648    c.paramIds = args.paramIds.split('/')
649
650    if args.start_step:
651        c.start_step = int(args.start_step)
652    else:
653        c.start_step = 0
654
655    if args.end_step:
656        c.end_step = int(args.end_step)
657    else:
658        c.end_step = 0
659
660    c.start_date = args.start_date
661    c.end_date = args.end_date
662
663    c.prefix = args.prefix
664
665    c.inputdir = args.inputdir
666
667    if args.outputdir:
668        c.outputdir = args.outputdir
669    else:
670        c.outputdir = c.inputdir
671
672    return args, c
673
674if __name__ == "__main__":
675    main()
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG