source: flex_extract.git/source/python/classes/GribTools.py @ 274f9ef

ctbtodev
Last change on this file since 274f9ef was 274f9ef, checked in by Anne Philipp <anne.philipp@…>, 5 years ago

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