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 *  # noqa: F403
 10from .section1 import *  # noqa: F403
 11from .section3 import *  # noqa: F403
 12from .section4 import *  # noqa: F403
 13from .section5 import *  # noqa: F403
 14from .section6 import *  # noqa: F403
 15from .originating_centers import *  # noqa: F403
 16from .ndfd_additionals import *  # noqa: F403
 17from .cf import *  # noqa: F403
 18
 19_varinfo_tables_datastore = {}
 20
 21# Note 209 is MRMS
 22GRIB2_DISCIPLINES = [0, 1, 2, 3, 4, 10, 20, 209]
 23
 24AEROSOL_PDTNS = [44, 45, 46, 47, 48, 49, 50, 80, 81, 82, 83, 84, 85]
 25AEROSOL_PARAMS = list(itertools.chain(range(0, 19), range(50, 82), range(100, 113), range(192, 197)))
 26CHEMICAL_PDTNS = [40, 41, 42, 43, 57, 58, 67, 68, 76, 77, 78, 79]
 27
 28
 29def _load_varinfo_tables(modname: str):
 30    """
 31    Load variable information tables from sub modules into the local
 32    varinfo table datastore (_varinfo_tables_datastore).
 33
 34    Parameters
 35    ----------
 36    modname
 37        Module name to extract variable info tables from.
 38    """
 39    module = importlib.import_module(modname, package=__name__)
 40    names = getattr(module, "__all__", [name for name in dir(module) if name.startswith("table_")])
 41    _varinfo_tables_datastore.update({name: getattr(module, name) for name in names})
 42
 43
 44def get_table(table: str, expand: bool = False) -> dict:
 45    """
 46    Return GRIB2 code table as a dictionary.
 47
 48    Parameters
 49    ----------
 50    table
 51        Code table number (e.g. '1.0').
 52
 53        NOTE: Code table '4.1' requires a 3rd value representing the product
 54        discipline (e.g. '4.1.0').
 55    expand
 56        If `True`, expand output dictionary wherever keys are a range.
 57
 58    Returns
 59    -------
 60    get_table
 61        GRIB2 code table as a dictionary.
 62    """
 63    if len(table) == 3 and table == "4.1":
 64        raise Exception("GRIB2 Code Table 4.1 requires a 3rd value representing the discipline.")
 65    if len(table) == 3 and table.startswith("4.2"):
 66        raise Exception("Use function get_varinfo_from_table() for GRIB2 Code Table 4.2")
 67    try:
 68        tbl = globals()["table_" + table.replace(".", "_")]
 69        if expand:
 70            _tbl = {}
 71            for k, v in tbl.items():
 72                if "-" in k:
 73                    irng = [int(i) for i in k.split("-")]
 74                    for i in range(irng[0], irng[1] + 1):
 75                        _tbl[str(i)] = v
 76                else:
 77                    _tbl[k] = v
 78            tbl = _tbl
 79        return tbl
 80    except KeyError:
 81        return {}
 82
 83
 84def get_value_from_table(
 85    value: Union[int, str],
 86    table: str,
 87    expand: bool = False,
 88) -> Optional[Union[float, int]]:
 89    """
 90    Return the definition given a GRIB2 code table.
 91
 92    Parameters
 93    ----------
 94    value
 95        Code table value.
 96    table
 97        Code table number.
 98    expand
 99        If `True`, expand output dictionary where keys are a range.
100
101    Returns
102    -------
103    get_value_from_table
104        Table value or `None` if not found.
105    """
106    try:
107        tbl = get_table(table, expand=expand)
108        value = str(value)
109        return tbl[value]
110    except KeyError:
111        for k in tbl.keys():
112            if "-" in k:
113                bounds = k.split("-")
114                if value >= bounds[0] and value <= bounds[1]:
115                    return tbl[k]
116        return None
117
118
119def get_varinfo_from_table(
120    discipline: Union[int, str],
121    parmcat: Union[int, str],
122    parmnum: Union[int, str],
123    isNDFD: bool = False,
124):
125    """
126    Return the GRIB2 variable information.
127
128    NOTE: This functions allows for all arguments to be converted to a string
129    type if arguments are integer.
130
131    Parameters
132    ----------
133    discipline
134        Discipline code value of a GRIB2 message.
135    parmcat
136        Parameter Category value of a GRIB2 message.
137    parmnum
138        Parameter Number value of a GRIB2 message.
139    isNDFD: optional
140        If `True`, then attempt to get variable information from the
141        supplemental NDFD tables, otherwise fallback to default tables.
142
143    Returns
144    -------
145    full_name
146        Full name of the GRIB2 variable. "Unknown" if variable is not found.
147    units
148        Units of the GRIB2 variable. "Unknown" if variable is not found.
149    shortName
150        Abbreviated name of the GRIB2 variable. "Unknown" if variable is not
151        found.
152    """
153    template = f"table_4_2_{discipline}_{parmcat}"
154    tbls = []
155    if isNDFD:
156        tbls.append(template + "_ndfd")
157    tbls.append(template)
158    for tbl in tbls:
159        vartbl = None
160        modname = f".section4_discipline{discipline}"
161        if tbl not in _varinfo_tables_datastore.keys():
162            _load_varinfo_tables(modname)
163        try:
164            vartbl = _varinfo_tables_datastore[tbl][str(parmnum)]
165        except KeyError:
166            pass
167        if vartbl is not None:
168            break
169    if vartbl is None:
170        vartbl = ["Unknown", "Unknown", "Unknown"]
171    return vartbl
172
173
174@lru_cache(maxsize=None)
175def get_shortnames(
176    discipline: Optional[Union[int, str]] = None,
177    parmcat: Optional[Union[int, str]] = None,
178    parmnum: Optional[Union[int, str]] = None,
179    isNDFD: bool = False,
180) -> List[str]:
181    """
182    Return a list of variable shortNames.
183
184    If all 3 args are None, then shortNames from all disciplines, parameter
185    categories, and numbers will be returned.
186
187    Parameters
188    ----------
189    discipline
190        GRIB2 discipline code value.
191    parmcat
192        GRIB2 parameter category value.
193    parmnum
194        Parameter Number value of a GRIB2 message.
195    isNDFD: optional
196        If `True`, signals function to try to get variable information from the
197        supplemental NDFD tables.
198
199    Returns
200    -------
201    get_shortnames
202        list of GRIB2 shortNames.
203    """
204    shortnames = list()
205    if discipline is None:
206        discipline = GRIB2_DISCIPLINES
207    else:
208        discipline = [discipline]
209    if parmcat is None:
210        parmcat = list()
211        for d in discipline:
212            parmcat += list(get_table(f"4.1.{d}").keys())
213    else:
214        parmcat = [parmcat]
215    if parmnum is None:
216        parmnum = list(range(256))
217    else:
218        parmnum = [parmnum]
219    for d in discipline:
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(
260                        dict(
261                            discipline=d,
262                            parameterCategory=pc,
263                            parameterNumber=pn,
264                            fullName=varinfo[0],
265                            units=varinfo[1],
266                        )
267                    )
268    return metadata
269
270
271def get_wgrib2_level_string(pdtn: int, pdt: ArrayLike) -> str:
272    """
273    Return a string that describes the level or layer of the GRIB2 message.
274
275    The format and language of the string is an exact replica of how wgrib2
276    produces the level/layer string in its inventory output.
277
278    Contents of wgrib2 source,
279    [Level.c](https://github.com/NOAA-EMC/NCEPLIBS-wgrib2/blob/develop/wgrib2/Level.c),
280    were converted into a Python dictionary and stored in grib2io as table
281    'wgrib2_level_string'.
282
283    Parameters
284    ----------
285    pdtn
286        GRIB2 Product Definition Template Number
287    pdt
288        Sequence containing GRIB2 Product Definition Template (Section 4).
289
290    Returns
291    -------
292    get_wgrib2_level_string
293        wgrib2-formatted level/layer string.
294    """
295    lvlstr = ""
296    if pdtn == 32:
297        return "no_level"
298    elif pdtn == 48:
299        idxs = slice(20, 26)
300    else:
301        idxs = slice(9, 15)
302    type1, sfac1, sval1, type2, sfac2, sval2 = map(int, pdt[idxs])
303    val1 = sval1 / 10**sfac1
304    if type1 in [100, 108]:
305        val1 *= 0.01
306    if type2 != 255:
307        # Layer
308        # assert type2 == type1, "Surface types are not equal: %g - %g" % (type1,type2)
309        val2 = sval2 / 10**sfac2
310        if type2 in [100, 108]:
311            val2 *= 0.01
312        lvlstr = get_value_from_table(type1, table="wgrib2_level_string")[1]
313        vals = (val1, val2)
314    else:
315        # Level
316        lvlstr = get_value_from_table(type1, table="wgrib2_level_string")[0]
317        vals = val1
318    if "%g" in lvlstr:
319        lvlstr %= vals
320    return lvlstr
321
322
323def _build_chemical_shortname(obj) -> str:
324    """
325    Build shortname for chemical constituent GRIB2 messages.
326    """
327    _CHEMICAL_MAPPING = get_table("4.230")
328    _AERO_MAPPING = get_table("4.233")
329    _PARAMETER_MAPPING = get_table("aerosol_parameter")
330    _LEVEL_MAPPING = get_table("aerosol_level")
331
332    # Get constituent type
333    const_metadata = getattr(obj, "constituentType", None)
334    constituent_type = str(const_metadata.value) if const_metadata is not None else ""
335    chemical_abbr = ""
336    if constituent_type in _CHEMICAL_MAPPING:
337        chemical_abbr = _CHEMICAL_MAPPING[constituent_type][1]
338    elif constituent_type in _AERO_MAPPING:
339        chemical_abbr = _AERO_MAPPING[constituent_type][1]
340
341    if chemical_abbr == "unknown":
342        chemical_abbr = ""
343
344    # Get parameter type
345    param = ""
346    if hasattr(obj, "parameterNumber"):
347        param_num = str(obj.parameterNumber)
348        # Use a specific mapping for chemical parameters
349        chemical_params = {
350            "0": "mr",  # Mass Density
351            "1": "colmd",  # Column-Integrated Mass Density
352            "2": "mmr",  # Mass Mixing Ratio
353            "52": "vmr",  # Volume Mixing Ratio
354        }
355        if param_num in chemical_params:
356            param = chemical_params[param_num]
357        elif param_num in _PARAMETER_MAPPING:
358            param = _PARAMETER_MAPPING[param_num]
359        else:
360            param = f"param{param_num}"
361
362    # Handle source/sink
363    source_sink = ""
364    if hasattr(obj, "sourceSinkIndicator") and obj.sourceSinkIndicator is not None:
365        ss_val = str(obj.sourceSinkIndicator.value)
366        if ss_val != "255":
367            source_sink = f"ss{ss_val}"
368
369    # Build the final shortname
370    parts = []
371    if chemical_abbr:
372        parts.append(chemical_abbr)
373    if param:
374        parts.append(param)
375    if source_sink:
376        parts.append(source_sink)
377
378    if not parts:
379        return get_varinfo_from_table(obj.section0[2], *obj.section4[2:4], isNDFD=obj._isNDFD)[2]
380
381    return "_".join(parts)
382
383
384def _build_aerosol_shortname(obj) -> str:
385    """ """
386    _OPTICAL_WAVELENGTH_MAPPING = get_table("aerosol_optical_wavelength")
387    _LEVEL_MAPPING = get_table("aerosol_level")
388    _PARAMETER_MAPPING = get_table("aerosol_parameter")
389    _AERO_TYPE_MAPPING = get_table("aerosol_type")
390    _AERO_TABLE_4233 = get_table("4.233")
391
392    # Build shortname from aerosol components
393    parts = []
394
395    # Get aerosol type
396    aero_metadata = getattr(obj, "typeOfAerosol", None)
397    aero_type = str(aero_metadata.value) if aero_metadata is not None else ""
398
399    # Try to get abbreviation from optimized mapping or Table 4.233
400    aero_abbr = ""
401    if aero_type in _AERO_TYPE_MAPPING:
402        aero_abbr = _AERO_TYPE_MAPPING[aero_type]
403    elif aero_type in _AERO_TABLE_4233:
404        aero_abbr = _AERO_TABLE_4233[aero_type][1]
405        if aero_abbr == "unknown":
406            aero_abbr = ""
407
408    # Add size information if applicable
409    aero_size = ""
410    if hasattr(obj, "scaledValueOfFirstSize"):
411        if float(obj.scaledValueOfFirstSize) > 0:
412            first_size = float(obj.scaledValueOfFirstSize)
413
414            # Map common PM sizes
415            size_map = {1: "pm1", 2.5: "pm25", 25: "pm25", 10: "pm10", 20: "pm20"}
416            aero_size = size_map.get(first_size, f"pm{int(first_size)}")
417
418            # Check for size intervals
419            if (
420                hasattr(obj, "scaledValueOfSecondSize")
421                and obj.scaledValueOfSecondSize is not None
422                and hasattr(obj, "typeOfIntervalForAerosolSize")
423                and obj.typeOfIntervalForAerosolSize.value == 6
424            ):
425                second_size = float(obj.scaledValueOfSecondSize)
426                if second_size > 0:
427                    if first_size == 2.5 and second_size == 10:
428                        aero_size = "pm25to10"
429                    elif first_size == 10 and second_size == 20:
430                        aero_size = "pm10to20"
431                    else:
432                        aero_size = f"pm{int(first_size)}to{int(second_size)}"
433
434    # Add optical and wavelength information
435    var_wavelength = ""
436    if hasattr(obj, "parameterNumber") and hasattr(obj, "scaledValueOfFirstWavelength") and hasattr(obj, "scaledValueOfSecondWavelength"):
437        optical_type = str(obj.parameterNumber)
438        if obj.scaledValueOfFirstWavelength > 0:
439            first_wl = obj.scaledValueOfFirstWavelength
440            second_wl = obj.scaledValueOfSecondWavelength
441
442            # Special case for AE between 440-870nm
443            if optical_type == "111" and first_wl == 440 and second_wl == 870:
444                key = (optical_type, "440TO870")
445            else:
446                # Find matching wavelength band
447                for wl_key, wl_info in _OPTICAL_WAVELENGTH_MAPPING.items():
448                    if wl_key[0] != optical_type:
449                        continue
450                    if int(wl_key[1]) != first_wl:
451                        continue
452                    if len(wl_key) == 3:
453                        if second_wl is not None and int(wl_key[2]) == second_wl:
454                            key = wl_key
455                            break
456                    elif len(wl_key) == 2:
457                        if second_wl is None:
458                            key = wl_key
459                            break
460                else:
461                    # If no match found, use raw values
462                    key = (
463                        optical_type,
464                        str(first_wl),
465                        str(second_wl) if second_wl is not None else "",
466                    )
467
468            # FIX THIS...
469            if key in _OPTICAL_WAVELENGTH_MAPPING.keys():
470                var_wavelength = _OPTICAL_WAVELENGTH_MAPPING[key]
471
472    # Get parameter type
473    param = ""
474    if hasattr(obj, "parameterNumber"):
475        param_num = str(obj.parameterNumber)
476        if param_num in _PARAMETER_MAPPING:
477            param = _PARAMETER_MAPPING[param_num]
478        else:
479            param = f"param{param_num}"
480
481    # Handle source/sink
482    source_sink = ""
483    if hasattr(obj, "sourceSinkIndicator") and obj.sourceSinkIndicator is not None:
484        ss_val = str(obj.sourceSinkIndicator.value)
485        if ss_val != "255":
486            source_sink = f"ss{ss_val}"
487
488    # Build the final shortname
489    if var_wavelength and aero_abbr:
490        shortname = f"{aero_abbr}{var_wavelength}"
491        if source_sink:
492            shortname = f"{shortname}_{source_sink}"
493    elif aero_abbr:
494        parts = []
495        parts.append(aero_abbr)
496        if aero_size:
497            parts.append(aero_size)
498        if param:
499            parts.append(param)
500        if source_sink:
501            parts.append(source_sink)
502        shortname = "_".join(parts) if len(parts) > 1 else parts[0]
503    else:
504        return get_varinfo_from_table(obj.section0[2], *obj.section4[2:4], isNDFD=obj._isNDFD)[2]
505
506    return shortname
507
508
509@lru_cache(maxsize=None)
510def get_table_names(var_tables: bool = False) -> tuple:
511    """
512    Return the names of available GRIB2 tables.
513
514    The table dictionaries as they exist in grib2io are prefixed
515    with "table_". When accessing a table through the available
516    functions, the prefix is not necessary.
517
518    Parameters
519    ----------
520    var_tables : bool, optional
521        If True, include names of variable tables.
522        If False (default), include only non-variable tables.
523
524    Returns
525    -------
526    tuple of str
527        A tuple sorted names of available tables.
528    """
529    tables = [name.replace("table_", "") for name, val in globals().items() if isinstance(val, dict) and name.startswith("table_")]
530
531    tables = [t.replace("_", ".") if t.startswith(tuple("0123456789")) else t for t in tables]
532
533    if var_tables:
534        for d in GRIB2_DISCIPLINES:
535            modname = f".section4_discipline{d}"
536            _load_varinfo_tables(modname)
537        vt = [t.replace("table_", "").replace("_", ".") for t in _varinfo_tables_datastore.keys()]
538        tables.extend(vt)
539
540    return tuple(sorted(tables))