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
 14try:
 15    from .. import iplib
 16except ImportError:
 17    pass
 18from .. import tables
 19from .. import templates
 20
 21
 22def decimal_to_scaled_int(
 23    value: Union[float, str, int],
 24    scale_factor: Optional[int] = None,
 25) -> Tuple[int, int]:
 26    """
 27    Convert a float-like value to a scaled integer using the minimal decimal scaling factor.
 28
 29    The input value is internally converted to a `Decimal` to ensure precise scaling.
 30
 31    Parameters
 32    ----------
 33    value : float, str, or int
 34        The numeric value to scale.
 35    scaled_value : int
 36        The integer result of scaling the original value by `10**scale_factor`.
 37
 38    Returns
 39    -------
 40    scale_factor : int
 41        The smallest power of 10 such that `value * 10**scale_factor` is an exact integer.
 42    scaled_value : int
 43        The integer result of scaling the original value by `10**scale_factor`.
 44    """
 45    dec_value = Decimal(str(value))  # Preserve exact decimal representation
 46
 47    with localcontext() as ctx:
 48        ctx.prec = 28
 49
 50        if scale_factor is not None:
 51            scaled = dec_value * (10**scale_factor)
 52            if scaled != scaled.to_integral_value():
 53                raise ValueError(f"Value {value} cannot be exactly scaled by 10^{scale_factor}")
 54            return scale_factor, int(scaled)
 55        else:
 56            scale_factor = 0
 57            while dec_value != dec_value.to_integral_value():
 58                dec_value *= 10
 59                scale_factor += 1
 60                if scale_factor > 20:
 61                    raise ValueError(f"Could not find exact scale factor for value {value} within bounds.")
 62            return scale_factor, int(dec_value)
 63
 64
 65def int2bin(i: int, nbits: int = 8, output: Union[Type[str], Type[List]] = str):
 66    """
 67    Convert integer to binary string or list
 68
 69    The struct module unpack using ">i" will unpack a 32-bit integer from a
 70    binary string.
 71
 72    Parameters
 73    ----------
 74    i
 75        Integer value to convert to binary representation.
 76    nbits : default=8
 77        Number of bits to return.  Valid values are 8 [DEFAULT], 16, 32, and
 78        64.
 79    output : default=str
 80        Return data as `str` [DEFAULT] or `list` (list of ints).
 81
 82    Returns
 83    -------
 84    int2bin
 85        `str` or `list` (list of ints) of binary representation of the integer
 86        value.
 87    """
 88    i = int(i) if not isinstance(i, int) else i
 89    assert nbits in [8, 16, 32, 64]
 90    bitstr = "{0:b}".format(i).zfill(nbits)
 91    if output is str:
 92        return bitstr
 93    elif output is list:
 94        return [int(b) for b in bitstr]
 95
 96
 97def ieee_float_to_int(f):
 98    """
 99    Convert an IEEE 754 32-bit float to a 32-bit integer.
100
101    Parameters
102    ----------
103    f : float
104        Floating-point value.
105
106    Returns
107    -------
108    ieee_float_to_int
109        `numpy.int32` representation of an IEEE 32-bit float.
110    """
111    i = struct.unpack(">i", struct.pack(">f", np.float32(f)))[0]
112    return np.int32(i)
113
114
115def ieee_int_to_float(i):
116    """
117    Convert a 32-bit integer to an IEEE 32-bit float.
118
119    Parameters
120    ----------
121    i : int
122        Integer value.
123
124    Returns
125    -------
126    ieee_int_to_float
127        `numpy.float32` representation of a 32-bit int.
128    """
129    f = struct.unpack(">f", struct.pack(">i", np.int32(i)))[0]
130    return np.float32(f)
131
132
133def get_leadtime(pdtn: int, pdt: ArrayLike) -> datetime.timedelta:
134    """
135    Compute lead time as a datetime.timedelta object.
136
137    Using information from GRIB2 Product Definition Template
138    Number, and Product Definition Template (Section 4).
139
140    Parameters
141    ----------
142    pdtn
143        GRIB2 Product Definition Template Number
144    pdt
145        Sequence containing GRIB2 Product Definition Template (Section 4).
146
147    Returns
148    -------
149    leadTime
150        datetime.timedelta object representing the lead time of the GRIB2 message.
151    """
152    lt = tables.get_value_from_table(pdt[templates.UnitOfForecastTime._key[pdtn]], "scale_time_seconds")
153    lt *= pdt[templates.ValueOfForecastTime._key[pdtn]]
154    return datetime.timedelta(seconds=int(lt))
155
156
157def get_duration(pdtn: int, pdt: ArrayLike) -> datetime.timedelta:
158    """
159    Compute a time duration as a datetime.timedelta.
160
161    Uses information from Product Definition Template Number, and Product
162    Definition Template (Section 4).
163
164    Parameters
165    ----------
166    pdtn
167        GRIB2 Product Definition Template Number
168    pdt
169        Sequence containing GRIB2 Product Definition Template (Section 4).
170
171    Returns
172    -------
173    get_duration
174        datetime.timedelta object representing the time duration of the GRIB2
175        message.
176    """
177    if pdtn in templates._timeinterval_pdtns:
178        ntime = pdt[templates.NumberOfTimeRanges._key[pdtn]]
179        duration_unit = tables.get_value_from_table(
180            pdt[templates.UnitOfTimeRangeOfStatisticalProcess._key[pdtn]],
181            "scale_time_seconds",
182        )
183        d = min(ntime, 1) * duration_unit * pdt[templates.TimeRangeOfStatisticalProcess._key[pdtn]]
184    else:
185        d = 0
186    return datetime.timedelta(seconds=int(d))
187
188
189def decode_wx_strings(lus: bytes) -> Dict[int, str]:
190    """
191    Decode GRIB2 Local Use Section to obtain NDFD/MDL Weather Strings.
192
193    The decode procedure is defined
194    [here](https://vlab.noaa.gov/web/mdl/nbm-gmos-grib2-wx-info).
195
196    Parameters
197    ----------
198    lus
199        GRIB2 Local Use Section containing NDFD weather strings.
200
201    Returns
202    -------
203    decode_wx_strings
204        Dict of NDFD/MDL weather strings. Keys are an integer value that
205        represent the sequential order of the key in the packed local use
206        section and the value is the weather key.
207    """
208    assert lus[0] == 1
209    # Unpack information related to the simple packing method
210    # the packed weather string data.
211    struct.unpack(">H", lus[1:3])[0]
212    struct.unpack(">i", lus[3:7])[0]
213    refvalue = struct.unpack(">i", lus[7:11])[0]
214    dsf = struct.unpack(">h", lus[11:13])[0]
215    nbits = lus[13]
216    datatype = lus[14]
217    if datatype == 0:  # Floating point
218        refvalue = np.float32(ieee_int_to_float(refvalue) * 10**-dsf)
219    elif datatype == 1:  # Integer
220        refvalue = np.int32(ieee_int_to_float(refvalue) * 10**-dsf)
221    # Upack each byte starting at byte 15 to end of the local use
222    # section, create a binary string and append to the full
223    # binary string.
224    b = ""
225    for i in range(15, len(lus)):
226        iword = struct.unpack(">B", lus[i : i + 1])[0]
227        b += bin(iword).split("b")[1].zfill(8)
228    # Iterate over the binary string (b). For each nbits
229    # chunk, convert to an integer, including the refvalue,
230    # and then convert the int to an ASCII character, then
231    # concatenate to wxstring.
232    wxstring = ""
233    for i in range(0, len(b), nbits):
234        wxstring += chr(int(b[i : i + nbits], 2) + refvalue)
235    # Return string as list, split by null character.
236    # return list(filter(None,wxstring.split('\0')))
237    return {n: k for n, k in enumerate(list(filter(None, wxstring.split("\0"))))}
238
239
240def get_wgrib2_prob_string(
241    probtype: int,
242    sfacl: int,
243    svall: int,
244    sfacu: int,
245    svalu: int,
246) -> str:
247    """
248    Return a wgrib2-styled string of probabilistic threshold information.
249
250    Logic from wgrib2 source,
251    [Prob.c](https://github.com/NOAA-EMC/NCEPLIBS-wgrib2/blob/develop/wgrib2/Prob.c),
252    is replicated here.
253
254    Parameters
255    ----------
256    probtype
257        Type of probability (Code Table 4.9).
258    sfacl
259        Scale factor of lower limit.
260    svall
261        Scaled value of lower limit.
262    sfacu
263        Scale factor of upper limit.
264    svalu
265        Scaled value of upper limit.
266
267    Returns
268    -------
269    get_wgrib2_prob_string
270        wgrib2-formatted string of probability threshold.
271    """
272    probstr = ""
273    if sfacl < 0:
274        sfacl = 0
275    if sfacu < 0:
276        sfacu = 0
277    lower = svall / (10**sfacl)
278    upper = svalu / (10**sfacu)
279    if probtype == 0:
280        probstr = "prob <%g" % (lower)
281    elif probtype == 1:
282        probstr = "prob >%g" % (upper)
283    elif probtype == 2:
284        if lower == upper:
285            probstr = "prob =%g" % (lower)
286        else:
287            probstr = "prob >=%g <%g" % (lower, upper)
288    elif probtype == 3:
289        probstr = "prob >%g" % (lower)
290    elif probtype == 4:
291        probstr = "prob <%g" % (upper)
292    else:
293        probstr = ""
294    return probstr
295
296
297def latlon_to_ij(
298    gdtn,
299    gdt,
300    lats,
301    lons,
302    missing_value=np.nan,
303):
304    """
305    Convert latitude/longitude coordinates to grid (i, j) indices using the
306    GRIB2 Grid Definition Section (GDS).
307
308    This function calls the grib2io iplib Cython extension module function,
309    `grib2io.iplib.latlon_to_ij`.
310
311    Parameters
312    ----------
313    gdtn : int
314        GRIB2 grid definition template number.
315    gdt : ndarray of int32
316        GRIB2 grid definition template values.
317    lats : numpy.ndarray or list
318        Latitude coordinates in degrees.
319    lons : numpy.ndarray or list
320        Longitude coordinates in degrees.
321    missing_value : float, optional
322        Missing value to represent when latitude/longitude coordinate is
323        outside the grid domain.
324
325    Returns
326    -------
327    xpts : ndarray of float32
328        Grid x-coordinates (i-indices) corresponding to the input
329        latitude/longitude points.
330    ypts : ndarray of float32
331        Grid y-coordinates (j-indices) corresponding to the input
332        latitude/longitude points.
333    """
334    # Check lats and lons
335    if isinstance(lats, list):
336        nlats = len(lats)
337    elif isinstance(lats, np.ndarray) and len(lats.shape) == 1:
338        nlats = lats.shape[0]
339    else:
340        raise ValueError("Latitudes must be a list or 1-D NumPy array.")
341    if isinstance(lons, list):
342        nlons = len(lons)
343    elif isinstance(lons, np.ndarray) and len(lons.shape) == 1:
344        nlons = lons.shape[0]
345    else:
346        raise ValueError("Longitudes must be a list or 1-D NumPy array.")
347    if nlats != nlons:
348        raise ValueError("Latitudes and longitudes same length.")
349    xpts, ypts = iplib.latlon_to_ij(
350        gdtn.astype(np.int32),
351        gdt.astype(np.int32),
352        np.array(lats, dtype=np.float64),
353        np.array(lons, dtype=np.float64),
354        missing_value,
355    )
356
357    return xpts.astype(np.float32), ypts.astype(np.float32)
358
359
360def compute_with_retries(obj, *, max_attempts: int = 6, base_sleep: float = 2.0):
361    """
362    Compute a Dask-backed object (Xarray Dataset or DataArray) with retries for transient errors.
363
364    Transient errors typically include network timeouts, connection resets, and
365    cloud-specific errors (e.g. IcechunkError, S3 SlowDown 429).
366
367    Parameters
368    ----------
369    obj : xarray.Dataset, xarray.DataArray, or dask.delayed.Delayed
370        The object to compute.
371    max_attempts : int, optional
372        Maximum number of attempts. Defaults to 6.
373    base_sleep : float, optional
374        Base sleep time in seconds for exponential backoff. Defaults to 2.0.
375
376    Returns
377    -------
378    The computed result (typically a NumPy-backed Xarray object or a scalar).
379
380    Raises
381    ------
382    Exception
383        The last exception encountered if all attempts fail, or any non-transient
384        exception.
385    """
386    import time
387
388    for attempt in range(1, max_attempts + 1):
389        try:
390            return obj.compute()
391        except Exception as exc:
392            msg = str(exc).lower()
393            # Basic transient error detection across typical remote backends
394            exc_name = type(exc).__name__.lower()
395            transient = (
396                "icechunkerror" in exc_name
397                or "timeout" in msg
398                or "timed out" in msg
399                or "connect" in msg
400                or "slowdown" in msg
401                or "429" in msg
402                or "503" in msg
403                or "eof occurred in violation of protocol" in msg
404                or "connection reset" in msg
405            )
406            if (not transient) or attempt == max_attempts:
407                raise
408            sleep_s = base_sleep**attempt
409            print(f"Transient read error ({type(exc).__name__}) on attempt {attempt}/{max_attempts}; retrying in {sleep_s}s...")
410            time.sleep(sleep_s)
411
412
413def percentile_string(pct):
414    """
415    Return a percentile string with the proper English ordinal suffix.
416
417    Parameters
418    ----------
419    pct : int
420        Percentile value in the range [0, 100].
421
422    Returns
423    -------
424    str
425        Percentile string, e.g., ``"1st percentile"``,
426        ``"21st percentile"``, or ``"90th percentile"``.
427
428    Raises
429    ------
430    ValueError
431        If `pct` is not in the range [0, 100].
432
433    Examples
434    --------
435    >>> percentile_string(1)
436    '1st percentile'
437    >>> percentile_string(21)
438    '21st percentile'
439    >>> percentile_string(90)
440    '90th percentile'
441    """
442    if not (0 <= pct <= 100):
443        raise ValueError("percentile must be between 0 and 100")
444
445    if 11 <= (pct % 100) <= 13:
446        suffix = "th"
447    else:
448        suffix = {1: "st", 2: "nd", 3: "rd"}.get(pct % 10, "th")
449
450    return f"{pct}{suffix} percentile"