source: flex_extract.git/source/pythontest/TestInstallTar/flex_extract_v7.1_ecgate/source/python/mods/plot_retrieved.py @ 25b14be

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

changed whole tree structure of flex_extract to have better overview

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