source: flex_extract.git/source/python/classes/GribTools.py @ 25b14be

dev
Last change on this file since 25b14be was 25b14be, checked in by Anne Philipp <anne.philipp@…>, 13 months ago

changed whole tree structure of flex_extract to have better overview

  • Property mode set to 100644
File size: 10.0 KB
Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
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__
32#    - get_keys
33#    - set_keys
34#    - copy
35#    - index
36#
37# @Class Attributes:
38#    - filenames
39#
40#*******************************************************************************
41
42# ------------------------------------------------------------------------------
43# MODULES
44# ------------------------------------------------------------------------------
45import os
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
50
51# ------------------------------------------------------------------------------
52# CLASS
53# ------------------------------------------------------------------------------
54class GribTools(object):
55    '''
56    Class for GRIB utilities (new methods) based on GRIB API
57    '''
58    # --------------------------------------------------------------------------
59    # CLASS FUNCTIONS
60    # --------------------------------------------------------------------------
61    def __init__(self, filenames):
62        '''
63        @Description:
64            Initialise an object of GribTools and assign a list
65            of filenames.
66
67        @Input:
68            filenames: list of strings
69                A list of filenames.
70
71        @Return:
72            <nothing>
73        '''
74
75        self.filenames = filenames
76
77        return
78
79
80    def get_keys(self, keynames, wherekeynames=[], wherekeyvalues=[]):
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
90            wherekeynames: list of strings, optional
91                Default value is an empty list.
92
93            wherekeyvalues: list of strings, optional
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
101        fileid = open(self.filenames, 'r')
102
103        return_list = []
104
105        while 1:
106            gid_in = grib_new_from_file(fileid)
107
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
117            for wherekey in wherekeynames:
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
125            if select:
126                llist = []
127                for key in keynames:
128                    llist.extend([str(grib_get(gid_in, key))])
129                return_list.append(llist)
130
131            grib_release(gid_in)
132
133        fileid.close()
134
135        return return_list
136
137
138    def set_keys(self, fromfile, keynames, keyvalues, wherekeynames=[],
139                 wherekeyvalues=[], strict=False, filemode='w'):
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
152            keynames: list of strings
153                List of keynames. Default is an empty list.
154
155            keyvalues: list of strings
156                List of keynames. Default is an empty list.
157
158            wherekeynames: list of strings, optional
159                Default value is an empty list.
160
161            wherekeyvalues: list of strings, optional
162                Default value is an empty list.
163
164            strict: boolean, optional
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
169            filemode: string, optional
170                Sets the mode for the output file. Default is "w".
171
172        @Return:
173            <nothing>
174
175        '''
176        fout = open(self.filenames, filemode)
177        fin = open(fromfile)
178
179        while 1:
180            gid_in = grib_new_from_file(fin)
181
182            if gid_in is None:
183                break
184
185            if len(wherekeynames) != len(wherekeyvalues):
186                raise Exception("Give a value for each keyname!")
187
188            select = True
189            i = 0
190            for wherekey in wherekeynames:
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
198            if select:
199                i = 0
200                for key in keynames:
201                    grib_set(gid_in, key, keyvalues[i])
202                    i += 1
203
204            grib_write(gid_in, fout)
205
206            grib_release(gid_in)
207
208        fin.close()
209        fout.close()
210
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
225            selectWhere: boolean, optional
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
230            keynames: list of strings, optional
231                List of keynames. Default is an empty list.
232
233            keyvalues: list of strings, optional
234                List of keynames. Default is an empty list.
235
236            filemode: string, optional
237                Sets the mode for the output file. Default is "w".
238
239        @Return:
240            <nothing>
241        '''
242
243        fin = open(filename_in)
244        fout = open(self.filenames, filemode)
245
246        while 1:
247            gid_in = grib_new_from_file(fin)
248
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
257            for key in keynames:
258                if not grib_is_defined(gid_in, key):
259                    raise Exception("Key was not defined")
260
261                if selectWhere:
262                    select = (select and (str(keyvalues[i]) ==
263                                          str(grib_get(gid_in, key))))
264                else:
265                    select = (select and (str(keyvalues[i]) !=
266                                          str(grib_get(gid_in, key))))
267                i += 1
268
269            if select:
270                grib_write(gid_in, fout)
271
272            grib_release(gid_in)
273
274        fin.close()
275        fout.close()
276
277        return
278
279    def index(self, index_keys=["mars"], index_file="my.idx"):
280        '''
281        @Description:
282            Create index file from a list of files if it does not exist or
283            read an index file.
284
285        @Input:
286            index_keys: list of strings, optional
287                Contains the list of key parameter names from
288                which the index is to be created.
289                Default is a list with a single entry string "mars".
290
291            index_file: string, optional
292                Filename where the indices are stored.
293                Default is "my.idx".
294
295        @Return:
296            iid: integer
297                Grib index id.
298        '''
299        print("... index will be done")
300        iid = None
301
302        if os.path.exists(index_file):
303            iid = grib_index_read(index_file)
304            print("Use existing index file: %s " % (index_file))
305        else:
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)
310                else:
311                    grib_index_add_file(iid, filename)
312
313            if iid is not None:
314                grib_index_write(iid, index_file)
315
316        print('... index done')
317
318        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG