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}')

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

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