95class open():
 96    """
 97    GRIB2 File Object.
 98
 99    This class can accommodate a physical file with one or more GRIB2
100    messages or a bytes object containing a GRIB2 messages.
101
102    A physical file can contain one or more GRIB2 messages.  When instantiated,
103    class `grib2io.open`, the file named `filename` is opened for reading (`mode
104    = 'r'`) and is automatically indexed.  The indexing procedure reads some of
105    the GRIB2 metadata for all GRIB2 Messages.  A GRIB2 Message may contain
106    submessages whereby Section 2-7 can be repeated.  grib2io accommodates for
107    this by flattening any GRIB2 submessages into multiple individual messages.
108
109    It is important to note that GRIB2 files from some Meteorological agencies
110    contain other data than GRIB2 messages.  GRIB2 files from ECMWF can contain
111    GRIB1 and GRIB2 messages.  grib2io checks for these and safely ignores them.
112
113    Attributes
114    ----------
115    closed : bool
116        `True` is file handle is close; `False` otherwise.
117    current_message : int
118        Current position of the file in units of GRIB2 Messages. (read only)
119    levels : tuple
120        Tuple containing a unique list of wgrib2-formatted level/layer strings.
121    messages : int
122        Count of GRIB2 Messages contained in the file.
123    mode : str
124        File IO mode of opening the file.
125    name : str
126        Full path name of the GRIB2 file.
127    save_index : bool
128        Whether to save a pickle-based index file for the GRIB2 file. Default is `True`.
129    size : int
130        Size of the file in units of bytes.
131    use_index
132        Whether to use an existing pickle-based index file for the GRIB2 file. Default is `True`.
133    variables : tuple
134        Tuple containing a unique list of variable short names (i.e. GRIB2
135        abbreviation names).
136    """
137
138    __slots__ = ('_fileid', '_filehandle', '_hasindex', '_index',
139                 '_pos', 'closed', 'current_message', 'messages', 'mode',
140                 'name', 'size', 'save_index', 'use_index', '_msgs')
141
142    def __init__(
143        self,
144        filename: Union[bytes, str, Path],
145        mode: Literal["r", "w", "x"] = "r",
146        *,
147        save_index = True,
148        use_index = True,
149        _xarray_backend = False,
150        **kwargs,
151        ):
152        """
153        Initialize GRIB2 File object instance.
154
155        Parameters
156        ----------
157        filename
158            File path containing GRIB2 messages OR bytes.
159        mode
160            File access mode where "r" opens the files for reading only; "w"
161            opens the file for overwriting and "x" for writing to a new file.
162        save_index
163            Whether to save a pickle-based index file for the GRIB2 file. Default is True.
164        use_index
165            Whether to use an existing pickle-based index file for the GRIB2 file. Default is True.
166        """
167
168        # All write modes are read/write.
169        # All modes are binary.
170        if mode in ("a", "x", "w"):
171            mode += "+"
172        mode = mode + "b"
173
174        self.mode = mode
175        self.save_index = save_index
176        self.use_index = use_index
177
178        if isinstance(filename, bytes):
179            if 'r' not in self.mode:
180                raise ValueError("bytes grib2 can only be opened as read")
181            self.current_message = 0
182            if filename[:2] == _GZIP_HEADER:
183                filename = gzip.decompress(filename)
184            if filename[:4]+filename[-4:] != b'GRIB7777':
185                raise ValueError("Invalid GRIB bytes")
186            self._filehandle = BytesIO(filename)
187            self.name = "<in-memory-file>"
188            self.size = len(filename)
189            self._fileid = hashlib.sha1((self.name+str(self.size)).encode('ASCII')).hexdigest()
190            self._index = build_index(self._filehandle)
191            self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
192            self.messages = len(self._msgs)
193
194        else:
195            self.current_message = 0
196            if 'r' in mode:
197                self._filehandle = builtins.open(filename, mode=mode)
198                # Some GRIB2 files are gzipped, so check for that here, but
199                # raise error when using xarray backend.
200                # Gzip files contain a 2-byte header b'\x1f\x8b'.
201                if self._filehandle.read(2) == _GZIP_HEADER:
202                    self._filehandle.close()
203                    if _xarray_backend:
204                        raise RuntimeError('Gzip GRIB2 files are not supported by the Xarray backend.')
205                    self._filehandle = gzip.open(filename, mode=mode)
206                else:
207                    self._filehandle = builtins.open(filename, mode=mode)
208            else:
209                self._filehandle = builtins.open(filename, mode=mode)
210                self.name = os.path.abspath(filename)
211            self.name = os.path.abspath(filename)
212            fstat = os.stat(self.name)
213            self.size = fstat.st_size
214            self._fileid = hashlib.sha1((self.name+str(fstat.st_ino)+
215                                     str(self.size)).encode('ASCII')).hexdigest()
216
217            if 'r' in self.mode:
218
219                indexfile = f"{self.name}_{self._fileid}.grib2ioidx"
220                if self.use_index and os.path.exists(indexfile):
221                    try:
222                        with builtins.open(indexfile, 'rb') as file:
223                            index = pickle.load(file)
224                        self._index = index
225                    except Exception as e:
226                        print(f"found indexfile: {indexfile}, but unable to load it: {e}")
227                        print(f"re-forming index from grib2file, but not writing indexfile")
228                        self._index = build_index(self._filehandle)
229                else:
230                    self._index = build_index(self._filehandle)
231                    if self.save_index:
232                        try:
233                            serialize_index(self._index, indexfile)
234                        except Exception as e:
235                            print(f"index was not serialized for future use: {e}")
236
237                self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
238
239                self.messages = len(self._msgs)
240            elif 'w' or 'x' in self.mode:
241                self.messages = 0
242                self.current_message = None
243
244        self.closed = self._filehandle.closed
245
246
247    def __delete__(self, instance):
248        self.close()
249        del self._index
250
251
252    def __enter__(self):
253        return self
254
255
256    def __exit__(self, atype, value, traceback):
257        self.close()
258
259
260    def __iter__(self):
261        yield from self._msgs
262
263
264    def __len__(self):
265        return self.messages
266
267
268    def __repr__(self):
269        strings = []
270        for k in self.__slots__:
271            if k.startswith('_'): continue
272            strings.append('%s = %s\n'%(k,eval('self.'+k)))
273        return ''.join(strings)
274
275
276    def __getitem__(self, key):
277        if isinstance(key,int):
278            if abs(key) >= len(self._msgs):
279                raise IndexError("index out of range")
280            else:
281                return self._msgs[key]
282        elif isinstance(key,str):
283            return self.select(shortName=key)
284        elif isinstance(key,slice):
285            return self._msgs[key]
286        else:
287            raise KeyError('Key must be an integer, slice, or GRIB2 variable shortName.')
288
289
290
291    @property
292    def levels(self):
293        """Provides a unique tuple of level strings."""
294        if self._hasindex:
295            return tuple(sorted(set([msg.level for msg in self._msgs])))
296        else:
297            return None
298
299
300    @property
301    def variables(self):
302        """Provides a unique tuple of variable shortName strings."""
303        if self._hasindex:
304            return tuple(sorted(set([msg.shortName for msg in self._msgs])))
305        else:
306            return None
307
308
309    def close(self):
310        """Close the file handle."""
311        if not self._filehandle.closed:
312            self.messages = 0
313            self.current_message = 0
314            self._filehandle.close()
315            self.closed = self._filehandle.closed
316
317
318    def read(self, size: Optional[int]=None):
319        """
320        Read size amount of GRIB2 messages from the current position.
321
322        If no argument is given, then size is None and all messages are returned
323        from the current position in the file. This read method follows the
324        behavior of Python's builtin open() function, but whereas that operates
325        on units of bytes, we operate on units of GRIB2 messages.
326
327        Parameters
328        ----------
329        size: default=None
330            The number of GRIB2 messages to read from the current position. If
331            no argument is give, the default value is None and remainder of
332            the file is read.
333
334        Returns
335        -------
336        read
337            ``Grib2Message`` object when size = 1 or a list of Grib2Messages
338            when size > 1.
339        """
340        if size is not None and size < 0:
341            size = None
342        if size is None or size > 1:
343            start = self.tell()
344            stop = self.messages if size is None else start+size
345            if size is None:
346                self.current_message = self.messages-1
347            else:
348                self.current_message += size
349            return self._msgs[slice(start,stop,1)]
350        elif size == 1:
351            self.current_message += 1
352            return self._msgs[self.current_message]
353        else:
354            None
355
356
357    def seek(self, pos: int):
358        """
359        Set the position within the file in units of GRIB2 messages.
360
361        Parameters
362        ----------
363        pos
364            The GRIB2 Message number to set the file pointer to.
365        """
366        if self._hasindex:
367            self._filehandle.seek(self._index['sectionOffset'][0][pos])
368            self.current_message = pos
369
370
371    def tell(self):
372        """Returns the position of the file in units of GRIB2 Messages."""
373        return self.current_message
374
375
376    def select(self, **kwargs):
377        """Select GRIB2 messages by `Grib2Message` attributes."""
378        # TODO: Added ability to process multiple values for each keyword (attribute)
379        idxs = []
380        nkeys = len(kwargs.keys())
381        for k,v in kwargs.items():
382            for m in self._msgs:
383                if hasattr(m,k) and getattr(m,k) == v: idxs.append(m._msgnum)
384        idxs = np.array(idxs,dtype=np.int32)
385        return [self._msgs[i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
386
387
388    def write(self, msg):
389        """
390        Writes GRIB2 message object to file.
391
392        Parameters
393        ----------
394        msg
395            GRIB2 message objects to write to file.
396        """
397        if isinstance(msg,list):
398            for m in msg:
399                self.write(m)
400            return
401
402        if issubclass(msg.__class__,_Grib2Message):
403            # TODO: We can consider letting pack return packed bytes instead of associating with message object
404            if hasattr(msg,'_msg'):
405                # write already packed bytes
406                self._filehandle.write(msg._msg)
407            else:
408                if msg._signature == msg._generate_signature() and \
409                   msg._data is None and \
410                   hasattr(msg._ondiskarray,'filehandle'):
411                   # write unchanged message from input
412                   offset = msg._ondiskarray.filehandle.tell()
413                   msg._ondiskarray.filehandle.seek(msg._ondiskarray.offset)
414                   self._filehandle.write(msg._ondiskarray.filehandle.read(msg.section0[-1]))
415                   msg._ondiskarray.filehandle.seek(offset)
416                else:
417                    msg.pack()
418                    self._filehandle.write(msg._msg)
419            self.size = os.path.getsize(self.name)
420            self.messages += 1
421        else:
422            raise TypeError("msg must be a Grib2Message object.")
423        return
424
425
426    def flush(self):
427        """Flush the file object buffer."""
428        self._filehandle.flush()
429
430
431    def levels_by_var(self, name: str):
432        """
433        Return a list of level strings given a variable shortName.
434
435        Parameters
436        ----------
437        name
438            Grib2Message variable shortName
439
440        Returns
441        -------
442        levels_by_var
443            A list of unique level strings.
444        """
445        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
446
447
448    def vars_by_level(self, level: str):
449        """
450        Return a list of variable shortName strings given a level.
451
452        Parameters
453        ----------
454        level
455            Grib2Message variable level
456
457        Returns
458        -------
459        vars_by_level
460            A list of unique variable shortName strings.
461        """
462        return list(sorted(set([msg.shortName for msg in self.select(level=level)])))

34def show_config():
35    """Print grib2io build configuration information."""
36    print(f'grib2io version {__version__} Configuration:')
37    print(f'')
38    print(f'NCEPLIBS-g2c library version: {__g2clib_version__}')
39    print(f'\tStatic library: {g2c_static}')
40    print(f'\tJPEG compression support: {has_jpeg_support}')
41    print(f'\tPNG compression support: {has_png_support}')
42    print(f'\tAEC compression support: {has_aec_support}')
43    print(f'')
44    print(f'NCEPLIPS-ip support: {has_interpolation}')
45    print(f'\tStatic library: {ip_static}')
46    print(f'\tOpenMP support: {has_openmp_support}')
47    print(f'')
48    print(f'Static libs:')
49    for lib in extra_objects:
50        print(f'\t{lib}')
51    print(f'')
52    print(f'NCEP GRIB2 Table Version: {_ncep_grib2_table_version}')

1795def interpolate(a,
1796    method: Union[int, str],
1797    grid_def_in,
1798    grid_def_out,
1799    method_options=None,
1800    num_threads=1,
1801):
1802    """
1803    This is the module-level interpolation function.
1804
1805    This interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
1806    through grib2io's internal iplib Cython extension module.
1807
1808    Parameters
1809    ----------
1810    a : numpy.ndarray or tuple
1811        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
1812        performed.  If `a` is a `tuple`, then vector interpolation will be
1813        performed with the assumption that u = a[0] and v = a[1] and are both
1814        `numpy.ndarray`.
1815
1816        These data are expected to be in 2-dimensional form with shape (ny, nx)
1817        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
1818        spatial, temporal, or classification (i.e. ensemble members) dimension.
1819        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
1820        acceptable for input into the interpolation subroutines. If needed, these
1821        data will be converted to `np.float32`.
1822    method
1823        Interpolate method to use. This can either be an integer or string using
1824        the following mapping:
1825
1826        | Interpolate Scheme | Integer Value |
1827        | :---:              | :---:         |
1828        | 'bilinear'         | 0             |
1829        | 'bicubic'          | 1             |
1830        | 'neighbor'         | 2             |
1831        | 'budget'           | 3             |
1832        | 'spectral'         | 4             |
1833        | 'neighbor-budget'  | 6             |
1834
1835    grid_def_in : grib2io.Grib2GridDef
1836        Grib2GridDef object for the input grid.
1837    grid_def_out : grib2io.Grib2GridDef
1838        Grib2GridDef object for the output grid or station points.
1839    method_options : list of ints, optional
1840        Interpolation options. See the NCEPLIBS-ip documentation for
1841        more information on how these are used.
1842    num_threads : int, optional
1843        Number of OpenMP threads to use for interpolation. The default
1844        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
1845        was not built with OpenMP, then this keyword argument and value
1846        will have no impact.
1847
1848    Returns
1849    -------
1850    interpolate
1851        Returns a `numpy.ndarray` of dtype `np.float32` when scalar interpolation
1852        is performed or a `tuple` of `numpy.ndarray`s when vector interpolation is
1853        performed with the assumptions that 0-index is the interpolated u and
1854        1-index is the interpolated v.
1855    """
1856
1857    from . import iplib
1858
1859    prev_num_threads = 1
1860    try:
1861        prev_num_threads = iplib.openmp_get_num_threads()
1862        iplib.openmp_set_num_threads(num_threads)
1863    except(AttributeError):
1864        pass
1865
1866    print(f"grib2io.interpolate thread report: OpenMP num threads = {iplib.openmp_get_num_threads()}")
1867
1868    if isinstance(method,int) and method not in _interp_schemes.values():
1869        raise ValueError('Invalid interpolation method.')
1870    elif isinstance(method,str):
1871        if method in _interp_schemes.keys():
1872            method = _interp_schemes[method]
1873        else:
1874            raise ValueError('Invalid interpolation method.')
1875
1876    if method_options is None:
1877        method_options = np.zeros((20),dtype=np.int32)
1878        if method in {3,6}:
1879            method_options[0:2] = -1
1880
1881    mi = grid_def_in.npoints
1882    mo = grid_def_out.npoints
1883
1884    # Adjust shape of input array(s)
1885    a, newshp = _adjust_array_shape_for_interp(a, grid_def_in, grid_def_out)
1886
1887    # Call interpolation subroutines according to type of a.
1888    if isinstance(a,np.ndarray):
1889        # Scalar
1890        km = a.shape[0]
1891        if np.any(np.isnan(a)):
1892            ibi = np.ones((km), dtype=np.int32)
1893            li = np.where(np.isnan(a),0,1).astype(np.uint8)
1894        else:
1895            ibi = np.zeros((km), dtype=np.int32)
1896            li = np.zeros(a.shape,dtype=np.uint8)
1897        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
1898            method,
1899            method_options,
1900            grid_def_in.gdtn,
1901            np.array(grid_def_in.gdt, dtype=np.int32),
1902            grid_def_out.gdtn,
1903            np.array(grid_def_out.gdt, dtype=np.int32),
1904            mi,
1905            mo,
1906            km,
1907            ibi,
1908            li,
1909            a,
1910        )
1911        out = np.where(lo==0,np.nan,go).reshape(newshp)
1912    elif isinstance(a,tuple):
1913        # Vector
1914        km = a[0].shape[0]
1915        if np.any(np.isnan(a)):
1916            ibi = np.ones((km), dtype=np.int32)
1917            li = np.where(np.isnan(a),0,1).astype(np.uint8)
1918        else:
1919            ibi = np.zeros((km), dtype=np.int32)
1920            li = np.zeros(a[0].shape,dtype=np.uint8)
1921        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
1922            method,
1923            method_options,
1924            grid_def_in.gdtn,
1925            np.array(grid_def_in.gdt, dtype=np.int32),
1926            grid_def_out.gdtn,
1927            np.array(grid_def_out.gdt, dtype=np.int32),
1928            mi,
1929            mo,
1930            km,
1931            ibi,
1932            li,
1933            a[0],
1934            a[1],
1935        )
1936        uo = np.where(lo==0,np.nan,uo).reshape(newshp)
1937        vo = np.where(lo==0,np.nan,vo).reshape(newshp)
1938        out = (uo,vo)
1939
1940    try:
1941        iplib.openmp_set_num_threads(prev_num_threads)
1942    except(AttributeError):
1943        pass
1944
1945    return out

1948def interpolate_to_stations(
1949    a,
1950    method: Union[int, str],
1951    grid_def_in,
1952    lats,
1953    lons,
1954    method_options=None,
1955    num_threads=1
1956):
1957    """
1958    Module-level interpolation function for interpolation to stations.
1959
1960    Interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
1961    via grib2io's iplib Cython exntension module. It supports scalar and
1962    vector interpolation according to the type of object a.
1963
1964    Parameters
1965    ----------
1966    a : numpy.ndarray or tuple
1967        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
1968        performed.  If `a` is a `tuple`, then vector interpolation will be
1969        performed with the assumption that u = a[0] and v = a[1] and are both
1970        `numpy.ndarray`.
1971
1972        These data are expected to be in 2-dimensional form with shape (ny, nx)
1973        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
1974        spatial, temporal, or classification (i.e. ensemble members) dimension.
1975        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
1976        acceptable for input into the interpolation subroutines. If needed, these
1977        data will be converted to `np.float32`.
1978    method
1979        Interpolate method to use. This can either be an integer or string using
1980        the following mapping:
1981
1982        | Interpolate Scheme | Integer Value |
1983        | :---:              | :---:         |
1984        | 'bilinear'         | 0             |
1985        | 'bicubic'          | 1             |
1986        | 'neighbor'         | 2             |
1987        | 'budget'           | 3             |
1988        | 'spectral'         | 4             |
1989        | 'neighbor-budget'  | 6             |
1990
1991    grid_def_in : grib2io.Grib2GridDef
1992        Grib2GridDef object for the input grid.
1993    lats : numpy.ndarray or list
1994        Latitudes for station points
1995    lons : numpy.ndarray or list
1996        Longitudes for station points
1997    method_options : list of ints, optional
1998        Interpolation options. See the NCEPLIBS-ip documentation for
1999        more information on how these are used.
2000    num_threads : int, optional
2001        Number of OpenMP threads to use for interpolation. The default
2002        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
2003        was not built with OpenMP, then this keyword argument and value
2004        will have no impact.
2005
2006    Returns
2007    -------
2008    interpolate_to_stations
2009        Returns a `numpy.ndarray` of dtype `np.float32` when scalar
2010        interpolation is performed or a `tuple` of `numpy.ndarray`s
2011        when vector interpolation is performed with the assumptions
2012        that 0-index is the interpolated u and 1-index is the
2013        interpolated v.
2014    """
2015    from . import iplib
2016
2017    prev_num_threads = 1
2018    try:
2019        prev_num_threads = iplib.openmp_get_num_threads()
2020        iplib.openmp_set_num_threads(num_threads)
2021    except(AttributeError):
2022        pass
2023
2024    if isinstance(method,int) and method not in _interp_schemes.values():
2025        raise ValueError('Invalid interpolation method.')
2026    elif isinstance(method,str):
2027        if method in _interp_schemes.keys():
2028            method = _interp_schemes[method]
2029        else:
2030            raise ValueError('Invalid interpolation method.')
2031
2032    if method_options is None:
2033        method_options = np.zeros((20),dtype=np.int32)
2034        if method in {3,6}:
2035            method_options[0:2] = -1
2036
2037    # Check lats and lons
2038    if isinstance(lats,list):
2039        nlats = len(lats)
2040    elif isinstance(lats,np.ndarray) and len(lats.shape) == 1:
2041        nlats = lats.shape[0]
2042    else:
2043        raise ValueError('Station latitudes must be a list or 1-D NumPy array.')
2044    if isinstance(lons,list):
2045        nlons = len(lons)
2046    elif isinstance(lons,np.ndarray) and len(lons.shape) == 1:
2047        nlons = lons.shape[0]
2048    else:
2049        raise ValueError('Station longitudes must be a list or 1-D NumPy array.')
2050    if nlats != nlons:
2051        raise ValueError('Station lats and lons must be same size.')
2052
2053    mi = grid_def_in.npoints
2054    mo = nlats
2055
2056    # Adjust shape of input array(s)
2057    a, newshp = _adjust_array_shape_for_interp_stations(a, grid_def_in, mo)
2058
2059    # Use gdtn = -1 for stations and an empty template array
2060    gdtn = -1
2061    gdt = np.zeros((200),dtype=np.int32)
2062
2063    # Call interpolation subroutines according to type of a.
2064    if isinstance(a,np.ndarray):
2065        # Scalar
2066        km = a.shape[0]
2067        ibi = np.zeros((km), dtype=np.int32)
2068        li = np.zeros(a.shape,dtype=np.uint8)
2069        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
2070            method,
2071            method_options,
2072            grid_def_in.gdtn,
2073            np.array(grid_def_in.gdt, dtype=np.int32),
2074            gdtn,
2075            gdt,
2076            mi,
2077            mo,
2078            km,
2079            ibi,
2080            li,
2081            a,
2082            lats=np.array(lats, dtype=np.float32),
2083            lons=np.array(lons, dtype=np.float32),
2084        )
2085        out = go.reshape(newshp)
2086
2087    elif isinstance(a,tuple):
2088        # Vector
2089        km = a[0].shape[0]
2090        ibi = np.zeros((km), dtype=np.int32)
2091        li = np.zeros(a[0].shape,dtype=np.uint8)
2092        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
2093            method,
2094            method_options,
2095            grid_def_in.gdtn,
2096            np.array(grid_def_in.gdt, dtype=np.int32),
2097            gdtn,
2098            gdt,
2099            mi,
2100            mo,
2101            km,
2102            ibi,
2103            li,
2104            a[0],
2105            a[1],
2106            lats=np.array(lats, dtype=np.float32),
2107            lons=np.array(lons, dtype=np.float32),
2108        )
2109        out = (uo.reshape(newshp),vo.reshape(newshp))
2110
2111    try:
2112        iplib.openmp_set_num_threads(prev_num_threads)
2113    except(AttributeError):
2114        pass
2115
2116    return out