source: flex_extract.git/Source/Python/Classes/GribUtil.py @ 026b359

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

updated license status (year)
  • Property mode set to 100644
File size: 10.0 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, keynames, keyvalues, wherekeynames,
125                 wherekeyvalues, filemode='wb'):
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
165        with open(fromfile, 'rb') as fin:
166            gid = codes_grib_new_from_file(fin)
167
168            select = True
169            i = 0
170            for wherekey in wherekeynames:
171                if not codes_is_defined(gid, wherekey):
172                    raise Exception("wherekey was not defined")
173
174                select = (select and (str(wherekeyvalues[i]) ==
175                                      str(codes_get(gid, wherekey))))
176                i += 1
177
178            if select:
179                i = 0
180                for key in keynames:
181                    if key == 'values':
182                        codes_set_values(gid, keyvalues[i])
183                    else:
184                        codes_set(gid, key, keyvalues[i])
185                    i += 1
186
187                codes_write(gid, fout)
188
189            codes_release(gid)
190
191        fout.close()
192
193        return
194
195    def copy_dummy_msg(self, filename_in, keynames, keyvalues,
196                       selectwhere=True, filemode='wb'):
197        '''Add the content of another input grib file to the objects file but
198        only messages corresponding to keys/values passed to the function.
199        The selectWhere switch decides if to copy the keys equal to (True) or
200        different to (False) the keynames/keyvalues list passed to the function.
201
202        Parameters
203        ----------
204        filename_in : :obj:`string`
205            Filename of the input file to read the grib messages from.
206
207        selectwhere : :obj:`boolean`, optional
208            Decides if to copy the keynames and values equal to (True) or
209            different to (False) the keynames/keyvalues list passed to the
210            function. Default is True.
211
212        keynames : :obj:`list` of :obj:`string`
213            List of keynames.
214
215        keyvalues : :obj:`list` of :obj:`string`
216            List of keyvalues.
217
218        filemode : :obj:`string`, optional
219            Sets the mode for the output file. Default is "wb".
220
221        Return
222        ------
223
224        '''
225        from eccodes import (codes_grib_new_from_file, codes_is_defined,
226                             codes_get, codes_release, codes_write)
227
228        if len(keynames) != len(keyvalues):
229            raise Exception("Give a value for each keyname!")
230
231
232        fout = open(self.filenames, filemode)
233
234        fields = 0
235
236        with open(filename_in, 'rb') as fin:
237            if fields >= 1:
238                fout.close()
239                return
240
241            gid = codes_grib_new_from_file(fin)
242
243            select = True
244            i = 0
245            for key in keynames:
246                if not codes_is_defined(gid, key):
247                    raise Exception("Key was not defined")
248
249                if selectwhere:
250                    select = (select and (str(keyvalues[i]) ==
251                                          str(codes_get(gid, key))))
252                else:
253                    select = (select and (str(keyvalues[i]) !=
254                                          str(codes_get(gid, key))))
255                i += 1
256
257            if select:
258                fields = fields + 1
259                codes_write(gid, fout)
260
261            codes_release(gid)
262
263        fout.close()
264
265        return
266
267    def index(self, index_keys, index_file="my.idx"):
268        '''Create index file from a list of files if it does not exist or
269        read an index file.
270
271        Parameters
272        ----------
273        index_keys: :obj:`list` of :obj:`string`
274            Contains the list of key parameter names from
275            which the index is to be created.
276
277        index_file: :obj:`string`, optional
278            Filename where the indices are stored.
279            Default is "my.idx".
280
281        Return
282        ------
283        iid: :obj:`integer`
284            Grib index id.
285        '''
286        from eccodes import (codes_index_read, codes_index_new_from_file,
287                             codes_index_add_file, codes_index_write)
288
289        print("... index will be done")
290        iid = None
291
292        if os.path.exists(index_file):
293            iid = codes_index_read(index_file)
294            print("Use existing index file: %s " % (index_file))
295        else:
296            for filename in self.filenames:
297                print("Inputfile: %s " % (filename))
298                if iid is None:
299                    iid = codes_index_new_from_file(filename, index_keys)
300                else:
301                    codes_index_add_file(iid, filename)
302
303            if iid is not None:
304                codes_index_write(iid, index_file)
305
306        print('... index done')
307
308        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG