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            vartbl = ['Unknown','Unknown','Unknown']
166        if vartbl is not None:
167            break
168    return vartbl
169
170
171@lru_cache(maxsize=None)
172def get_shortnames(
173    discipline: Optional[Union[int, str]] = None,
174    parmcat: Optional[Union[int, str]] = None,
175    parmnum: Optional[Union[int, str]] = None,
176    isNDFD: bool = False,
177) -> List[str]:
178    """
179    Return a list of variable shortNames.
180
181    If all 3 args are None, then shortNames from all disciplines, parameter
182    categories, and numbers will be returned.
183
184    Parameters
185    ----------
186    discipline
187        GRIB2 discipline code value.
188    parmcat
189        GRIB2 parameter category value.
190    parmnum
191        Parameter Number value of a GRIB2 message.
192    isNDFD: optional
193        If `True`, signals function to try to get variable information from the
194        supplemental NDFD tables.
195
196    Returns
197    -------
198    get_shortnames
199        list of GRIB2 shortNames.
200    """
201    shortnames = list()
202    if discipline is None:
203        discipline = GRIB2_DISCIPLINES
204    else:
205        discipline = [discipline]
206    if parmcat is None:
207        parmcat = list()
208        for d in discipline:
209            parmcat += list(get_table(f'4.1.{d}').keys())
210    else:
211        parmcat = [parmcat]
212    if parmnum is None:
213        parmnum = list(range(256))
214    else:
215        parmnum = [parmnum]
216    for d in discipline:
217
218        for pc in parmcat:
219            for pn in parmnum:
220                shortnames.append(get_varinfo_from_table(d,pc,pn,isNDFD)[2])
221
222    shortnames = sorted(set(shortnames))
223    try:
224        shortnames.remove('unknown')
225        shortnames.remove('Unknown')
226    except(ValueError):
227        pass
228    return shortnames
229
230
231@lru_cache(maxsize=None)
232def get_metadata_from_shortname(shortname: str):
233    """
234    Provide GRIB2 variable metadata attributes given a GRIB2 shortName.
235
236    Parameters
237    ----------
238    shortname
239        GRIB2 variable shortName.
240
241    Returns
242    -------
243    get_metadata_from_shortname
244        list of dictionary items where each dictionary item contains the
245        variable metadata key:value pairs.
246
247        NOTE: Some variable shortNames will exist in multiple parameter
248        category/number tables according to the GRIB2 discipline.
249    """
250    metadata = []
251    for d in GRIB2_DISCIPLINES:
252        parmcat = list(get_table(f'4.1.{d}').keys())
253        for pc in parmcat:
254            for pn in range(256):
255                varinfo = get_varinfo_from_table(d,pc,pn,False)
256                if shortname == varinfo[2]:
257                    metadata.append(dict(discipline=d,parameterCategory=pc,parameterNumber=pn,
258                                         fullName=varinfo[0],units=varinfo[1]))
259    return metadata
260
261
262def get_wgrib2_level_string(pdtn: int, pdt: ArrayLike) -> str:
263    """
264    Return a string that describes the level or layer of the GRIB2 message.
265
266    The format and language of the string is an exact replica of how wgrib2
267    produces the level/layer string in its inventory output.
268
269    Contents of wgrib2 source,
270    [Level.c](https://github.com/NOAA-EMC/NCEPLIBS-wgrib2/blob/develop/wgrib2/Level.c),
271    were converted into a Python dictionary and stored in grib2io as table
272    'wgrib2_level_string'.
273
274    Parameters
275    ----------
276    pdtn
277        GRIB2 Product Definition Template Number
278    pdt
279        Sequence containing GRIB2 Product Definition Template (Section 4).
280
281    Returns
282    -------
283    get_wgrib2_level_string
284        wgrib2-formatted level/layer string.
285    """
286    lvlstr = ''
287    if pdtn == 32:
288        return 'no_level'
289    elif pdtn == 48:
290        idxs = slice(20,26)
291    else:
292        idxs = slice(9,15)
293    type1, sfac1, sval1, type2, sfac2, sval2 = map(int,pdt[idxs])
294    val1 = sval1/10**sfac1
295    if type1 in [100,108]: val1 *= 0.01
296    if type2 != 255:
297        # Layer
298        #assert type2 == type1, "Surface types are not equal: %g - %g" % (type1,type2)
299        val2 = sval2/10**sfac2
300        if type2 in [100,108]: val2 *= 0.01
301        lvlstr = get_value_from_table(type1,table='wgrib2_level_string')[1]
302        vals = (val1,val2)
303    else:
304        # Level
305        lvlstr = get_value_from_table(type1,table='wgrib2_level_string')[0]
306        vals = (val1)
307    if '%g' in lvlstr: lvlstr %= vals
308    return lvlstr
309
310
311def _build_aerosol_shortname(obj) -> str:
312    """
313    """
314    _OPTICAL_WAVELENGTH_MAPPING = get_table('aerosol_optical_wavelength')
315    _LEVEL_MAPPING = get_table('aerosol_level')
316    _PARAMETER_MAPPING = get_table('aerosol_parameter')
317    _AERO_TYPE_MAPPING = get_table('aerosol_type')
318
319    # Build shortname from aerosol components
320    parts = []
321
322    # Get aerosol type
323    aero_type = str(obj.typeOfAerosol.value) if obj.typeOfAerosol is not None else ""
324
325    # Add size information if applicable
326    aero_size = ""
327    if hasattr(obj, 'scaledValueOfFirstSize'):
328        if float(obj.scaledValueOfFirstSize) > 0:
329            first_size = float(obj.scaledValueOfFirstSize)
330
331            # Map common PM sizes
332            size_map = {1: 'pm1', 25: 'pm25', 10: 'pm10', 20: 'pm20'}
333            aero_size = size_map.get(first_size, f"pm{int(first_size)}")
334
335            # Check for size intervals
336            if (hasattr(obj, 'scaledValueOfSecondSize') and
337                obj.scaledValueOfSecondSize is not None and
338                hasattr(obj, 'typeOfIntervalForAerosolSize') and
339                obj.typeOfIntervalForAerosolSize.value == 6):
340
341                second_size = float(obj.scaledValueOfSecondSize)
342                if second_size > 0:
343                    if (first_size == 2.5 and second_size == 10):
344                        aero_size = 'PM25to10'
345                    elif (first_size == 10 and second_size == 20):
346                        aero_size = 'PM10to20'
347                    else:
348                        aero_size = f"PM{int(first_size)}to{int(second_size)}"
349
350    # Add optical and wavelength information
351    var_wavelength = ''
352    if (hasattr(obj, 'parameterNumber') and
353        hasattr(obj, 'scaledValueOfFirstWavelength') and
354        hasattr(obj, 'scaledValueOfSecondWavelength')):
355
356        optical_type = str(obj.parameterNumber)
357        if obj.scaledValueOfFirstWavelength > 0:
358            first_wl = obj.scaledValueOfFirstWavelength
359            second_wl = obj.scaledValueOfSecondWavelength
360
361            # Special case for AE between 440-870nm
362            if optical_type == '111' and first_wl == 440 and second_wl == 870:
363                key = (optical_type, '440TO870')
364            else:
365                # Find matching wavelength band
366                for wl_key, wl_info in _OPTICAL_WAVELENGTH_MAPPING.items():
367                    if (wl_key[0] == optical_type and  # Check optical_type first
368                        int(wl_key[1]) == first_wl and
369                        (second_wl is None or int(wl_key[2]) == second_wl)):
370                        key = wl_key
371                        break
372                else:
373                    # If no match found, use raw values
374                    key = (optical_type, str(first_wl),
375                          str(second_wl) if second_wl is not None else '')
376
377            # FIX THIS...
378            if key in _OPTICAL_WAVELENGTH_MAPPING.keys():
379                var_wavelength = _OPTICAL_WAVELENGTH_MAPPING[key]
380
381    # Add level information
382    level_str = ''
383    if hasattr(obj, 'typeOfFirstFixedSurface'):
384        first_level = str(obj.typeOfFirstFixedSurface.value)
385        first_value = str(obj.scaledValueOfFirstFixedSurface) if obj.scaledValueOfFirstFixedSurface > 0 else ''
386        if first_level in _LEVEL_MAPPING:
387            level_str = f"{_LEVEL_MAPPING[first_level]}{first_value}"
388
389    # Get parameter type
390    param = ''
391    if hasattr(obj, 'parameterNumber'):
392        param_num = str(obj.parameterNumber)
393        if param_num in _PARAMETER_MAPPING:
394            param = _PARAMETER_MAPPING[param_num]
395
396    # Build the final shortname
397    if var_wavelength and aero_type in _AERO_TYPE_MAPPING:
398        shortname = f"{_AERO_TYPE_MAPPING[aero_type]}{var_wavelength}"
399    elif aero_type in _AERO_TYPE_MAPPING:
400        parts = []
401        if level_str:
402            parts.append(level_str)
403        parts.append(_AERO_TYPE_MAPPING[aero_type])
404        if aero_size:
405            parts.append(aero_size)
406        if param:
407            parts.append(param)
408        shortname = '_'.join(parts) if len(parts) > 1 else parts[0]
409    else:
410        return get_varinfo_from_table(obj.section0[2], *obj.section4[2:4], isNDFD=obj._isNDFD)[2]
411
412    return shortname
413
414
415@lru_cache(maxsize=None)
416def get_table_names(var_tables: bool=False) -> tuple:
417    """
418    Return the names of available GRIB2 tables.
419
420    The table dictionaries as they exist in grib2io are prefixed
421    with "table_". When accessing a table through the available
422    functions, the prefix is not necessary.
423
424    Parameters
425    ----------
426    var_tables : bool, optional
427        If True, include names of variable tables.
428        If False (default), include only non-variable tables.
429
430    Returns
431    -------
432    tuple of str
433        A tuple sorted names of available tables.
434    """
435    tables = [
436        name.replace("table_","") for name, val in globals().items()
437        if isinstance(val, dict) and name.startswith("table_")
438    ]
439
440    tables = [
441        t.replace("_",".") if t.startswith(tuple("0123456789")) else t for t in tables
442    ]
443
444    if var_tables:
445        for d in GRIB2_DISCIPLINES:
446            modname = f'.section4_discipline{d}'
447            _load_varinfo_tables(modname)
448        vt = [t.replace("table_","").replace("_",".") for t in _varinfo_tables_datastore.keys()]
449        tables.extend(vt)
450
451    return tuple(sorted(tables))