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

2030def interpolate(a,
2031    method: Union[int, str],
2032    grid_def_in,
2033    grid_def_out,
2034    method_options=None,
2035    num_threads=1,
2036):
2037    """
2038    This is the module-level interpolation function.
2039
2040    This interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
2041    through grib2io's internal iplib Cython extension module.
2042
2043    Parameters
2044    ----------
2045    a : numpy.ndarray or tuple
2046        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
2047        performed.  If `a` is a `tuple`, then vector interpolation will be
2048        performed with the assumption that u = a[0] and v = a[1] and are both
2049        `numpy.ndarray`.
2050
2051        These data are expected to be in 2-dimensional form with shape (ny, nx)
2052        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
2053        spatial, temporal, or classification (i.e. ensemble members) dimension.
2054        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
2055        acceptable for input into the interpolation subroutines. If needed, these
2056        data will be converted to `np.float32`.
2057    method
2058        Interpolate method to use. This can either be an integer or string using
2059        the following mapping:
2060
2061        | Interpolate Scheme | Integer Value |
2062        | :---:              | :---:         |
2063        | 'bilinear'         | 0             |
2064        | 'bicubic'          | 1             |
2065        | 'neighbor'         | 2             |
2066        | 'budget'           | 3             |
2067        | 'spectral'         | 4             |
2068        | 'neighbor-budget'  | 6             |
2069
2070    grid_def_in : grib2io.Grib2GridDef
2071        Grib2GridDef object for the input grid.
2072    grid_def_out : grib2io.Grib2GridDef
2073        Grib2GridDef object for the output grid or station points.
2074    method_options : list of ints, optional
2075        Interpolation options. See the NCEPLIBS-ip documentation for
2076        more information on how these are used.
2077    num_threads : int, optional
2078        Number of OpenMP threads to use for interpolation. The default
2079        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
2080        was not built with OpenMP, then this keyword argument and value
2081        will have no impact.
2082
2083    Returns
2084    -------
2085    interpolate
2086        Returns a `numpy.ndarray` of dtype `np.float32` when scalar interpolation
2087        is performed or a `tuple` of `numpy.ndarray`s when vector interpolation is
2088        performed with the assumptions that 0-index is the interpolated u and
2089        1-index is the interpolated v.
2090    """
2091
2092    from . import iplib
2093
2094    prev_num_threads = 1
2095    try:
2096        prev_num_threads = iplib.openmp_get_num_threads()
2097        iplib.openmp_set_num_threads(num_threads)
2098    except(AttributeError):
2099        pass
2100
2101    print(f"grib2io.interpolate thread report: OpenMP num threads = {iplib.openmp_get_num_threads()}")
2102
2103    if isinstance(method,int) and method not in _interp_schemes.values():
2104        raise ValueError('Invalid interpolation method.')
2105    elif isinstance(method,str):
2106        if method in _interp_schemes.keys():
2107            method = _interp_schemes[method]
2108        else:
2109            raise ValueError('Invalid interpolation method.')
2110
2111    if method_options is None:
2112        method_options = np.zeros((20),dtype=np.int32)
2113        if method in {3,6}:
2114            method_options[0:2] = -1
2115
2116    mi = grid_def_in.npoints
2117    mo = grid_def_out.npoints
2118
2119    # Adjust shape of input array(s)
2120    a, newshp = _adjust_array_shape_for_interp(a, grid_def_in, grid_def_out)
2121
2122    # Call interpolation subroutines according to type of a.
2123    if isinstance(a,np.ndarray):
2124        # Scalar
2125        km = a.shape[0]
2126        if np.any(np.isnan(a)):
2127            ibi = np.ones((km), dtype=np.int32)
2128            li = np.where(np.isnan(a),0,1).astype(np.uint8)
2129        else:
2130            ibi = np.zeros((km), dtype=np.int32)
2131            li = np.zeros(a.shape,dtype=np.uint8)
2132        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
2133            method,
2134            method_options,
2135            grid_def_in.gdtn,
2136            np.array(grid_def_in.gdt, dtype=np.int32),
2137            grid_def_out.gdtn,
2138            np.array(grid_def_out.gdt, dtype=np.int32),
2139            mi,
2140            mo,
2141            km,
2142            ibi,
2143            li,
2144            a,
2145        )
2146        out = np.where(lo==0,np.nan,go).reshape(newshp)
2147    elif isinstance(a,tuple):
2148        # Vector
2149        km = a[0].shape[0]
2150        if np.any(np.isnan(a)):
2151            ibi = np.ones((km), dtype=np.int32)
2152            li = np.where(np.isnan(a),0,1).astype(np.uint8)
2153        else:
2154            ibi = np.zeros((km), dtype=np.int32)
2155            li = np.zeros(a[0].shape,dtype=np.uint8)
2156        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
2157            method,
2158            method_options,
2159            grid_def_in.gdtn,
2160            np.array(grid_def_in.gdt, dtype=np.int32),
2161            grid_def_out.gdtn,
2162            np.array(grid_def_out.gdt, dtype=np.int32),
2163            mi,
2164            mo,
2165            km,
2166            ibi,
2167            li,
2168            a[0],
2169            a[1],
2170        )
2171        uo = np.where(lo==0,np.nan,uo).reshape(newshp)
2172        vo = np.where(lo==0,np.nan,vo).reshape(newshp)
2173        out = (uo,vo)
2174
2175    try:
2176        iplib.openmp_set_num_threads(prev_num_threads)
2177    except(AttributeError):
2178        pass
2179
2180    return out

2183def interpolate_to_stations(
2184    a,
2185    method: Union[int, str],
2186    grid_def_in,
2187    lats,
2188    lons,
2189    method_options=None,
2190    num_threads=1
2191):
2192    """
2193    Module-level interpolation function for interpolation to stations.
2194
2195    Interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
2196    via grib2io's iplib Cython exntension module. It supports scalar and
2197    vector interpolation according to the type of object a.
2198
2199    Parameters
2200    ----------
2201    a : numpy.ndarray or tuple
2202        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
2203        performed.  If `a` is a `tuple`, then vector interpolation will be
2204        performed with the assumption that u = a[0] and v = a[1] and are both
2205        `numpy.ndarray`.
2206
2207        These data are expected to be in 2-dimensional form with shape (ny, nx)
2208        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
2209        spatial, temporal, or classification (i.e. ensemble members) dimension.
2210        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
2211        acceptable for input into the interpolation subroutines. If needed, these
2212        data will be converted to `np.float32`.
2213    method
2214        Interpolate method to use. This can either be an integer or string using
2215        the following mapping:
2216
2217        | Interpolate Scheme | Integer Value |
2218        | :---:              | :---:         |
2219        | 'bilinear'         | 0             |
2220        | 'bicubic'          | 1             |
2221        | 'neighbor'         | 2             |
2222        | 'budget'           | 3             |
2223        | 'spectral'         | 4             |
2224        | 'neighbor-budget'  | 6             |
2225
2226    grid_def_in : grib2io.Grib2GridDef
2227        Grib2GridDef object for the input grid.
2228    lats : numpy.ndarray or list
2229        Latitudes for station points
2230    lons : numpy.ndarray or list
2231        Longitudes for station points
2232    method_options : list of ints, optional
2233        Interpolation options. See the NCEPLIBS-ip documentation for
2234        more information on how these are used.
2235    num_threads : int, optional
2236        Number of OpenMP threads to use for interpolation. The default
2237        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
2238        was not built with OpenMP, then this keyword argument and value
2239        will have no impact.
2240
2241    Returns
2242    -------
2243    interpolate_to_stations
2244        Returns a `numpy.ndarray` of dtype `np.float32` when scalar
2245        interpolation is performed or a `tuple` of `numpy.ndarray`s
2246        when vector interpolation is performed with the assumptions
2247        that 0-index is the interpolated u and 1-index is the
2248        interpolated v.
2249    """
2250    from . import iplib
2251
2252    prev_num_threads = 1
2253    try:
2254        prev_num_threads = iplib.openmp_get_num_threads()
2255        iplib.openmp_set_num_threads(num_threads)
2256    except(AttributeError):
2257        pass
2258
2259    if isinstance(method,int) and method not in _interp_schemes.values():
2260        raise ValueError('Invalid interpolation method.')
2261    elif isinstance(method,str):
2262        if method in _interp_schemes.keys():
2263            method = _interp_schemes[method]
2264        else:
2265            raise ValueError('Invalid interpolation method.')
2266
2267    if method_options is None:
2268        method_options = np.zeros((20),dtype=np.int32)
2269        if method in {3,6}:
2270            method_options[0:2] = -1
2271
2272    # Check lats and lons
2273    if isinstance(lats,list):
2274        nlats = len(lats)
2275    elif isinstance(lats,np.ndarray) and len(lats.shape) == 1:
2276        nlats = lats.shape[0]
2277    else:
2278        raise ValueError('Station latitudes must be a list or 1-D NumPy array.')
2279    if isinstance(lons,list):
2280        nlons = len(lons)
2281    elif isinstance(lons,np.ndarray) and len(lons.shape) == 1:
2282        nlons = lons.shape[0]
2283    else:
2284        raise ValueError('Station longitudes must be a list or 1-D NumPy array.')
2285    if nlats != nlons:
2286        raise ValueError('Station lats and lons must be same size.')
2287
2288    mi = grid_def_in.npoints
2289    mo = nlats
2290
2291    # Adjust shape of input array(s)
2292    a, newshp = _adjust_array_shape_for_interp_stations(a, grid_def_in, mo)
2293
2294    # Use gdtn = -1 for stations and an empty template array
2295    gdtn = -1
2296    gdt = np.zeros((200),dtype=np.int32)
2297
2298    # Call interpolation subroutines according to type of a.
2299    if isinstance(a,np.ndarray):
2300        # Scalar
2301        km = a.shape[0]
2302        ibi = np.zeros((km), dtype=np.int32)
2303        li = np.zeros(a.shape,dtype=np.uint8)
2304        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
2305            method,
2306            method_options,
2307            grid_def_in.gdtn,
2308            np.array(grid_def_in.gdt, dtype=np.int32),
2309            gdtn,
2310            gdt,
2311            mi,
2312            mo,
2313            km,
2314            ibi,
2315            li,
2316            a,
2317            lats=np.array(lats, dtype=np.float32),
2318            lons=np.array(lons, dtype=np.float32),
2319        )
2320        out = go.reshape(newshp)
2321
2322    elif isinstance(a,tuple):
2323        # Vector
2324        km = a[0].shape[0]
2325        ibi = np.zeros((km), dtype=np.int32)
2326        li = np.zeros(a[0].shape,dtype=np.uint8)
2327        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
2328            method,
2329            method_options,
2330            grid_def_in.gdtn,
2331            np.array(grid_def_in.gdt, dtype=np.int32),
2332            gdtn,
2333            gdt,
2334            mi,
2335            mo,
2336            km,
2337            ibi,
2338            li,
2339            a[0],
2340            a[1],
2341            lats=np.array(lats, dtype=np.float32),
2342            lons=np.array(lons, dtype=np.float32),
2343        )
2344        out = (uo.reshape(newshp),vo.reshape(newshp))
2345
2346    try:
2347        iplib.openmp_set_num_threads(prev_num_threads)
2348    except(AttributeError):
2349        pass
2350
2351    return out