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