grib2io.utils

  1"""
  2Collection of utility functions to assist in the encoding and decoding
  3of GRIB2 Messages.
  4"""
  5
  6import datetime
  7import struct
  8from decimal import Decimal, localcontext
  9from typing import Dict, List, Optional, Tuple, Type, Union
 10
 11import numpy as np
 12from numpy.typing import ArrayLike
 13
 14from .. import iplib
 15from .. import tables
 16from .. import templates
 17
 18
 19def decimal_to_scaled_int(
 20    value: Union[float, str, int],
 21    scale_factor: Optional[int] = None,
 22) -> Tuple[int, int]:
 23    """
 24    Convert a float-like value to a scaled integer using the minimal decimal scaling factor.
 25
 26    The input value is internally converted to a `Decimal` to ensure precise scaling.
 27
 28    Parameters
 29    ----------
 30    value : float, str, or int
 31        The numeric value to scale.
 32    scaled_value : int
 33        The integer result of scaling the original value by `10**scale_factor`.
 34
 35    Returns
 36    -------
 37    scale_factor : int
 38        The smallest power of 10 such that `value * 10**scale_factor` is an exact integer.
 39    scaled_value : int
 40        The integer result of scaling the original value by `10**scale_factor`.
 41    """
 42    dec_value = Decimal(str(value))  # Preserve exact decimal representation
 43
 44    with localcontext() as ctx:
 45        ctx.prec = 28
 46
 47        if scale_factor is not None:
 48            scaled = dec_value * (10 ** scale_factor)
 49            if scaled != scaled.to_integral_value():
 50                raise ValueError(
 51                    f"Value {value} cannot be exactly scaled by 10^{scale_factor}"
 52                )
 53            return scale_factor, int(scaled)
 54        else:
 55            scale_factor = 0
 56            while dec_value != dec_value.to_integral_value():
 57                dec_value *= 10
 58                scale_factor += 1
 59                if scale_factor > 20:
 60                    raise ValueError(
 61                        f"Could not find exact scale factor for value {value} within bounds."
 62                    )
 63            return scale_factor, int(dec_value)
 64
 65
 66def int2bin(i: int, nbits: int=8, output: Union[Type[str], Type[List]]=str):
 67    """
 68    Convert integer to binary string or list
 69
 70    The struct module unpack using ">i" will unpack a 32-bit integer from a
 71    binary string.
 72
 73    Parameters
 74    ----------
 75    i
 76        Integer value to convert to binary representation.
 77    nbits : default=8
 78        Number of bits to return.  Valid values are 8 [DEFAULT], 16, 32, and
 79        64.
 80    output : default=str
 81        Return data as `str` [DEFAULT] or `list` (list of ints).
 82
 83    Returns
 84    -------
 85    int2bin
 86        `str` or `list` (list of ints) of binary representation of the integer
 87        value.
 88    """
 89    i = int(i) if not isinstance(i,int) else i
 90    assert nbits in [8,16,32,64]
 91    bitstr = "{0:b}".format(i).zfill(nbits)
 92    if output is str:
 93        return bitstr
 94    elif output is list:
 95        return [int(b) for b in bitstr]
 96
 97
 98def ieee_float_to_int(f):
 99    """
100    Convert an IEEE 754 32-bit float to a 32-bit integer.
101
102    Parameters
103    ----------
104    f : float
105        Floating-point value.
106
107    Returns
108    -------
109    ieee_float_to_int
110        `numpy.int32` representation of an IEEE 32-bit float.
111    """
112    i = struct.unpack('>i',struct.pack('>f',np.float32(f)))[0]
113    return np.int32(i)
114
115
116def ieee_int_to_float(i):
117    """
118    Convert a 32-bit integer to an IEEE 32-bit float.
119
120    Parameters
121    ----------
122    i : int
123        Integer value.
124
125    Returns
126    -------
127    ieee_int_to_float
128        `numpy.float32` representation of a 32-bit int.
129    """
130    f = struct.unpack('>f',struct.pack('>i',np.int32(i)))[0]
131    return np.float32(f)
132
133
134def get_leadtime(pdtn: int, pdt: ArrayLike) -> datetime.timedelta:
135    """
136    Compute lead time as a datetime.timedelta object.
137
138    Using information from GRIB2 Product Definition Template
139    Number, and Product Definition Template (Section 4).
140
141    Parameters
142    ----------
143    pdtn
144        GRIB2 Product Definition Template Number
145    pdt
146        Sequence containing GRIB2 Product Definition Template (Section 4).
147
148    Returns
149    -------
150    leadTime
151        datetime.timedelta object representing the lead time of the GRIB2 message.
152    """
153    lt = tables.get_value_from_table(pdt[templates.UnitOfForecastTime._key[pdtn]], 'scale_time_seconds')
154    lt *= pdt[templates.ValueOfForecastTime._key[pdtn]]
155    return datetime.timedelta(seconds=int(lt))
156
157
158def get_duration(pdtn: int, pdt: ArrayLike) -> datetime.timedelta:
159    """
160    Compute a time duration as a datetime.timedelta.
161
162    Uses information from Product Definition Template Number, and Product
163    Definition Template (Section 4).
164
165    Parameters
166    ----------
167    pdtn
168        GRIB2 Product Definition Template Number
169    pdt
170        Sequence containing GRIB2 Product Definition Template (Section 4).
171
172    Returns
173    -------
174    get_duration
175        datetime.timedelta object representing the time duration of the GRIB2
176        message.
177    """
178    if pdtn in templates._timeinterval_pdtns:
179        ntime = pdt[templates.NumberOfTimeRanges._key[pdtn]]
180        duration_unit = tables.get_value_from_table(
181            pdt[templates.UnitOfTimeRangeOfStatisticalProcess._key[pdtn]],
182            'scale_time_seconds')
183        d = ntime * duration_unit * pdt[
184            templates.TimeRangeOfStatisticalProcess._key[pdtn]]
185    else:
186        d = 0
187    return datetime.timedelta(seconds=int(d))
188
189
190def decode_wx_strings(lus: bytes) -> Dict[int, str]:
191    """
192    Decode GRIB2 Local Use Section to obtain NDFD/MDL Weather Strings.
193
194    The decode procedure is defined
195    [here](https://vlab.noaa.gov/web/mdl/nbm-gmos-grib2-wx-info).
196
197    Parameters
198    ----------
199    lus
200        GRIB2 Local Use Section containing NDFD weather strings.
201
202    Returns
203    -------
204    decode_wx_strings
205        Dict of NDFD/MDL weather strings. Keys are an integer value that
206        represent the sequential order of the key in the packed local use
207        section and the value is the weather key.
208    """
209    assert lus[0] == 1
210    # Unpack information related to the simple packing method
211    # the packed weather string data.
212    ngroups = struct.unpack('>H',lus[1:3])[0]
213    nvalues = struct.unpack('>i',lus[3:7])[0]
214    refvalue = struct.unpack('>i',lus[7:11])[0]
215    dsf = struct.unpack('>h',lus[11:13])[0]
216    nbits = lus[13]
217    datatype = lus[14]
218    if datatype == 0: # Floating point
219        refvalue = np.float32(ieee_int_to_float(refvalue)*10**-dsf)
220    elif datatype == 1: # Integer
221        refvalue = np.int32(ieee_int_to_float(refvalue)*10**-dsf)
222    # Upack each byte starting at byte 15 to end of the local use
223    # section, create a binary string and append to the full
224    # binary string.
225    b = ''
226    for i in range(15,len(lus)):
227        iword = struct.unpack('>B',lus[i:i+1])[0]
228        b += bin(iword).split('b')[1].zfill(8)
229    # Iterate over the binary string (b). For each nbits
230    # chunk, convert to an integer, including the refvalue,
231    # and then convert the int to an ASCII character, then
232    # concatenate to wxstring.
233    wxstring = ''
234    for i in range(0,len(b),nbits):
235        wxstring += chr(int(b[i:i+nbits],2)+refvalue)
236    # Return string as list, split by null character.
237    #return list(filter(None,wxstring.split('\0')))
238    return {n:k for n,k in enumerate(list(filter(None,wxstring.split('\0'))))}
239
240
241def get_wgrib2_prob_string(
242    probtype: int,
243    sfacl: int,
244    svall: int,
245    sfacu: int,
246    svalu: int,
247) -> str:
248    """
249    Return a wgrib2-styled string of probabilistic threshold information.
250
251    Logic from wgrib2 source,
252    [Prob.c](https://github.com/NOAA-EMC/NCEPLIBS-wgrib2/blob/develop/wgrib2/Prob.c),
253    is replicated here.
254
255    Parameters
256    ----------
257    probtype
258        Type of probability (Code Table 4.9).
259    sfacl
260        Scale factor of lower limit.
261    svall
262        Scaled value of lower limit.
263    sfacu
264        Scale factor of upper limit.
265    svalu
266        Scaled value of upper limit.
267
268    Returns
269    -------
270    get_wgrib2_prob_string
271        wgrib2-formatted string of probability threshold.
272    """
273    probstr = ''
274    if sfacl < 0: sfacl = 0
275    if sfacu < 0: sfacu = 0
276    lower = svall/(10**sfacl)
277    upper = svalu/(10**sfacu)
278    if probtype == 0:
279        probstr = 'prob <%g' % (lower)
280    elif probtype == 1:
281        probstr = 'prob >%g' % (upper)
282    elif probtype == 2:
283        if lower == upper:
284            probstr = 'prob =%g' % (lower)
285        else:
286            probstr = 'prob >=%g <%g' % (lower,upper)
287    elif probtype == 3:
288        probstr = 'prob >%g' % (lower)
289    elif probtype == 4:
290        probstr = 'prob <%g' % (upper)
291    else:
292        probstr = ''
293    return probstr
294
295
296def latlon_to_ij(
297    gdtn,
298    gdt,
299    lats,
300    lons,
301    missing_value = np.nan,
302):
303    """
304    Convert latitude/longitude coordinates to grid (i, j) indices using the
305    GRIB2 Grid Definition Section (GDS).
306
307    This function calls the grib2io iplib Cython extension module function,
308    `grib2io.iplib.latlon_to_ij`.
309
310    Parameters
311    ----------
312    gdtn : int
313        GRIB2 grid definition template number.
314    gdt : ndarray of int32
315        GRIB2 grid definition template values.
316    lats : numpy.ndarray or list
317        Latitude coordinates in degrees.
318    lons : numpy.ndarray or list
319        Longitude coordinates in degrees.
320    missing_value : float, optional
321        Missing value to represent when latitude/longitude coordinate is
322        outside the grid domain.
323
324    Returns
325    -------
326    xpts : ndarray of float32
327        Grid x-coordinates (i-indices) corresponding to the input
328        latitude/longitude points.
329    ypts : ndarray of float32
330        Grid y-coordinates (j-indices) corresponding to the input
331        latitude/longitude points.
332    """
333    # Check lats and lons
334    if isinstance(lats,list):
335        nlats = len(lats)
336    elif isinstance(lats,np.ndarray) and len(lats.shape) == 1:
337        nlats = lats.shape[0]
338    else:
339        raise ValueError("Latitudes must be a list or 1-D NumPy array.")
340    if isinstance(lons,list):
341        nlons = len(lons)
342    elif isinstance(lons,np.ndarray) and len(lons.shape) == 1:
343        nlons = lons.shape[0]
344    else:
345        raise ValueError("Longitudes must be a list or 1-D NumPy array.")
346    if nlats != nlons:
347        raise ValueError("Latitudes and longitudes same length.")
348    return iplib.latlon_to_ij(
349        gdtn.astype(np.int32),
350        gdt.astype(np.int32),
351        np.array(lats, dtype=np.float32),
352        np.array(lons, dtype=np.float32),
353        missing_value,
354    )