source: flex_extract.git/source/python/classes/GribTools.py @ 0934db1

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

rechanged eccodes implementation, does not work with fortran grib_api lib
  • Property mode set to 100644
File size: 10.1 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# from eccodes import (codes_grib_new_from_file, codes_is_defined, codes_get,
53                     # codes_release, codes_set, codes_write, codes_index_read,
54                     # codes_index_new_from_file, codes_index_add_file,
55                     # codes_index_write)
56
57# ------------------------------------------------------------------------------
58# CLASS
59# ------------------------------------------------------------------------------
60class GribTools(object):
61    '''
62    Class for GRIB utilities (new methods) based on GRIB API
63    '''
64    # --------------------------------------------------------------------------
65    # CLASS FUNCTIONS
66    # --------------------------------------------------------------------------
67    def __init__(self, filenames):
68        '''
69        @Description:
70            Initialise an object of GribTools and assign a list
71            of filenames.
72
73        @Input:
74            filenames: list of strings
75                A list of filenames.
76
77        @Return:
78            <nothing>
79        '''
80
81        self.filenames = filenames
82
83        return
84
85
86    def get_keys(self, keynames, wherekeynames=[], wherekeyvalues=[]):
87        '''
88        @Description:
89            get keyvalues for a given list of keynames
90            a where statement can be given (list of key and list of values)
91
92        @Input:
93            keynames: list of strings
94                List of keynames.
95
96            wherekeynames: list of strings, optional
97                Default value is an empty list.
98
99            wherekeyvalues: list of strings, optional
100                Default value is an empty list.
101
102        @Return:
103            return_list: list of strings
104                List of keyvalues for given keynames.
105        '''
106
107        fileid = open(self.filenames, 'r')
108
109        return_list = []
110
111        while 1:
112            gid = grib_new_from_file(fileid)
113
114            if gid is None:
115                break
116
117            if len(wherekeynames) != len(wherekeyvalues):
118                raise Exception("Number of key values and key names must be \
119                                 the same. Give a value for each keyname!")
120
121            select = True
122            i = 0
123            for wherekey in wherekeynames:
124                if not grib_is_defined(gid, wherekey):
125                    raise Exception("where key was not defined")
126
127                select = (select and (str(wherekeyvalues[i]) ==
128                                      str(grib_get(gid, wherekey))))
129                i += 1
130
131            if select:
132                llist = []
133                for key in keynames:
134                    llist.extend([str(grib_get(gid, key))])
135                return_list.append(llist)
136
137            grib_release(gid)
138
139        fileid.close()
140
141        return return_list
142
143
144    def set_keys(self, fromfile, keynames, keyvalues, wherekeynames=[],
145                 wherekeyvalues=[], strict=False, filemode='w'):
146        '''
147        @Description:
148            Opens the file to read the grib messages and then write
149            them to a new output file. By default all messages are
150            written out. Also, the keyvalues of the passed list of
151            keynames are set or only those meeting the where statement.
152            (list of key and list of values).
153
154        @Input:
155            fromfile: string
156                Filename of the input file to read the grib messages from.
157
158            keynames: list of strings
159                List of keynames. Default is an empty list.
160
161            keyvalues: list of strings
162                List of keynames. Default is an empty list.
163
164            wherekeynames: list of strings, optional
165                Default value is an empty list.
166
167            wherekeyvalues: list of strings, optional
168                Default value is an empty list.
169
170            strict: boolean, optional
171                Decides if everything from keynames and keyvalues
172                is written out the grib file (False) or only those
173                meeting the where statement (True). Default is False.
174
175            filemode: string, optional
176                Sets the mode for the output file. Default is "w".
177
178        @Return:
179            <nothing>
180
181        '''
182        fout = open(self.filenames, filemode)
183        fin = open(fromfile)
184
185        while 1:
186            gid = grib_new_from_file(fin)
187
188            if gid is None:
189                break
190
191            if len(wherekeynames) != len(wherekeyvalues):
192                raise Exception("Give a value for each keyname!")
193
194            select = True
195            i = 0
196            for wherekey in wherekeynames:
197                if not grib_is_defined(gid, wherekey):
198                    raise Exception("where Key was not defined")
199
200                select = (select and (str(wherekeyvalues[i]) ==
201                                      str(grib_get(gid, wherekey))))
202                i += 1
203
204            if select:
205                i = 0
206                for key in keynames:
207                    grib_set(gid, key, keyvalues[i])
208                    i += 1
209
210            grib_write(gid, fout)
211
212            grib_release(gid)
213
214        fin.close()
215        fout.close()
216
217        return
218
219    def copy(self, filename_in, selectWhere=True,
220             keynames=[], keyvalues=[], filemode='w'):
221        '''
222        Add the content of another input grib file to the objects file but
223        only messages corresponding to keys/values passed to the function.
224        The selectWhere switch decides if to copy the keys equal to (True) or
225        different to (False) the keynames/keyvalues list passed to the function.
226
227        @Input:
228            filename_in: string
229                Filename of the input file to read the grib messages from.
230
231            selectWhere: boolean, optional
232                Decides if to copy the keynames and values equal to (True) or
233                different to (False) the keynames/keyvalues list passed to the
234                function. Default is True.
235
236            keynames: list of strings, optional
237                List of keynames. Default is an empty list.
238
239            keyvalues: list of strings, optional
240                List of keynames. Default is an empty list.
241
242            filemode: string, optional
243                Sets the mode for the output file. Default is "w".
244
245        @Return:
246            <nothing>
247        '''
248
249        fin = open(filename_in)
250        fout = open(self.filenames, filemode)
251
252        while 1:
253            gid = grib_new_from_file(fin)
254
255            if gid is None:
256                break
257
258            if len(keynames) != len(keyvalues):
259                raise Exception("Give a value for each keyname!")
260
261            select = True
262            i = 0
263            for key in keynames:
264                if not grib_is_defined(gid, key):
265                    raise Exception("Key was not defined")
266
267                if selectWhere:
268                    select = (select and (str(keyvalues[i]) ==
269                                          str(grib_get(gid, key))))
270                else:
271                    select = (select and (str(keyvalues[i]) !=
272                                          str(grib_get(gid, key))))
273                i += 1
274
275            if select:
276                grib_write(gid, fout)
277
278            grib_release(gid)
279
280        fin.close()
281        fout.close()
282
283        return
284
285    def index(self, index_keys=["mars"], index_file="my.idx"):
286        '''
287        @Description:
288            Create index file from a list of files if it does not exist or
289            read an index file.
290
291        @Input:
292            index_keys: list of strings, optional
293                Contains the list of key parameter names from
294                which the index is to be created.
295                Default is a list with a single entry string "mars".
296
297            index_file: string, optional
298                Filename where the indices are stored.
299                Default is "my.idx".
300
301        @Return:
302            iid: integer
303                Grib index id.
304        '''
305        print("... index will be done")
306        iid = None
307
308        if os.path.exists(index_file):
309            iid = grib_index_read(index_file)
310            print("Use existing index file: %s " % (index_file))
311        else:
312            for filename in self.filenames:
313                print("Inputfile: %s " % (filename))
314                if iid is None:
315                    iid = grib_index_new_from_file(filename, index_keys)
316                else:
317                    grib_index_add_file(iid, filename)
318
319            if iid is not None:
320                grib_index_write(iid, index_file)
321
322        print('... index done')
323
324        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG