source: flex_extract.git/Source/Python/Classes/GribUtil.py @ 2c3c135

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

eliminated bug with eccodes file opening for python3 for rrint feature
  • Property mode set to 100644
File size: 10.4 KB
Line 
1#!/usr/bin/env python3
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-2020.
16#    Anne Philipp, Leopold Haimberger
17#
18#    SPDX-License-Identifier: CC-BY-4.0
19#
20#    This work is licensed under the Creative Commons Attribution 4.0
21#    International License. To view a copy of this license, visit
22#    http://creativecommons.org/licenses/by/4.0/ or send a letter to
23#    Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
24#*******************************************************************************
25
26# ------------------------------------------------------------------------------
27# MODULES
28# ------------------------------------------------------------------------------
29from __future__ import print_function
30
31import os
32
33# ------------------------------------------------------------------------------
34# CLASS
35# ------------------------------------------------------------------------------
36class GribUtil(object):
37    '''
38    Class for GRIB utilities (new methods) based on GRIB API
39
40    The GRIB API provides all necessary tools to work directly with the
41    grib files. Nevertheless, the GRIB API tools are very basic and are in
42    direct connection with the grib files. This class provides some higher
43    functions which apply a set of GRIB API tools together in the respective
44    context. So, the class initially contains a list of grib files (their
45    names) and the using program then applies the methods directly on the
46    class objects without having to think about how the actual GRIB API
47    tools have to be arranged.
48    '''
49    # --------------------------------------------------------------------------
50    # CLASS FUNCTIONS
51    # --------------------------------------------------------------------------
52    def __init__(self, filenames):
53        '''Initialise an object of GribUtil and assign a list of filenames.
54
55        Parameters
56        ----------
57        filenames : :obj:`list` of :obj:`strings`
58             A list of filenames.
59
60        Return
61        ------
62
63        '''
64
65        self.filenames = filenames
66
67        return
68
69
70    def get_keys(self, keynames, wherekeynames, wherekeyvalues):
71        '''Get keyvalues for a given list of keynames a where statement
72        can be given (list of key and list of values)
73
74        Parameters
75        ----------
76        keynames : :obj:`list` of :obj:`string`
77            List of keynames.
78
79        wherekeynames : :obj:`list` of :obj:`string`
80            List of key names for indexing grib message parameter.
81
82        wherekeyvalues : :obj:`list` of :obj:`string`
83            List of key values corresponding the key names.
84
85        Return
86        ------
87        return_list : :obj:`list` of :obj:`string`
88            List of keyvalues for given keynames.
89        '''
90        from eccodes import (codes_new_from_file, codes_is_defined, codes_get,
91                             codes_release)
92
93        return_list = []
94
95        with open(self.filenames, 'rb') as fileid:
96
97            gid = codes_new_from_file(fileid)
98
99            if len(wherekeynames) != len(wherekeyvalues):
100                raise Exception("Number of key values and key names must be \
101                                 the same. Give a value for each keyname!")
102
103            select = True
104            i = 0
105            for wherekey in wherekeynames:
106                if not codes_is_defined(gid, wherekey):
107                    raise Exception("where key was not defined")
108
109                select = (select and (str(wherekeyvalues[i]) ==
110                                      str(codes_get(gid, wherekey))))
111                i += 1
112
113            if select:
114                llist = []
115                for key in keynames:
116                    llist.extend([str(codes_get(gid, key))])
117                return_list.append(llist)
118
119            codes_release(gid)
120
121        return return_list
122
123
124    def set_keys(self, fromfile, filemode='wb', keynames=[], keyvalues=[], 
125                 wherekeynames=[], wherekeyvalues=[]):
126        '''Opens the file to read the grib messages and then write
127        the selected messages (with wherekeys) to a new output file.
128        Also, the keyvalues of the passed list of keynames are set.
129
130        Parameters
131        ----------
132        fromfile : :obj:`string`
133            Filename of the input file to read the grib messages from.
134
135        keynames : :obj:`list` of :obj:`string`
136            List of keynames to set in the selected messages.
137            Default is an empty list.
138
139        keyvalues : :obj:`list` of :obj:`string`
140            List of keyvalues to set in the selected messages.
141            Default is an empty list.
142
143        wherekeynames : :obj:`list` of :obj:`string`
144            List of keynames to select correct message.
145
146        wherekeyvalues : :obj:`list` of :obj:`string`
147            List of keyvalues for keynames to select correct message.
148
149        filemode : :obj:`string`, optional
150            Sets the mode for the output file. Default is "wb".
151
152        Return
153        ------
154
155        '''
156        from eccodes import (codes_grib_new_from_file, codes_is_defined,
157                             codes_get, codes_set, codes_write,
158                             codes_set_values, codes_release)
159
160        if len(wherekeynames) != len(wherekeyvalues):
161            raise Exception("Give a value for each keyname!")
162
163        fout = open(self.filenames, filemode)
164       # print(fout)
165        fin = open(fromfile, 'rb')
166       # print(fin)
167        while True:
168            gid = codes_grib_new_from_file(fin)
169           # print('test')
170            if gid is None:
171                break           
172           
173            select = True
174            #print(str(codes_get(gid,'paramId')))
175            for i, wherekey in enumerate(wherekeynames):
176                if not codes_is_defined(gid, wherekey):
177                    raise Exception("wherekey was not defined")
178
179                select = (select and (str(wherekeyvalues[i]) ==
180                                      str(codes_get(gid, wherekey))))
181
182            if select:
183                for i, key in enumerate(keynames):
184                    if key == 'values':
185                        codes_set_values(gid, keyvalues[i])
186                    else:
187                        codes_set(gid, key, keyvalues[i])
188
189                codes_write(gid, fout)
190
191            codes_release(gid)
192
193        fout.close()
194        fin.close()
195
196        return
197
198    def copy_dummy_msg(self, filename_in, keynames, keyvalues,
199                       selectwhere=True, filemode='wb'):
200        '''Add the content of another input grib file to the objects file but
201        only messages corresponding to keys/values passed to the function.
202        The selectWhere switch decides if to copy the keys equal to (True) or
203        different to (False) the keynames/keyvalues list passed to the function.
204
205        Parameters
206        ----------
207        filename_in : :obj:`string`
208            Filename of the input file to read the grib messages from.
209
210        selectwhere : :obj:`boolean`, optional
211            Decides if to copy the keynames and values equal to (True) or
212            different to (False) the keynames/keyvalues list passed to the
213            function. Default is True.
214
215        keynames : :obj:`list` of :obj:`string`
216            List of keynames.
217
218        keyvalues : :obj:`list` of :obj:`string`
219            List of keyvalues.
220
221        filemode : :obj:`string`, optional
222            Sets the mode for the output file. Default is "wb".
223
224        Return
225        ------
226
227        '''
228        from eccodes import (codes_grib_new_from_file, codes_is_defined,
229                             codes_get, codes_release, codes_write)
230
231        if len(keynames) != len(keyvalues):
232            raise Exception("Give a value for each keyname!")
233
234        fout = open(self.filenames, filemode)
235        fin = open(filename_in, 'rb')
236       
237        fields = 0
238        while True:
239            if fields >= len(keyvalues):
240                fout.close()
241                fin.close()
242                return
243               
244            gid = codes_grib_new_from_file(fin)
245            if gid is None:
246                break           
247
248            for i, key in enumerate(keynames):
249               
250                select = True
251               
252                if not codes_is_defined(gid, key):
253                    raise Exception("Key was not defined")
254
255                if selectwhere:
256                    select = (select and (str(keyvalues[i]) ==
257                                          str(codes_get(gid, key))))
258                else:
259                    select = (select and (str(keyvalues[i]) !=
260                                          str(codes_get(gid, key))))
261               
262                if select:
263                    fields = fields + 1
264                    codes_write(gid, fout)
265                   
266            codes_release(gid)
267
268        fout.close()
269        fin.close()
270
271        return
272
273    def index(self, index_keys, index_file="my.idx"):
274        '''Create index file from a list of files if it does not exist or
275        read an index file.
276
277        Parameters
278        ----------
279        index_keys: :obj:`list` of :obj:`string`
280            Contains the list of key parameter names from
281            which the index is to be created.
282
283        index_file: :obj:`string`, optional
284            Filename where the indices are stored.
285            Default is "my.idx".
286
287        Return
288        ------
289        iid: :obj:`integer`
290            Grib index id.
291        '''
292        from eccodes import (codes_index_read, codes_index_new_from_file,
293                             codes_index_add_file, codes_index_write)
294
295        print("... index will be done")
296        iid = None
297
298        if os.path.exists(index_file):
299            iid = codes_index_read(index_file)
300            print("Use existing index file: %s " % (index_file))
301        else:
302            for filename in self.filenames:
303                print("Inputfile: %s " % (filename))
304                if iid is None:
305                    iid = codes_index_new_from_file(filename, index_keys)
306                else:
307                    codes_index_add_file(iid, filename)
308
309            if iid is not None:
310                codes_index_write(iid, index_file)
311
312        print('... index done')
313
314        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG