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