grib2io.tables

  1"""Functions for retrieving data from NCEP GRIB2 Tables."""
  2
  3from functools import lru_cache
  4from typing import Optional, Union, List
  5from numpy.typing import ArrayLike
  6import itertools
  7import importlib
  8import numpy as np
  9
 10from .section0 import *
 11from .section1 import *
 12from .section3 import *
 13from .section4 import *
 14from .section5 import *
 15from .section6 import *
 16from .originating_centers import *
 17from .ndfd_additionals import *
 18from .cf import *
 19
 20_varinfo_tables_datastore = {}
 21
 22# Note 209 is MRMS
 23GRIB2_DISCIPLINES = [0, 1, 2, 3, 4, 10, 20, 209]
 24
 25AEROSOL_PDTNS = [46, 48]  # 47, 49, 80, 81, 82, 83, 84, 85] <- these don't seem to be working
 26AEROSOL_PARAMS = list(itertools.chain(range(0,19),range(50,82),range(100,113),range(192,197)))
 27
 28def _load_varinfo_tables(modname: str):
 29    """
 30    Load variable information tables from sub modules into the local
 31    varinfo table datastore (_varinfo_tables_datastore).
 32
 33    Parameters
 34    ----------
 35    modname
 36        Module name to extract variable info tables from.
 37    """
 38    module = importlib.import_module(modname, package=__name__)
 39    names = getattr(module, '__all__', [name for name in dir(module) if name.startswith('table_')])
 40    _varinfo_tables_datastore.update({name: getattr(module, name) for name in names})
 41
 42
 43def get_table(table: str, expand: bool=False) -> dict:
 44    """
 45    Return GRIB2 code table as a dictionary.
 46
 47    Parameters
 48    ----------
 49    table
 50        Code table number (e.g. '1.0').
 51
 52        NOTE: Code table '4.1' requires a 3rd value representing the product
 53        discipline (e.g. '4.1.0').
 54    expand
 55        If `True`, expand output dictionary wherever keys are a range.
 56
 57    Returns
 58    -------
 59    get_table
 60        GRIB2 code table as a dictionary.
 61    """
 62    if len(table) == 3 and table == '4.1':
 63        raise Exception('GRIB2 Code Table 4.1 requires a 3rd value representing the discipline.')
 64    if len(table) == 3 and table.startswith('4.2'):
 65        raise Exception('Use function get_varinfo_from_table() for GRIB2 Code Table 4.2')
 66    try:
 67        tbl = globals()['table_'+table.replace('.','_')]
 68        if expand:
 69            _tbl = {}
 70            for k,v in tbl.items():
 71                if '-' in k:
 72                    irng = [int(i) for i in k.split('-')]
 73                    for i in range(irng[0],irng[1]+1):
 74                        _tbl[str(i)] = v
 75                else:
 76                    _tbl[k] = v
 77            tbl = _tbl
 78        return tbl
 79    except(KeyError):
 80        return {}
 81
 82
 83def get_value_from_table(
 84    value: Union[int, str],
 85    table: str,
 86    expand: bool = False,
 87) -> Optional[Union[float, int]]:
 88    """
 89    Return the definition given a GRIB2 code table.
 90
 91    Parameters
 92    ----------
 93    value
 94        Code table value.
 95    table
 96        Code table number.
 97    expand
 98        If `True`, expand output dictionary where keys are a range.
 99
100    Returns
101    -------
102    get_value_from_table
103        Table value or `None` if not found.
104    """
105    try:
106        tbl = get_table(table,expand=expand)
107        value = str(value)
108        return tbl[value]
109    except(KeyError):
110        for k in tbl.keys():
111            if '-' in k:
112                bounds = k.split('-')
113                if value >= bounds[0] and value <= bounds[1]:
114                    return tbl[k]
115        return None
116
117
118def get_varinfo_from_table(
119    discipline: Union[int, str],
120    parmcat: Union[int, str],
121    parmnum: Union[int, str],
122    isNDFD: bool = False,
123):
124    """
125    Return the GRIB2 variable information.
126
127    NOTE: This functions allows for all arguments to be converted to a string
128    type if arguments are integer.
129
130    Parameters
131    ----------
132    discipline
133        Discipline code value of a GRIB2 message.
134    parmcat
135        Parameter Category value of a GRIB2 message.
136    parmnum
137        Parameter Number value of a GRIB2 message.
138    isNDFD: optional
139        If `True`, then attempt to get variable information from the
140        supplemental NDFD tables, otherwise fallback to default tables.
141
142    Returns
143    -------
144    full_name
145        Full name of the GRIB2 variable. "Unknown" if variable is not found.
146    units
147        Units of the GRIB2 variable. "Unknown" if variable is not found.
148    shortName
149        Abbreviated name of the GRIB2 variable. "Unknown" if variable is not
150        found.
151    """
152    template = f'table_4_2_{discipline}_{parmcat}'
153    tbls = []
154    if isNDFD:
155        tbls.append(template+'_ndfd')
156    tbls.append(template)
157    for tbl in tbls:
158        vartbl = None
159        modname = f'.section4_discipline{discipline}'
160        if tbl not in _varinfo_tables_datastore.keys():
161            _load_varinfo_tables(modname)
162        try:
163            vartbl = _varinfo_tables_datastore[tbl][str(parmnum)]
164        except(KeyError):
165            pass
166        if vartbl is not None:
167            break
168    if vartbl is None:
169        vartbl = ['Unknown','Unknown','Unknown']
170    return vartbl
171
172
173@lru_cache(maxsize=None)
174def get_shortnames(
175    discipline: Optional[Union[int, str]] = None,
176    parmcat: Optional[Union[int, str]] = None,
177    parmnum: Optional[Union[int, str]] = None,
178    isNDFD: bool = False,
179) -> List[str]:
180    """
181    Return a list of variable shortNames.
182
183    If all 3 args are None, then shortNames from all disciplines, parameter
184    categories, and numbers will be returned.
185
186    Parameters
187    ----------
188    discipline
189        GRIB2 discipline code value.
190    parmcat
191        GRIB2 parameter category value.
192    parmnum
193        Parameter Number value of a GRIB2 message.
194    isNDFD: optional
195        If `True`, signals function to try to get variable information from the
196        supplemental NDFD tables.
197
198    Returns
199    -------
200    get_shortnames
201        list of GRIB2 shortNames.
202    """
203    shortnames = list()
204    if discipline is None:
205        discipline = GRIB2_DISCIPLINES
206    else:
207        discipline = [discipline]
208    if parmcat is None:
209        parmcat = list()
210        for d in discipline:
211            parmcat += list(get_table(f'4.1.{d}').keys())
212    else:
213        parmcat = [parmcat]
214    if parmnum is None:
215        parmnum = list(range(256))
216    else:
217        parmnum = [parmnum]
218    for d in discipline:
219
220        for pc in parmcat:
221            for pn in parmnum:
222                shortnames.append(get_varinfo_from_table(d,pc,pn,isNDFD)[2])
223
224    shortnames = sorted(set(shortnames))
225    try:
226        shortnames.remove('unknown')
227        shortnames.remove('Unknown')
228    except(ValueError):
229        pass
230    return shortnames
231
232
233@lru_cache(maxsize=None)
234def get_metadata_from_shortname(shortname: str):
235    """
236    Provide GRIB2 variable metadata attributes given a GRIB2 shortName.
237
238    Parameters
239    ----------
240    shortname
241        GRIB2 variable shortName.
242
243    Returns
244    -------
245    get_metadata_from_shortname
246        list of dictionary items where each dictionary item contains the
247        variable metadata key:value pairs.
248
249        NOTE: Some variable shortNames will exist in multiple parameter
250        category/number tables according to the GRIB2 discipline.
251    """
252    metadata = []
253    for d in GRIB2_DISCIPLINES:
254        parmcat = list(get_table(f'4.1.{d}').keys())
255        for pc in parmcat:
256            for pn in range(256):
257                varinfo = get_varinfo_from_table(d,pc,pn,False)
258                if shortname == varinfo[2]:
259                    metadata.append(dict(discipline=d,parameterCategory=pc,parameterNumber=pn,
260                                         fullName=varinfo[0],units=varinfo[1]))
261    return metadata
262
263
264def get_wgrib2_level_string(pdtn: int, pdt: ArrayLike) -> str:
265    """
266    Return a string that describes the level or layer of the GRIB2 message.
267
268    The format and language of the string is an exact replica of how wgrib2
269    produces the level/layer string in its inventory output.
270
271    Contents of wgrib2 source,
272    [Level.c](https://github.com/NOAA-EMC/NCEPLIBS-wgrib2/blob/develop/wgrib2/Level.c),
273    were converted into a Python dictionary and stored in grib2io as table
274    'wgrib2_level_string'.
275
276    Parameters
277    ----------
278    pdtn
279        GRIB2 Product Definition Template Number
280    pdt
281        Sequence containing GRIB2 Product Definition Template (Section 4).
282
283    Returns
284    -------
285    get_wgrib2_level_string
286        wgrib2-formatted level/layer string.
287    """
288    lvlstr = ''
289    if pdtn == 32:
290        return 'no_level'
291    elif pdtn == 48:
292        idxs = slice(20,26)
293    else:
294        idxs = slice(9,15)
295    type1, sfac1, sval1, type2, sfac2, sval2 = map(int,pdt[idxs])
296    val1 = sval1/10**sfac1
297    if type1 in [100,108]: val1 *= 0.01
298    if type2 != 255:
299        # Layer
300        #assert type2 == type1, "Surface types are not equal: %g - %g" % (type1,type2)
301        val2 = sval2/10**sfac2
302        if type2 in [100,108]: val2 *= 0.01
303        lvlstr = get_value_from_table(type1,table='wgrib2_level_string')[1]
304        vals = (val1,val2)
305    else:
306        # Level
307        lvlstr = get_value_from_table(type1,table='wgrib2_level_string')[0]
308        vals = (val1)
309    if '%g' in lvlstr: lvlstr %= vals
310    return lvlstr
311
312
313def _build_aerosol_shortname(obj) -> str:
314    """
315    """
316    _OPTICAL_WAVELENGTH_MAPPING = get_table('aerosol_optical_wavelength')
317    _LEVEL_MAPPING = get_table('aerosol_level')
318    _PARAMETER_MAPPING = get_table('aerosol_parameter')
319    _AERO_TYPE_MAPPING = get_table('aerosol_type')
320
321    # Build shortname from aerosol components
322    parts = []
323
324    # Get aerosol type
325    aero_type = str(obj.typeOfAerosol.value) if obj.typeOfAerosol is not None else ""
326
327    # Add size information if applicable
328    aero_size = ""
329    if hasattr(obj, 'scaledValueOfFirstSize'):
330        if float(obj.scaledValueOfFirstSize) > 0:
331            first_size = float(obj.scaledValueOfFirstSize)
332
333            # Map common PM sizes
334            size_map = {1: 'pm1', 25: 'pm25', 10: 'pm10', 20: 'pm20'}
335            aero_size = size_map.get(first_size, f"pm{int(first_size)}")
336
337            # Check for size intervals
338            if (hasattr(obj, 'scaledValueOfSecondSize') and
339                obj.scaledValueOfSecondSize is not None and
340                hasattr(obj, 'typeOfIntervalForAerosolSize') and
341                obj.typeOfIntervalForAerosolSize.value == 6):
342
343                second_size = float(obj.scaledValueOfSecondSize)
344                if second_size > 0:
345                    if (first_size == 2.5 and second_size == 10):
346                        aero_size = 'PM25to10'
347                    elif (first_size == 10 and second_size == 20):
348                        aero_size = 'PM10to20'
349                    else:
350                        aero_size = f"PM{int(first_size)}to{int(second_size)}"
351
352    # Add optical and wavelength information
353    var_wavelength = ''
354    if (hasattr(obj, 'parameterNumber') and
355        hasattr(obj, 'scaledValueOfFirstWavelength') and
356        hasattr(obj, 'scaledValueOfSecondWavelength')):
357
358        optical_type = str(obj.parameterNumber)
359        if obj.scaledValueOfFirstWavelength > 0:
360            first_wl = obj.scaledValueOfFirstWavelength
361            second_wl = obj.scaledValueOfSecondWavelength
362
363            # Special case for AE between 440-870nm
364            if optical_type == '111' and first_wl == 440 and second_wl == 870:
365                key = (optical_type, '440TO870')
366            else:
367                # Find matching wavelength band
368                for wl_key, wl_info in _OPTICAL_WAVELENGTH_MAPPING.items():
369                    if (wl_key[0] == optical_type and  # Check optical_type first
370                        int(wl_key[1]) == first_wl and
371                        (second_wl is None or int(wl_key[2]) == second_wl)):
372                        key = wl_key
373                        break
374                else:
375                    # If no match found, use raw values
376                    key = (optical_type, str(first_wl),
377                          str(second_wl) if second_wl is not None else '')
378
379            # FIX THIS...
380            if key in _OPTICAL_WAVELENGTH_MAPPING.keys():
381                var_wavelength = _OPTICAL_WAVELENGTH_MAPPING[key]
382
383    # Add level information
384    level_str = ''
385    if hasattr(obj, 'typeOfFirstFixedSurface'):
386        first_level = str(obj.typeOfFirstFixedSurface.value)
387        first_value = str(obj.scaledValueOfFirstFixedSurface) if obj.scaledValueOfFirstFixedSurface > 0 else ''
388        if first_level in _LEVEL_MAPPING:
389            level_str = f"{_LEVEL_MAPPING[first_level]}{first_value}"
390
391    # Get parameter type
392    param = ''
393    if hasattr(obj, 'parameterNumber'):
394        param_num = str(obj.parameterNumber)
395        if param_num in _PARAMETER_MAPPING:
396            param = _PARAMETER_MAPPING[param_num]
397
398    # Build the final shortname
399    if var_wavelength and aero_type in _AERO_TYPE_MAPPING:
400        shortname = f"{_AERO_TYPE_MAPPING[aero_type]}{var_wavelength}"
401    elif aero_type in _AERO_TYPE_MAPPING:
402        parts = []
403        if level_str:
404            parts.append(level_str)
405        parts.append(_AERO_TYPE_MAPPING[aero_type])
406        if aero_size:
407            parts.append(aero_size)
408        if param:
409            parts.append(param)
410        shortname = '_'.join(parts) if len(parts) > 1 else parts[0]
411    else:
412        return get_varinfo_from_table(obj.section0[2], *obj.section4[2:4], isNDFD=obj._isNDFD)[2]
413
414    return shortname
415
416
417@lru_cache(maxsize=None)
418def get_table_names(var_tables: bool=False) -> tuple:
419    """
420    Return the names of available GRIB2 tables.
421
422    The table dictionaries as they exist in grib2io are prefixed
423    with "table_". When accessing a table through the available
424    functions, the prefix is not necessary.
425
426    Parameters
427    ----------
428    var_tables : bool, optional
429        If True, include names of variable tables.
430        If False (default), include only non-variable tables.
431
432    Returns
433    -------
434    tuple of str
435        A tuple sorted names of available tables.
436    """
437    tables = [
438        name.replace("table_","") for name, val in globals().items()
439        if isinstance(val, dict) and name.startswith("table_")
440    ]
441
442    tables = [
443        t.replace("_",".") if t.startswith(tuple("0123456789")) else t for t in tables
444    ]
445
446    if var_tables:
447        for d in GRIB2_DISCIPLINES:
448            modname = f'.section4_discipline{d}'
449            _load_varinfo_tables(modname)
450        vt = [t.replace("table_","").replace("_",".") for t in _varinfo_tables_datastore.keys()]
451        tables.extend(vt)
452
453    return tuple(sorted(tables))