source: flex_extract.git/source/python/classes/GribTools.py @ aa275fc

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

substituted grib_api with eccodes

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