52class open():
 53    """
 54    GRIB2 File Object.  A physical file can contain one or more GRIB2 messages.  When instantiated,
 55    class `grib2io.open`, the file named `filename` is opened for reading (`mode = 'r'`) and is
 56    automatically indexed.  The indexing procedure reads some of the GRIB2 metadata for all GRIB2 Messages.
 57
 58    A GRIB2 Message may contain submessages whereby Section 2-7 can be repeated.  grib2io accommodates
 59    for this by flattening any GRIB2 submessages into multiple individual messages.
 60
 61    Attributes
 62    ----------
 63
 64    **`mode`** File IO mode of opening the file.
 65
 66    **`name`** Full path name of the GRIB2 file.
 67
 68    **`messages`** Count of GRIB2 Messages contained in the file.
 69
 70    **`current_message`** Current position of the file in units of GRIB2 Messages.
 71
 72    **`size`** Size of the file in units of bytes.
 73
 74    **`closed`** `True` is file handle is close; `False` otherwise.
 75
 76    **`variables`** Tuple containing a unique list of variable short names (i.e. GRIB2 abbreviation names).
 77
 78    **`levels`** Tuple containing a unique list of wgrib2-formatted level/layer strings.
 79    """
 80    __slots__ = ('_filehandle','_hasindex','_index','mode','name','messages',
 81                 'current_message','size','closed','variables','levels','_pos')
 82    def __init__(self, filename, mode='r', **kwargs):
 83        """
 84        `open` Constructor
 85
 86        Parameters
 87        ----------
 88
 89        **`filename : str`**
 90
 91        File name containing GRIB2 messages.
 92
 93        **`mode : str, optional`**
 94
 95        File access mode where `r` opens the files for reading only; `w` opens the file for writing.
 96        """
 97        if mode in {'a','r','w'}:
 98            mode = mode+'b'
 99            if 'w' in mode: mode += '+'
100        self._filehandle = builtins.open(filename,mode=mode,buffering=ONE_MB)
101        self._hasindex = False
102        self._index = {}
103        self.mode = mode
104        self.name = os.path.abspath(filename)
105        self.messages = 0
106        self.current_message = 0
107        self.size = os.path.getsize(self.name)
108        self.closed = self._filehandle.closed
109        self.levels = None
110        self.variables = None
111        if 'r' in self.mode:
112            try:
113                self._build_index(no_data=kwargs['_xarray_backend'])
114            except(KeyError):
115                self._build_index()
116        # FIX: Cannot perform reads on mode='a'
117        #if 'a' in self.mode and self.size > 0: self._build_index()
118
119
120    def __delete__(self, instance):
121        """
122        """
123        self.close()
124        del self._index
125
126
127    def __enter__(self):
128        """
129        """
130        return self
131
132
133    def __exit__(self, atype, value, traceback):
134        """
135        """
136        self.close()
137
138
139    def __iter__(self):
140        """
141        """
142        yield from self._index['msg']
143
144
145    def __len__(self):
146        """
147        """
148        return self.messages
149
150
151    def __repr__(self):
152        """
153        """
154        strings = []
155        for k in self.__slots__:
156            if k.startswith('_'): continue
157            strings.append('%s = %s\n'%(k,eval('self.'+k)))
158        return ''.join(strings)
159
160
161    def __getitem__(self, key):
162        """
163        """
164        if isinstance(key,int):
165            if abs(key) >= len(self._index['msg']):
166                raise IndexError("index out of range")
167            else:
168                return self._index['msg'][key]
169        elif isinstance(key,str):
170            return self.select(shortName=key)
171        elif isinstance(key,slice):
172            return self._index['msg'][key]
173        else:
174            raise KeyError('Key must be an integer, slice, or GRIB2 variable shortName.')
175
176
177    def _build_index(self, no_data=False):
178        """
179        Perform indexing of GRIB2 Messages.
180        """
181        # Initialize index dictionary
182        if not self._hasindex:
183            self._index['offset'] = []
184            self._index['bitmap_offset'] = []
185            self._index['data_offset'] = []
186            self._index['size'] = []
187            self._index['data_size'] = []
188            self._index['submessageOffset'] = []
189            self._index['submessageBeginSection'] = []
190            self._index['isSubmessage'] = []
191            self._index['messageNumber'] = []
192            self._index['msg'] = []
193            self._hasindex = True
194
195        # Iterate
196        while True:
197            try:
198                # Read first 4 bytes and decode...looking for "GRIB"
199                pos = self._filehandle.tell()
200                header = struct.unpack('>i',self._filehandle.read(4))[0]
201
202                # Test header. Then get information from GRIB2 Section 0: the discipline
203                # number, edition number (should always be 2), and GRIB2 message size.
204                # Then iterate to check for submessages.
205                if header.to_bytes(4,'big') == b'GRIB':
206
207                    _issubmessage = False
208                    _submsgoffset = 0
209                    _submsgbegin = 0
210                    _bmapflag = None
211
212                    # Read the rest of Section 0 using struct.
213                    section0 = np.concatenate(([header],list(struct.unpack('>HBBQ',self._filehandle.read(12)))),dtype=np.int64)
214                    assert section0[3] == 2
215
216                    # Read and unpack Section 1
217                    secsize = struct.unpack('>i',self._filehandle.read(4))[0]
218                    secnum = struct.unpack('>B',self._filehandle.read(1))[0]
219                    assert secnum == 1
220                    self._filehandle.seek(self._filehandle.tell()-5)
221                    _grbmsg = self._filehandle.read(secsize)
222                    _grbpos = 0
223                    section1,_grbpos = g2clib.unpack1(_grbmsg,_grbpos,np.empty)
224                    secrange = range(2,8)
225                    while 1:
226                        section2 = b''
227                        for num in secrange:
228                            secsize = struct.unpack('>i',self._filehandle.read(4))[0]
229                            secnum = struct.unpack('>B',self._filehandle.read(1))[0]
230                            if secnum == num:
231                                if secnum == 2:
232                                    if secsize > 0:
233                                        section2 = self._filehandle.read(secsize-5)
234                                elif secnum == 3:
235                                    self._filehandle.seek(self._filehandle.tell()-5)
236                                    _grbmsg = self._filehandle.read(secsize)
237                                    _grbpos = 0
238                                    # Unpack Section 3
239                                    _gds,_gdt,_deflist,_grbpos = g2clib.unpack3(_grbmsg,_grbpos,np.empty)
240                                    _gds = _gds.tolist()
241                                    _gdt = _gdt.tolist()
242                                    section3 = np.concatenate((_gds,_gdt))
243                                    section3 = np.where(section3==4294967295,-1,section3)
244                                elif secnum == 4:
245                                    self._filehandle.seek(self._filehandle.tell()-5)
246                                    _grbmsg = self._filehandle.read(secsize)
247                                    _grbpos = 0
248                                    # Unpack Section 4
249                                    _numcoord,_pdt,_pdtnum,_coordlist,_grbpos = g2clib.unpack4(_grbmsg,_grbpos,np.empty)
250                                    _pdt = _pdt.tolist()
251                                    section4 = np.concatenate((np.array((_numcoord,_pdtnum)),_pdt))
252                                elif secnum == 5:
253                                    self._filehandle.seek(self._filehandle.tell()-5)
254                                    _grbmsg = self._filehandle.read(secsize)
255                                    _grbpos = 0
256                                    # Unpack Section 5
257                                    _drt,_drtn,_npts,self._pos = g2clib.unpack5(_grbmsg,_grbpos,np.empty)
258                                    section5 = np.concatenate((np.array((_npts,_drtn)),_drt))
259                                elif secnum == 6:
260                                    # Unpack Section 6. Not really...just get the flag value.
261                                    _bmapflag = struct.unpack('>B',self._filehandle.read(1))[0]
262                                    if _bmapflag == 0:
263                                        _bmappos = self._filehandle.tell()-6
264                                    elif _bmapflag == 254:
265                                        pass # Do this to keep the previous position value
266                                    else:
267                                        _bmappos = None
268                                    self._filehandle.seek(self._filehandle.tell()+secsize-6)
269                                elif secnum == 7:
270                                    # Unpack Section 7. No need to read it, just index the position in file.
271                                    _datapos = self._filehandle.tell()-5
272                                    _datasize = secsize
273                                    self._filehandle.seek(self._filehandle.tell()+secsize-5)
274                                else:
275                                    self._filehandle.seek(self._filehandle.tell()+secsize-5)
276                            else:
277                                if num == 2 and secnum == 3:
278                                    pass # Allow this.  Just means no Local Use Section.
279                                else:
280                                    _issubmessage = True
281                                    _submsgoffset = (self._filehandle.tell()-5)-(self._index['offset'][-1])
282                                    _submsgbegin = secnum
283                                self._filehandle.seek(self._filehandle.tell()-5)
284                                continue
285                        trailer = struct.unpack('>4s',self._filehandle.read(4))[0]
286                        if trailer == b'7777':
287                            self.messages += 1
288                            self._index['offset'].append(pos)
289                            self._index['bitmap_offset'].append(_bmappos)
290                            self._index['data_offset'].append(_datapos)
291                            self._index['size'].append(section0[-1])
292                            self._index['data_size'].append(_datasize)
293                            self._index['messageNumber'].append(self.messages)
294                            self._index['isSubmessage'].append(_issubmessage)
295                            if _issubmessage:
296                                self._index['submessageOffset'].append(_submsgoffset)
297                                self._index['submessageBeginSection'].append(_submsgbegin)
298                            else:
299                                self._index['submessageOffset'].append(0)
300                                self._index['submessageBeginSection'].append(_submsgbegin)
301
302                            # Create Grib2Message with data.
303                            msg = Grib2Message(section0,section1,section2,section3,section4,section5,_bmapflag)
304                            msg._msgnum = self.messages-1
305                            msg._deflist = _deflist
306                            msg._coordlist = _coordlist
307                            if not no_data:
308                                shape = (msg.ny,msg.nx)
309                                ndim = 2
310                                if msg.typeOfValues == 0:
311                                    dtype = 'float32'
312                                elif msg.typeOfValues == 1:
313                                    dtype = 'int32'
314                                msg._data = Grib2MessageOnDiskArray(shape, ndim, dtype, self._filehandle,
315                                                                    msg, pos, _bmappos, _datapos)
316                            self._index['msg'].append(msg)
317
318                            break
319                        else:
320                            self._filehandle.seek(self._filehandle.tell()-4)
321                            self.messages += 1
322                            self._index['offset'].append(pos)
323                            self._index['bitmap_offset'].append(_bmappos)
324                            self._index['data_offset'].append(_datapos)
325                            self._index['size'].append(section0[-1])
326                            self._index['data_size'].append(_datasize)
327                            self._index['messageNumber'].append(self.messages)
328                            self._index['isSubmessage'].append(_issubmessage)
329                            self._index['submessageOffset'].append(_submsgoffset)
330                            self._index['submessageBeginSection'].append(_submsgbegin)
331
332                            # Create Grib2Message with data.
333                            msg = Grib2Message(section0,section1,section2,section3,section4,section5,_bmapflag)
334                            msg._msgnum = self.messages-1
335                            msg._deflist = _deflist
336                            msg._coordlist = _coordlist
337                            if not no_data:
338                                shape = (msg.ny,msg.nx)
339                                ndim = 2
340                                if msg.typeOfValues == 0:
341                                    dtype = 'float32'
342                                elif msg.typeOfValues == 1:
343                                    dtype = 'int32'
344                                msg._data = Grib2MessageOnDiskArray(shape, ndim, dtype, self._filehandle,
345                                                                    msg, pos, _bmappos, _datapos)
346                            self._index['msg'].append(msg)
347
348                            continue
349
350            except(struct.error):
351                if 'r' in self.mode:
352                    self._filehandle.seek(0)
353                break
354
355        # Index at end of _build_index()
356        if self._hasindex and not no_data:
357             self.variables = tuple(sorted(set([msg.shortName for msg in self._index['msg']])))
358             self.levels = tuple(sorted(set([msg.level for msg in self._index['msg']])))
359
360
361    def close(self):
362        """
363        Close the file handle
364        """
365        if not self._filehandle.closed:
366            self.messages = 0
367            self.current_message = 0
368            self._filehandle.close()
369            self.closed = self._filehandle.closed
370
371
372    def read(self, size=None):
373        """
374        Read size amount of GRIB2 messages from the current position. If no argument is
375        given, then size is None and all messages are returned from the current position
376        in the file. This read method follows the behavior of Python's builtin open()
377        function, but whereas that operates on units of bytes, we operate on units of
378        GRIB2 messages.
379
380        Parameters
381        ----------
382
383        **`size : int, optional`**
384
385        The number of GRIB2 messages to read from the current position. If no argument is
386        give, the default value is `None` and remainder of the file is read.
387
388        Returns
389        -------
390
391        `Grib2Message` object when size = 1 or a `list` of Grib2Messages when
392        size > 1.
393        """
394        if size is not None and size < 0:
395            size = None
396        if size is None or size > 1:
397            start = self.tell()
398            stop = self.messages if size is None else start+size
399            if size is None:
400                self.current_message = self.messages-1
401            else:
402                self.current_message += size
403            return self._index['msg'][slice(start,stop,1)]
404        elif size == 1:
405            self.current_message += 1
406            return self._index['msg'][self.current_message]
407        else:
408            None
409
410
411    def seek(self, pos):
412        """
413        Set the position within the file in units of GRIB2 messages.
414
415        Parameters
416        ----------
417
418        **`pos : int`**
419
420        The GRIB2 Message number to set the file pointer to.
421        """
422        if self._hasindex:
423            self._filehandle.seek(self._index['offset'][pos])
424            self.current_message = pos
425
426
427    def tell(self):
428        """
429        Returns the position of the file in units of GRIB2 Messages.
430        """
431        return self.current_message
432
433
434    def select(self,**kwargs):
435        """
436        Select GRIB2 messages by `Grib2Message` attributes.
437        """
438        # TODO: Added ability to process multiple values for each keyword (attribute)
439        idxs = []
440        nkeys = len(kwargs.keys())
441        for k,v in kwargs.items():
442            idxs += [msg._msgnum for msg in self._index['msg'] if getattr(msg,k) == v]
443        idxs = np.array(idxs,dtype=np.int32)
444        return [self._index['msg'][i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
445
446
447    def write(self, msg):
448        """
449        Writes GRIB2 message object to file.
450
451        Parameters
452        ----------
453
454        **`msg : Grib2Message or sequence of Grib2Messages`**
455
456        GRIB2 message objects to write to file.
457        """
458        if isinstance(msg,list):
459            for m in msg:
460                self.write(m)
461            return
462
463        if issubclass(msg.__class__,_Grib2Message):
464            if hasattr(msg,'_msg'):
465                self._filehandle.write(msg._msg)
466            else:
467                if msg._signature != msg._generate_signature():
468                    msg.pack()
469                    self._filehandle.write(msg._msg)
470                else:
471                    if hasattr(msg._data,'filehandle'):
472                        msg._data.filehandle.seek(msg._data.offset)
473                        self._filehandle.write(msg._data.filehandle.read(msg.section0[-1]))
474                    else:
475                        msg.pack()
476                        self._filehandle.write(msg._msg)
477            self.flush()
478            self.size = os.path.getsize(self.name)
479            self._filehandle.seek(self.size-msg.section0[-1])
480            self._build_index()
481        else:
482            raise TypeError("msg must be a Grib2Message object.")
483        return
484
485
486    def flush(self):
487        """
488        Flush the file object buffer.
489        """
490        self._filehandle.flush()
491
492
493    def levels_by_var(self,name):
494        """
495        Return a list of level strings given a variable shortName.
496
497        Parameters
498        ----------
499
500        **`name : str`**
501
502        Grib2Message variable shortName
503
504        Returns
505        -------
506
507        A list of strings of unique level strings.
508        """
509        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
510
511
512    def vars_by_level(self,level):
513        """
514        Return a list of variable shortName strings given a level.
515
516        Parameters
517        ----------
518
519        **`level : str`**
520
521        Grib2Message variable level
522
523        Returns
524        -------
525
526        A list of strings of variable shortName strings.
527        """
528        return list(sorted(set([msg.shortName for msg in self.select(level=level)])))

531class Grib2Message:
532    """
533    """
534    def __new__(self, section0: np.array = np.array([struct.unpack('>I',b'GRIB')[0],0,0,2,0]),
535                      section1: np.array = np.zeros((13),dtype=np.int64),
536                      section2: bytes = None,
537                      section3: np.array = None,
538                      section4: np.array = None,
539                      section5: np.array = None, *args, **kwargs):
540
541        bases = list()
542        if section3 is None:
543            if 'gdtn' in kwargs.keys():
544                gdtn = kwargs['gdtn']
545                Gdt = templates.gdt_class_by_gdtn(gdtn)
546                bases.append(Gdt)
547                section3 = np.zeros((Gdt._len+5),dtype=np.int64)
548            else:
549                raise ValueError("Must provide GRIB2 Grid Definition Template Number or section 3 array")
550        else:
551            gdtn = section3[4]
552            Gdt = templates.gdt_class_by_gdtn(gdtn)
553            bases.append(Gdt)
554
555        if section4 is None:
556            if 'pdtn' in kwargs.keys():
557                pdtn = kwargs['pdtn']
558                Pdt = templates.pdt_class_by_pdtn(pdtn)
559                bases.append(Pdt)
560                section4 = np.zeros((Pdt._len+2),dtype=np.int64)
561            else:
562                raise ValueError("Must provide GRIB2 Production Definition Template Number or section 4 array")
563        else:
564            pdtn = section4[1]
565            Pdt = templates.pdt_class_by_pdtn(pdtn)
566            bases.append(Pdt)
567
568        if section5 is None:
569            if 'drtn' in kwargs.keys():
570                drtn = kwargs['drtn']
571                Drt = templates.drt_class_by_drtn(drtn)
572                bases.append(Drt)
573                section5 = np.zeros((Drt._len+2),dtype=np.int64)
574            else:
575                raise ValueError("Must provide GRIB2 Data Representation Template Number or section 5 array")
576        else:
577            drtn = section5[1]
578            Drt = templates.drt_class_by_drtn(drtn)
579            bases.append(Drt)
580
581        # attempt to use existing Msg class if it has already been made with gdtn,pdtn,drtn combo
582        try:
583            Msg = _msg_class_store[f"{gdtn}:{pdtn}:{drtn}"]
584        except KeyError:
585            @dataclass(init=False, repr=False)
586            class Msg(_Grib2Message, *bases):
587                pass
588            _msg_class_store[f"{gdtn}:{pdtn}:{drtn}"] = Msg
589
590
591
592        return Msg(section0, section1, section2, section3, section4, section5, *args)

14def show_config():
15    """
16    Print grib2io build configuration information.
17    """
18    from g2clib import __version__ as g2clib_version
19    from g2clib import _has_png as have_png
20    from g2clib import _has_jpeg as have_jpeg
21    print("grib2io version %s Configuration:\n"%(__version__))
22    print("\tg2c library version:".expandtabs(4),g2clib_version)
23    print("\tJPEG compression support:".expandtabs(4),bool(have_jpeg))
24    print("\tPNG compression support:".expandtabs(4),bool(have_png))

1224def interpolate(a, method, grid_def_in, grid_def_out, method_options=None):
1225    """
1226    Perform grid spatial interpolation via the [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip).
1227
1228    Parameters
1229    ----------
1230
1231    **`a : numpy.ndarray`**
1232
1233    Array data to interpolate from. These data are expected to be in
1234    2-dimensional form with shape (ny, nx) or 3-dimensional where the
1235    3rd dimension represents another spatial, temporal, or classification
1236    (i.e. ensemble members) dimension. The function will properly flatten
1237    the array that is acceptable for the NCEPLIBS-ip interpolation
1238    subroutines.
1239
1240    **`method : int or str`**
1241
1242    Interpolate method to use. This can either be an integer or string using
1243    the following mapping:
1244
1245    | Interpolate Scheme | Integer Value |
1246    | :---:              | :---:         |
1247    | 'bilinear'         | 0             |
1248    | 'bicubic'          | 1             |
1249    | 'neighbor'         | 2             |
1250    | 'budget'           | 3             |
1251    | 'spectral'         | 4             |
1252    | 'neighbor-budget'  | 6             |
1253
1254    **`grid_def_in : grib2io.Grib2GridDef`**
1255
1256    Grib2GridDef object of the input grid.
1257
1258    **`grid_def_out : grib2io.Grib2GridDef`**
1259
1260    Grib2GridDef object of the output grid.
1261
1262    **`method_options : list of ints, optional`**
1263
1264    Interpolation options. See the NCEPLIBS-ip doucmentation for
1265    more information on how these are used.
1266    """
1267    from . import _interpolate
1268
1269    interp_schemes = {'bilinear':0, 'bicubic':1, 'neighbor':2,
1270                      'budget':3, 'spectral':4, 'neighbor-budget':6}
1271
1272    if isinstance(method,int) and method not in interp_schemes.values():
1273        raise ValueError('Invalid interpolation method.')
1274    elif isinstance(method,str):
1275        if method in interp_schemes.keys():
1276            method = interp_schemes[method]
1277        else:
1278            raise ValueError('Invalid interpolation method.')
1279
1280    if method_options is None:
1281        method_options = np.zeros((20),dtype=np.int32)
1282        if method == 3:
1283            method_options[0:2] = -1
1284
1285    ni = grid_def_in.nx*grid_def_in.ny
1286    no = grid_def_out.nx*grid_def_out.ny
1287
1288    if len(a.shape) == 2 and a.shape == (grid_def_in.ny,grid_def_in.nx):
1289        newshp = (grid_def_out.ny,grid_def_out.nx)
1290        a = np.expand_dims(a.flatten(),axis=0)
1291    elif len(a.shape) == 3 and a.shape[-2:] == (grid_def_in.ny,grid_def_in.nx):
1292        newshp = (a.shape[0],grid_def_out.ny,grid_def_out.nx)
1293        a = a.reshape(*a.shape[:-2],-1)
1294    else:
1295        raise ValueError("Array shape must be either (ny,nx) or (:,ny,nx).")
1296
1297    ibi = np.zeros((a.shape[0]),dtype=np.int32)
1298    li = np.zeros(a.shape,dtype=np.int32)
1299    go = np.zeros((a.shape[0],grid_def_out.ny*grid_def_out.nx),dtype=np.float32)
1300
1301    no,ibo,lo,iret = _interpolate.interpolate(method,method_options,
1302                                              grid_def_in.gdtn,grid_def_in.gdt,
1303                                              grid_def_out.gdtn,grid_def_out.gdt,
1304                                              ibi,li.T,a.T,go.T)
1305
1306    return go.reshape(newshp)

1309@dataclass
1310class Grib2GridDef:
1311    """
1312    Class to hold GRIB2 Grid Definition Template Number and Template as
1313    class attributes. This allows for cleaner looking code when passing these
1314    metadata around.  For example, the `grib2io._Grib2Message.interpolate`
1315    method and `grib2io.interpolate` function accepts these objects.
1316    """
1317    gdtn: int
1318    gdt: np.array
1319
1320    @classmethod
1321    def from_section3(cls, section3):
1322        return cls(section3[4],section3[5:])
1323
1324    @property
1325    def nx(self):
1326        return self.gdt[7]
1327
1328    @property
1329    def ny(self):
1330        return self.gdt[8]
1331
1332    @property
1333    def npoints(self):
1334        return self.gdt[7] * self.gdt[8]