source: flex_extract.git/python/GribTools.py @ e1228f3

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

resolved loop import between controlfile and tools

  • Property mode set to 100644
File size: 10.0 KB
RevLine 
[64cf353]1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
[991df6a]3#*******************************************************************************
4# @Author: Anne Fouilloux (University of Oslo)
5#
6# @Date: July 2014
7#
8# @Change History:
9#   February 2018 - Anne Philipp (University of Vienna):
10#        - applied PEP8 style guide
11#        - added documentation
12#        - changed some naming
13#
14# @License:
15#    (C) Copyright 2014-2018.
16#
17#    This software is licensed under the terms of the Apache Licence Version 2.0
18#    which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
19#
20# @Class Description:
21#    The GRIB API provides all necessary tools to work directly with the
22#    grib files. Nevertheless, the GRIB API tools are very basic and are in
23#    direct connection with the grib files. This class provides some higher
24#    functions which apply a set of GRIB API tools together in the respective
25#    context. So, the class initially contains a list of grib files (their
26#    names) and the using program then applies the methods directly on the
27#    class objects without having to think about how the actual GRIB API
28#    tools have to be arranged.
29#
30# @Class Content:
31#    - __init__
[ff99eae]32#    - get_keys
33#    - set_keys
[991df6a]34#    - copy
35#    - index
36#
[ff99eae]37# @Class Attributes:
38#    - filenames
39#
[991df6a]40#*******************************************************************************
[64cf353]41
42# ------------------------------------------------------------------------------
43# MODULES
44# ------------------------------------------------------------------------------
[efdb01a]45import os
[ff99eae]46from gribapi import grib_new_from_file, grib_is_defined, grib_get, \
47                    grib_release, grib_set, grib_write, grib_index_read, \
48                    grib_index_new_from_file, grib_index_add_file,  \
49                    grib_index_write
[991df6a]50
[64cf353]51# ------------------------------------------------------------------------------
52# CLASS
53# ------------------------------------------------------------------------------
[e1228f3]54class GribTools(object):
[64cf353]55    '''
56    Class for GRIB utilities (new methods) based on GRIB API
57    '''
58    # --------------------------------------------------------------------------
[efdb01a]59    # CLASS FUNCTIONS
[64cf353]60    # --------------------------------------------------------------------------
61    def __init__(self, filenames):
62        '''
63        @Description:
[e1228f3]64            Initialise an object of GribTools and assign a list
[64cf353]65            of filenames.
66
67        @Input:
68            filenames: list of strings
69                A list of filenames.
70
71        @Return:
72            <nothing>
73        '''
74
[ff99eae]75        self.filenames = filenames
[64cf353]76
77        return
78
79
[ff99eae]80    def get_keys(self, keynames, wherekeynames=[], wherekeyvalues=[]):
[64cf353]81        '''
82        @Description:
83            get keyvalues for a given list of keynames
84            a where statement can be given (list of key and list of values)
85
86        @Input:
87            keynames: list of strings
88                List of keynames.
89
[ff99eae]90            wherekeynames: list of strings, optional
[64cf353]91                Default value is an empty list.
92
[ff99eae]93            wherekeyvalues: list of strings, optional
[64cf353]94                Default value is an empty list.
95
96        @Return:
97            return_list: list of strings
98                List of keyvalues for given keynames.
99        '''
100
[ff99eae]101        fileid = open(self.filenames, 'r')
[64cf353]102
103        return_list = []
104
[d69b677]105        while 1:
[64cf353]106            gid_in = grib_new_from_file(fileid)
[d69b677]107
[64cf353]108            if gid_in is None:
109                break
110
111            if len(wherekeynames) != len(wherekeyvalues):
112                raise Exception("Number of key values and key names must be \
113                                 the same. Give a value for each keyname!")
114
115            select = True
116            i = 0
[d69b677]117            for wherekey in wherekeynames:
[64cf353]118                if not grib_is_defined(gid_in, wherekey):
119                    raise Exception("where key was not defined")
120
121                select = (select and (str(wherekeyvalues[i]) ==
122                                      str(grib_get(gid_in, wherekey))))
123                i += 1
124
[d69b677]125            if select:
126                llist = []
127                for key in keynames:
128                    llist.extend([str(grib_get(gid_in, key))])
129                return_list.append(llist)
[64cf353]130
[d69b677]131            grib_release(gid_in)
[64cf353]132
[d69b677]133        fileid.close()
[64cf353]134
[d69b677]135        return return_list
[64cf353]136
137
[ff99eae]138    def set_keys(self, fromfile, keynames, keyvalues, wherekeynames=[],
139                 wherekeyvalues=[], strict=False, filemode='w'):
[64cf353]140        '''
141        @Description:
142            Opens the file to read the grib messages and then write
143            them to a new output file. By default all messages are
144            written out. Also, the keyvalues of the passed list of
145            keynames are set or only those meeting the where statement.
146            (list of key and list of values).
147
148        @Input:
149            fromfile: string
150                Filename of the input file to read the grib messages from.
151
[ff99eae]152            keynames: list of strings
[64cf353]153                List of keynames. Default is an empty list.
154
[ff99eae]155            keyvalues: list of strings
[64cf353]156                List of keynames. Default is an empty list.
157
[ff99eae]158            wherekeynames: list of strings, optional
[64cf353]159                Default value is an empty list.
160
[ff99eae]161            wherekeyvalues: list of strings, optional
[64cf353]162                Default value is an empty list.
163
[02c8c50]164            strict: boolean, optional
[64cf353]165                Decides if everything from keynames and keyvalues
166                is written out the grib file (False) or only those
167                meeting the where statement (True). Default is False.
168
[02c8c50]169            filemode: string, optional
[64cf353]170                Sets the mode for the output file. Default is "w".
171
172        @Return:
173            <nothing>
174
175        '''
[ff99eae]176        fout = open(self.filenames, filemode)
[64cf353]177        fin = open(fromfile)
178
[d69b677]179        while 1:
[64cf353]180            gid_in = grib_new_from_file(fin)
181
182            if gid_in is None:
183                break
[d69b677]184
[64cf353]185            if len(wherekeynames) != len(wherekeyvalues):
186                raise Exception("Give a value for each keyname!")
187
188            select = True
189            i = 0
[d69b677]190            for wherekey in wherekeynames:
[64cf353]191                if not grib_is_defined(gid_in, wherekey):
192                    raise Exception("where Key was not defined")
193
194                select = (select and (str(wherekeyvalues[i]) ==
195                                      str(grib_get(gid_in, wherekey))))
196                i += 1
197
[d69b677]198            if select:
[64cf353]199                i = 0
[d69b677]200                for key in keynames:
201                    grib_set(gid_in, key, keyvalues[i])
[64cf353]202                    i += 1
203
[ff99eae]204            grib_write(gid_in, fout)
205
[d69b677]206            grib_release(gid_in)
[64cf353]207
[d69b677]208        fin.close()
209        fout.close()
210
[64cf353]211        return
212
213    def copy(self, filename_in, selectWhere=True,
214             keynames=[], keyvalues=[], filemode='w'):
215        '''
216        Add the content of another input grib file to the objects file but
217        only messages corresponding to keys/values passed to the function.
218        The selectWhere switch decides if to copy the keys equal to (True) or
219        different to (False) the keynames/keyvalues list passed to the function.
220
221        @Input:
222            filename_in: string
223                Filename of the input file to read the grib messages from.
224
[02c8c50]225            selectWhere: boolean, optional
[64cf353]226                Decides if to copy the keynames and values equal to (True) or
227                different to (False) the keynames/keyvalues list passed to the
228                function. Default is True.
229
[ff99eae]230            keynames: list of strings, optional
[64cf353]231                List of keynames. Default is an empty list.
232
[ff99eae]233            keyvalues: list of strings, optional
[64cf353]234                List of keynames. Default is an empty list.
235
[02c8c50]236            filemode: string, optional
[64cf353]237                Sets the mode for the output file. Default is "w".
238
239        @Return:
240            <nothing>
241        '''
242
243        fin = open(filename_in)
[ff99eae]244        fout = open(self.filenames, filemode)
[64cf353]245
[d69b677]246        while 1:
[64cf353]247            gid_in = grib_new_from_file(fin)
[d69b677]248
[64cf353]249            if gid_in is None:
250                break
251
252            if len(keynames) != len(keyvalues):
253                raise Exception("Give a value for each keyname!")
254
255            select = True
256            i = 0
[d69b677]257            for key in keynames:
[64cf353]258                if not grib_is_defined(gid_in, key):
259                    raise Exception("Key was not defined")
260
[d69b677]261                if selectWhere:
[64cf353]262                    select = (select and (str(keyvalues[i]) ==
263                                          str(grib_get(gid_in, key))))
[d69b677]264                else:
[64cf353]265                    select = (select and (str(keyvalues[i]) !=
266                                          str(grib_get(gid_in, key))))
267                i += 1
268
[d69b677]269            if select:
[64cf353]270                grib_write(gid_in, fout)
271
[d69b677]272            grib_release(gid_in)
[64cf353]273
[d69b677]274        fin.close()
275        fout.close()
276
[64cf353]277        return
278
279    def index(self, index_keys=["mars"], index_file="my.idx"):
280        '''
281        @Description:
[02c8c50]282            Create index file from a list of files if it does not exist or
[64cf353]283            read an index file.
284
285        @Input:
[02c8c50]286            index_keys: list of strings, optional
287                Contains the list of key parameter names from
288                which the index is to be created.
[64cf353]289                Default is a list with a single entry string "mars".
290
[02c8c50]291            index_file: string, optional
[64cf353]292                Filename where the indices are stored.
293                Default is "my.idx".
294
295        @Return:
296            iid: integer
297                Grib index id.
298        '''
[ff99eae]299        print "... index will be done"
300        iid = None
[64cf353]301
[ff99eae]302        if os.path.exists(index_file):
303            iid = grib_index_read(index_file)
304            print "Use existing index file: %s " % (index_file)
[d69b677]305        else:
[ff99eae]306            for filename in self.filenames:
307                print "Inputfile: %s " % (filename)
308                if iid is None:
309                    iid = grib_index_new_from_file(filename, index_keys)
[d69b677]310                else:
[ccab809]311                    print 'in else zweig'
[ff99eae]312                    grib_index_add_file(iid, filename)
[d69b677]313
[ff99eae]314            if iid is not None:
315                grib_index_write(iid, index_file)
[d69b677]316
[ff99eae]317        print '... index done'
[d69b677]318
[ff99eae]319        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG