76class open():
 77    """
 78    GRIB2 File Object.
 79
 80    A physical file can contain one or more GRIB2 messages.  When instantiated,
 81    class `grib2io.open`, the file named `filename` is opened for reading (`mode
 82    = 'r'`) and is automatically indexed.  The indexing procedure reads some of
 83    the GRIB2 metadata for all GRIB2 Messages.  A GRIB2 Message may contain
 84    submessages whereby Section 2-7 can be repeated.  grib2io accommodates for
 85    this by flattening any GRIB2 submessages into multiple individual messages.
 86
 87    It is important to note that GRIB2 files from some Meteorological agencies
 88    contain other data than GRIB2 messages.  GRIB2 files from ECMWF can contain
 89    GRIB1 and GRIB2 messages.  grib2io checks for these and safely ignores them.
 90
 91    Attributes
 92    ----------
 93    closed : bool
 94        `True` is file handle is close; `False` otherwise.
 95    current_message : int
 96        Current position of the file in units of GRIB2 Messages.
 97    levels : tuple
 98        Tuple containing a unique list of wgrib2-formatted level/layer strings.
 99    messages : int
100        Count of GRIB2 Messages contained in the file.
101    mode : str
102        File IO mode of opening the file.
103    name : str
104        Full path name of the GRIB2 file.
105    size : int
106        Size of the file in units of bytes.
107    variables : tuple
108        Tuple containing a unique list of variable short names (i.e. GRIB2
109        abbreviation names).
110    """
111
112    __slots__ = ('_fileid', '_filehandle', '_hasindex', '_index', '_nodata',
113                 '_pos', 'closed', 'current_message', 'messages', 'mode',
114                 'name', 'size')
115
116    def __init__(self, filename: str, mode: Literal["r", "w", "x"] = "r", **kwargs):
117        """
118        Initialize GRIB2 File object instance.
119
120        Parameters
121        ----------
122        filename
123            File name containing GRIB2 messages.
124        mode: default="r"
125            File access mode where "r" opens the files for reading only; "w"
126            opens the file for overwriting and "x" for writing to a new file.
127        """
128
129        # Manage keywords
130        if "_xarray_backend" not in kwargs:
131            kwargs["_xarray_backend"] = False
132            self._nodata = False
133        else:
134            self._nodata = kwargs["_xarray_backend"]
135
136        # All write modes are read/write.
137        # All modes are binary.
138        if mode in ("a", "x", "w"):
139            mode += "+"
140        mode = mode + "b"
141
142        # Some GRIB2 files are gzipped, so check for that here, but
143        # raise error when using xarray backend.
144        if 'r' in mode:
145            self._filehandle = builtins.open(filename, mode=mode)
146            # Gzip files contain a 2-byte header b'\x1f\x8b'.
147            if self._filehandle.read(2) == b'\x1f\x8b':
148                self._filehandle.close()
149                if kwargs["_xarray_backend"]:
150                    raise RuntimeError('Gzip GRIB2 files are not supported by the Xarray backend.')
151                import gzip
152                self._filehandle = gzip.open(filename, mode=mode)
153            else:
154                self._filehandle = builtins.open(filename, mode=mode)
155        else:
156            self._filehandle = builtins.open(filename, mode=mode)
157
158        self.name = os.path.abspath(filename)
159        fstat = os.stat(self.name)
160        self._hasindex = False
161        self._index = {}
162        self.mode = mode
163        self.messages = 0
164        self.current_message = 0
165        self.size = fstat.st_size
166        self.closed = self._filehandle.closed
167        self._fileid = hashlib.sha1((self.name+str(fstat.st_ino)+
168                                     str(self.size)).encode('ASCII')).hexdigest()
169        if 'r' in self.mode:
170            self._build_index()
171        # FIX: Cannot perform reads on mode='a'
172        #if 'a' in self.mode and self.size > 0: self._build_index()
173
174
175    def __delete__(self, instance):
176        self.close()
177        del self._index
178
179
180    def __enter__(self):
181        return self
182
183
184    def __exit__(self, atype, value, traceback):
185        self.close()
186
187
188    def __iter__(self):
189        yield from self._index['msg']
190
191
192    def __len__(self):
193        return self.messages
194
195
196    def __repr__(self):
197        strings = []
198        for k in self.__slots__:
199            if k.startswith('_'): continue
200            strings.append('%s = %s\n'%(k,eval('self.'+k)))
201        return ''.join(strings)
202
203
204    def __getitem__(self, key):
205        if isinstance(key,int):
206            if abs(key) >= len(self._index['msg']):
207                raise IndexError("index out of range")
208            else:
209                return self._index['msg'][key]
210        elif isinstance(key,str):
211            return self.select(shortName=key)
212        elif isinstance(key,slice):
213            return self._index['msg'][key]
214        else:
215            raise KeyError('Key must be an integer, slice, or GRIB2 variable shortName.')
216
217
218    def _build_index(self):
219        """Perform indexing of GRIB2 Messages."""
220        # Initialize index dictionary
221        if not self._hasindex:
222            self._index['sectionOffset'] = []
223            self._index['sectionSize'] = []
224            self._index['msgSize'] = []
225            self._index['msgNumber'] = []
226            self._index['msg'] = []
227            self._index['isSubmessage'] = []
228            self._hasindex = True
229
230        # Iterate
231        while True:
232            try:
233                # Initialize
234                bmapflag = None
235                pos = self._filehandle.tell()
236                section2 = b''
237                trailer = b''
238                _secpos = dict.fromkeys(range(8))
239                _secsize = dict.fromkeys(range(8))
240                _isSubmessage = False
241
242                # Ignore headers (usually text) that are not part of the GRIB2
243                # file.  For example, NAVGEM files have a http header at the
244                # beginning that needs to be ignored.
245
246                # Read a byte at a time until "GRIB" is found.  Using
247                # "wgrib2" on a NAVGEM file, the header was 421 bytes and
248                # decided to go to 2048 bytes to be safe. For normal GRIB2
249                # files this should be quick and break out of the first
250                # loop when "test_offset" is 0.
251                for test_offset in range(2048):
252                    self._filehandle.seek(pos + test_offset)
253                    header = struct.unpack(">I", self._filehandle.read(4))[0]
254                    if header.to_bytes(4, "big") == b"GRIB":
255                        pos = pos + test_offset
256                        break
257                else:
258                    # NOTE: Coming here means that no "GRIB" message identifier
259                    # was found in the previous 2048 bytes. So here we continue
260                    # the while True loop.
261                    continue
262
263                # Read the rest of Section 0 using struct.
264                _secpos[0] = self._filehandle.tell()-4
265                _secsize[0] = 16
266                secmsg = self._filehandle.read(12)
267                section0 = np.concatenate(([header],list(struct.unpack('>HBBQ',secmsg))),dtype=np.int64)
268
269                # Make sure message is GRIB2.
270                if section0[3] != 2:
271                    # Check for GRIB1 and ignore.
272                    if secmsg[3] == 1:
273                        warnings.warn("GRIB version 1 message detected.  Ignoring...")
274                        grib1_size = int.from_bytes(secmsg[:3],"big")
275                        self._filehandle.seek(self._filehandle.tell()+grib1_size-16)
276                        continue
277                    else:
278                        raise ValueError("Bad GRIB version number.")
279
280                # Read and unpack sections 1 through 8 which all follow a
281                # pattern of section size, number, and content.
282                while 1:
283                    # Read first 5 bytes of the section which contains the size
284                    # of the section (4 bytes) and the section number (1 byte).
285                    secmsg = self._filehandle.read(5)
286                    secsize, secnum = struct.unpack('>iB',secmsg)
287
288                    # Record the offset of the section number and "append" the
289                    # rest of the section to secmsg.
290                    _secpos[secnum] = self._filehandle.tell()-5
291                    _secsize[secnum] = secsize
292                    if secnum in {1,3,4,5}:
293                        secmsg += self._filehandle.read(secsize-5)
294                    grbpos = 0
295
296                    # Unpack section
297                    if secnum == 1:
298                        # Unpack Section 1
299                        section1, grbpos = g2clib.unpack1(secmsg)
300                    elif secnum == 2:
301                        # Unpack Section 2
302                        section2 = self._filehandle.read(secsize-5)
303                    elif secnum == 3:
304                        # Unpack Section 3
305                        gds, gdt, deflist, grbpos = g2clib.unpack3(secmsg)
306                        gds = gds.tolist()
307                        gdt = gdt.tolist()
308                        section3 = np.concatenate((gds,gdt))
309                        section3 = np.where(section3==4294967295,-1,section3)
310                    elif secnum == 4:
311                        # Unpack Section 4
312                        pdtnum, pdt, coordlist, numcoord, grbpos = g2clib.unpack4(secmsg)
313                        pdt = pdt.tolist()
314                        section4 = np.concatenate((np.array((numcoord,pdtnum)),pdt))
315                    elif secnum == 5:
316                        # Unpack Section 5
317                        drtn, drt, npts, self._pos = g2clib.unpack5(secmsg)
318                        section5 = np.concatenate((np.array((npts,drtn)),drt))
319                        section5 = np.where(section5==4294967295,-1,section5)
320                    elif secnum == 6:
321                        # Unpack Section 6 - Just the bitmap flag
322                        bmapflag = struct.unpack('>B',self._filehandle.read(1))[0]
323                        if bmapflag == 0:
324                            bmappos = self._filehandle.tell()-6
325                        elif bmapflag == 254:
326                            # Do this to keep the previous position value
327                            pass
328                        else:
329                            bmappos = None
330                        self._filehandle.seek(self._filehandle.tell()+secsize-6)
331                    elif secnum == 7:
332                        # Do not unpack section 7 here, but move the file pointer
333                        # to after section 7.
334                        self._filehandle.seek(self._filehandle.tell()+secsize-5)
335
336                        # Update the file index.
337                        self.messages += 1
338                        self._index['sectionOffset'].append(copy.deepcopy(_secpos))
339                        self._index['sectionSize'].append(copy.deepcopy(_secsize))
340                        self._index['msgSize'].append(section0[-1])
341                        self._index['msgNumber'].append(self.messages)
342                        self._index['isSubmessage'].append(_isSubmessage)
343
344                        # Create Grib2Message with data.
345                        msg = Grib2Message(section0,section1,section2,section3,section4,section5,bmapflag)
346                        msg._msgnum = self.messages-1
347                        msg._deflist = deflist
348                        msg._coordlist = coordlist
349                        if not self._nodata:
350                            msg._ondiskarray = Grib2MessageOnDiskArray((msg.ny,msg.nx), 2,
351                                                                TYPE_OF_VALUES_DTYPE[msg.typeOfValues],
352                                                                self._filehandle,
353                                                                msg, pos, _secpos[6], _secpos[7])
354                        self._index['msg'].append(msg)
355
356                        # If here, then we have moved through GRIB2 section 1-7.
357                        # Now we need to check the next 4 bytes after section 7.
358                        trailer = struct.unpack('>i',self._filehandle.read(4))[0]
359
360                        # If we reach the GRIB2 trailer string ('7777'), then we
361                        # can break begin processing the next GRIB2 message.  If
362                        # not, then we continue within the same iteration to
363                        # process a GRIB2 submessage.
364                        if trailer.to_bytes(4, "big") == b'7777':
365                            break
366                        else:
367                            # If here, trailer should be the size of the first
368                            # section of the next submessage, then the next byte
369                            # is the section number.  Check this value.
370                            nextsec = struct.unpack('>B',self._filehandle.read(1))[0]
371                            if nextsec not in {2,3,4}:
372                                raise ValueError("Bad GRIB2 message structure.")
373                            self._filehandle.seek(self._filehandle.tell()-5)
374                            _isSubmessage = True
375                            continue
376                    else:
377                        raise ValueError("Bad GRIB2 section number.")
378
379            except(struct.error):
380                if 'r' in self.mode:
381                    self._filehandle.seek(0)
382                break
383
384
385    @property
386    def levels(self):
387        """Provides a unique tuple of level strings."""
388        if self._hasindex and not self._nodata:
389            return tuple(sorted(set([msg.level for msg in self._index['msg']])))
390        else:
391            return None
392
393
394    @property
395    def variables(self):
396        """Provides a unique tuple of variable shortName strings."""
397        if self._hasindex and not self._nodata:
398            return tuple(sorted(set([msg.shortName for msg in self._index['msg']])))
399        else:
400            return None
401
402
403    def close(self):
404        """Close the file handle."""
405        if not self._filehandle.closed:
406            self.messages = 0
407            self.current_message = 0
408            self._filehandle.close()
409            self.closed = self._filehandle.closed
410
411
412    def read(self, size: Optional[int]=None):
413        """
414        Read size amount of GRIB2 messages from the current position.
415
416        If no argument is given, then size is None and all messages are returned
417        from the current position in the file. This read method follows the
418        behavior of Python's builtin open() function, but whereas that operates
419        on units of bytes, we operate on units of GRIB2 messages.
420
421        Parameters
422        ----------
423        size: default=None
424            The number of GRIB2 messages to read from the current position. If
425            no argument is give, the default value is None and remainder of
426            the file is read.
427
428        Returns
429        -------
430        read
431            ``Grib2Message`` object when size = 1 or a list of Grib2Messages
432            when size > 1.
433        """
434        if size is not None and size < 0:
435            size = None
436        if size is None or size > 1:
437            start = self.tell()
438            stop = self.messages if size is None else start+size
439            if size is None:
440                self.current_message = self.messages-1
441            else:
442                self.current_message += size
443            return self._index['msg'][slice(start,stop,1)]
444        elif size == 1:
445            self.current_message += 1
446            return self._index['msg'][self.current_message]
447        else:
448            None
449
450
451    def seek(self, pos: int):
452        """
453        Set the position within the file in units of GRIB2 messages.
454
455        Parameters
456        ----------
457        pos
458            The GRIB2 Message number to set the file pointer to.
459        """
460        if self._hasindex:
461            self._filehandle.seek(self._index['sectionOffset'][0][pos])
462            self.current_message = pos
463
464
465    def tell(self):
466        """Returns the position of the file in units of GRIB2 Messages."""
467        return self.current_message
468
469
470    def select(self, **kwargs):
471        """Select GRIB2 messages by `Grib2Message` attributes."""
472        # TODO: Added ability to process multiple values for each keyword (attribute)
473        idxs = []
474        nkeys = len(kwargs.keys())
475        for k,v in kwargs.items():
476            for m in self._index['msg']:
477                if hasattr(m,k) and getattr(m,k) == v: idxs.append(m._msgnum)
478        idxs = np.array(idxs,dtype=np.int32)
479        return [self._index['msg'][i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
480
481
482    def write(self, msg):
483        """
484        Writes GRIB2 message object to file.
485
486        Parameters
487        ----------
488        msg
489            GRIB2 message objects to write to file.
490        """
491        if isinstance(msg,list):
492            for m in msg:
493                self.write(m)
494            return
495
496        if issubclass(msg.__class__,_Grib2Message):
497            if hasattr(msg,'_msg'):
498                self._filehandle.write(msg._msg)
499            else:
500                if msg._signature != msg._generate_signature():
501                    msg.pack()
502                    self._filehandle.write(msg._msg)
503                else:
504                    if hasattr(msg._data,'filehandle'):
505                        msg._data.filehandle.seek(msg._data.offset)
506                        self._filehandle.write(msg._data.filehandle.read(msg.section0[-1]))
507                    else:
508                        msg.pack()
509                        self._filehandle.write(msg._msg)
510            self.flush()
511            self.size = os.path.getsize(self.name)
512            self._filehandle.seek(self.size-msg.section0[-1])
513            self._build_index()
514        else:
515            raise TypeError("msg must be a Grib2Message object.")
516        return
517
518
519    def flush(self):
520        """Flush the file object buffer."""
521        self._filehandle.flush()
522
523
524    def levels_by_var(self, name: str):
525        """
526        Return a list of level strings given a variable shortName.
527
528        Parameters
529        ----------
530        name
531            Grib2Message variable shortName
532
533        Returns
534        -------
535        levels_by_var
536            A list of unique level strings.
537        """
538        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
539
540
541    def vars_by_level(self, level: str):
542        """
543        Return a list of variable shortName strings given a level.
544
545        Parameters
546        ----------
547        level
548            Grib2Message variable level
549
550        Returns
551        -------
552        vars_by_level
553            A list of unique variable shortName strings.
554        """
555        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}')

1670def interpolate(a,
1671    method: Union[int, str],
1672    grid_def_in,
1673    grid_def_out,
1674    method_options=None,
1675    num_threads=1,
1676):
1677    """
1678    This is the module-level interpolation function.
1679
1680    This interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
1681    through grib2io's internal iplib Cython extension module.
1682
1683    Parameters
1684    ----------
1685    a : numpy.ndarray or tuple
1686        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
1687        performed.  If `a` is a `tuple`, then vector interpolation will be
1688        performed with the assumption that u = a[0] and v = a[1] and are both
1689        `numpy.ndarray`.
1690
1691        These data are expected to be in 2-dimensional form with shape (ny, nx)
1692        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
1693        spatial, temporal, or classification (i.e. ensemble members) dimension.
1694        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
1695        acceptable for input into the interpolation subroutines. If needed, these
1696        data will be converted to `np.float32`.
1697    method
1698        Interpolate method to use. This can either be an integer or string using
1699        the following mapping:
1700
1701        | Interpolate Scheme | Integer Value |
1702        | :---:              | :---:         |
1703        | 'bilinear'         | 0             |
1704        | 'bicubic'          | 1             |
1705        | 'neighbor'         | 2             |
1706        | 'budget'           | 3             |
1707        | 'spectral'         | 4             |
1708        | 'neighbor-budget'  | 6             |
1709
1710    grid_def_in : grib2io.Grib2GridDef
1711        Grib2GridDef object for the input grid.
1712    grid_def_out : grib2io.Grib2GridDef
1713        Grib2GridDef object for the output grid or station points.
1714    method_options : list of ints, optional
1715        Interpolation options. See the NCEPLIBS-ip documentation for
1716        more information on how these are used.
1717    num_threads : int, optional
1718        Number of OpenMP threads to use for interpolation. The default
1719        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
1720        was not built with OpenMP, then this keyword argument and value
1721        will have no impact.
1722
1723    Returns
1724    -------
1725    interpolate
1726        Returns a `numpy.ndarray` of dtype `np.float32` when scalar interpolation
1727        is performed or a `tuple` of `numpy.ndarray`s when vector interpolation is
1728        performed with the assumptions that 0-index is the interpolated u and
1729        1-index is the interpolated v.
1730    """
1731
1732    from . import iplib
1733
1734    prev_num_threads = 1
1735    try:
1736        prev_num_threads = iplib.get_openmp_threads()
1737        iplib.set_openmp_threads(num_threads)
1738    except(AttributeError):
1739        pass
1740
1741    if isinstance(method,int) and method not in _interp_schemes.values():
1742        raise ValueError('Invalid interpolation method.')
1743    elif isinstance(method,str):
1744        if method in _interp_schemes.keys():
1745            method = _interp_schemes[method]
1746        else:
1747            raise ValueError('Invalid interpolation method.')
1748
1749    if method_options is None:
1750        method_options = np.zeros((20),dtype=np.int32)
1751        if method in {3,6}:
1752            method_options[0:2] = -1
1753
1754    mi = grid_def_in.npoints
1755    mo = grid_def_out.npoints
1756
1757    # Adjust shape of input array(s)
1758    a, newshp = _adjust_array_shape_for_interp(a, grid_def_in, grid_def_out)
1759
1760    # Call interpolation subroutines according to type of a.
1761    if isinstance(a,np.ndarray):
1762        # Scalar
1763        km = a.shape[0]
1764        if np.any(np.isnan(a)):
1765            ibi = np.ones((km), dtype=np.int32)
1766            li = np.where(np.isnan(a),0,1).astype(np.uint8)
1767        else:
1768            ibi = np.zeros((km), dtype=np.int32)
1769            li = np.zeros(a.shape,dtype=np.uint8)
1770        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
1771            method,
1772            method_options,
1773            grid_def_in.gdtn,
1774            np.array(grid_def_in.gdt, dtype=np.int32),
1775            grid_def_out.gdtn,
1776            np.array(grid_def_out.gdt, dtype=np.int32),
1777            mi,
1778            mo,
1779            km,
1780            ibi,
1781            li,
1782            a,
1783        )
1784        out = np.where(lo==0,np.nan,go).reshape(newshp)
1785    elif isinstance(a,tuple):
1786        # Vector
1787        km = a[0].shape[0]
1788        if np.any(np.isnan(a)):
1789            ibi = np.ones((km), dtype=np.int32)
1790            li = np.where(np.isnan(a),0,1).astype(np.uint8)
1791        else:
1792            ibi = np.zeros((km), dtype=np.int32)
1793            li = np.zeros(a[0].shape,dtype=np.uint8)
1794        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
1795            method,
1796            method_options,
1797            grid_def_in.gdtn,
1798            np.array(grid_def_in.gdt, dtype=np.int32),
1799            grid_def_out.gdtn,
1800            np.array(grid_def_out.gdt, dtype=np.int32),
1801            mi,
1802            mo,
1803            km,
1804            ibi,
1805            li,
1806            a[0],
1807            a[1],
1808        )
1809        uo = np.where(lo==0,np.nan,uo).reshape(newshp)
1810        vo = np.where(lo==0,np.nan,vo).reshape(newshp)
1811        out = (uo,vo)
1812
1813    try:
1814        iplib.set_openmp_threads(prev_num_threads)
1815    except(AttributeError):
1816        pass
1817
1818    return out

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