1"""
   2grib2io Backend Engine for Xarray
   3=================================
   4grib2io provides a Xarray backend entrypoint for decoding many GRIB2 messages
   5from a single file or many files and represented as Xarray DataArray objects and
   6collected along common coordinates as Datasets and DataTrees.
   7
   8.. warning::
   9
  10   The ``grib2io.xarray_backend`` engine API is **experimental**.
  11   Its interface and behavior may change in future releases,
  12   which could affect backward compatibility.
  13
  14   Users are encouraged to treat this backend as subject to change
  15   and to pin their ``grib2io`` version if depending on its current
  16   implementation details.
  17"""
  18from grib2io._grib2io import _data
  19from grib2io import Grib2Message, Grib2GridDef, msgs_from_index
  20import grib2io
  21from xarray.backends.locks import SerializableLock
  22from xarray.core import indexing
  23from xarray.backends import (
  24    BackendArray,
  25    BackendEntrypoint,
  26)
  27from copy import copy
  28from collections import defaultdict
  29from dataclasses import dataclass, field, astuple
  30import importlib.metadata
  31import itertools
  32import logging
  33import typing
  34import warnings
  35
  36from . import tables
  37
  38import numpy as np
  39import pandas as pd
  40import xarray as xr
  41import re
  42from pyproj import CRS
  43import datetime
  44
  45# Check if xarray version supports DataTree
  46_HAS_DATATREE = False
  47try:
  48    # Try importing DataTree to check if it's available
  49    xarray_version = importlib.metadata.version('xarray')
  50    xarray_parts = [int(x) if x.isdigit() else x for x in xarray_version.split('.')]
  51    min_version_parts = [2024, 10, 0]
  52    _HAS_DATATREE = xarray_parts >= min_version_parts
  53except (ImportError, ValueError):
  54    _HAS_DATATREE = False
  55
  56_logger = logging.getLogger(__name__)
  57
  58_LOCK = SerializableLock()
  59
  60_LEVEL_NAME_MAPPING = grib2io.tables.get_table('4.5.grib2io.level.name')
  61
  62_TREE_HIERARCHY_LEVELS = [
  63    "typeOfFirstFixedSurface",
  64    "valueOfFirstFixedSurface",
  65    "productDefinitionTemplateNumber",
  66    "perturbationNumber",
  67    "leadTime",
  68    "duration",
  69    "percentileValue",
  70    "typeOfProbability",
  71    "thresholdLowerLimit",
  72    "thresholdUpperLimit"
  73]
  74
  75AVAILABLE_NON_GEO_COORDS = [
  76    "duration",
  77    "leadTime",
  78    "percentileValue",
  79    "perturbationNumber",
  80    "refDate",
  81    "thresholdLowerLimit",
  82    "thresholdUpperLimit",
  83    "valueOfFirstFixedSurface",
  84    "valueOfSecondFixedSurface",
  85    "aerosolType",
  86    "scaledValueOfFirstWavelength",
  87    "scaledValueOfSecondWavelength",
  88    "scaledValueOfCentralWaveNumber",
  89    "scaledValueOfFirstSize",
  90    "scaledValueOfSecondSize"
  91]
  92"""Available non-geographic coordinate names."""
  93
  94AVAILABLE_NON_GEO_DIMS = [
  95    "duration",
  96    "leadTime",
  97    "percentileValue",
  98    "perturbationNumber",
  99    "refDate",
 100    "threshold",
 101    "level",
 102]
 103"""Available non-geographic dimension names."""
 104
 105# Lookup table to define surface types that should be parsed as vertical coordinates
 106VERTICAL_COORDINATE_SURFACES = [
 107    "Ground or Water Surface",
 108    "Isothermal Level",
 109    "Specified radius from the centre of the Sun",
 110    "Isobaric Surface",
 111    "Mean Sea Level",
 112    "Specific Altitude Above Mean Sea Level",
 113    "Specified Height Level Above Ground",
 114    "Sigma Level",
 115    "Hybrid Level",
 116    "Depth Below Land Surface",
 117    "Isentropic (theta) Level",
 118    "Level at Specified Pressure Difference from Ground to Level",
 119    "Potential Vorticity Surface",
 120    "Eta Level",
 121    "Logarithmic Hybrid Level",
 122    "Sigma height level",
 123    "Hybrid Height Level",
 124    "Hybrid Pressure Level",
 125    "Soil level",
 126    "Sea-ice level",
 127    "Depth Below Sea Level",
 128    "Depth Below Water Surface",
 129    "Ocean Model Level",
 130    "Ocean level defined by water density (sigma-theta) difference from near-surface to level",
 131    "Ocean level defined by water potential temperature difference from near-surface to level",
 132    "Ocean level defined by vertical eddy diffusivity difference from near-surface to level",
 133    "Ocean level defined by water density (rho) difference from near-surface to level"
 134]
 135"""
 136Lookup table to define surface types that should be parsed as vertical coordinates
 137when `data_model="nws-viz"`.
 138"""
 139
 140def parse_data_model(ds, data_model):
 141    """
 142    Normalize a GRIB2-derived Dataset to a target data model (currently ``"nws-viz"``).
 143
 144    When ``data_model == "nws-viz"``, this function converts coordinate and
 145    variable names to snake_case, derives CF-like metadata, promotes select
 146    GRIB-derived quantities to coordinates, optionally swaps dimensions, and
 147    standardizes units/attributes. If ``data_model`` is anything else, the
 148    input dataset is returned unchanged.
 149
 150    Parameters
 151    ----------
 152    ds : xarray.Dataset
 153        GRIB2-derived dataset whose variables and attributes follow the
 154        conventions emitted by ``grib2io``. Expected to contain GRIB-related
 155        attributes such as ``typeOfFirstFixedSurface``,
 156        ``typeOfSecondFixedSurface``, and (for probabilistic variables)
 157        ``typeOfProbability``.
 158    data_model : str
 159        Target data model name. Only the value ``"nws-viz"`` triggers
 160        transformations.
 161
 162    Returns
 163    -------
 164    xarray.Dataset
 165        A new dataset with:
 166        * Selected coordinates renamed:
 167          ``refDate -> forecast_reference_time``,
 168          ``leadTime -> lead_time``,
 169          ``validDate -> time``,
 170          ``percentileValue -> percentile``,
 171          ``thresholdLowerLimit -> threshold_lower_limit``,
 172          ``thresholdUpperLimit -> threshold_upper_limit``.
 173        * Vertical coordinates derived from
 174          ``valueOfFirstFixedSurface`` / ``valueOfSecondFixedSurface`` and their
 175          corresponding ``typeOf*FixedSurface`` definitions. New coordinate
 176          names are generated from the surface definition (lowercased, spaces
 177          to underscores, punctuation removed). If the name already exists, a
 178          ``"_2"`` suffix is appended.
 179        * Possible dimension swaps:
 180          ``level -> <derived_vertical_coord>`` when present; and for
 181          probabilistic variables, ``threshold -> threshold_lower_limit`` or
 182          ``threshold -> threshold_upper_limit`` when
 183          ``typeOfProbability`` indicates the appropriate semantics.
 184        * Variable names lowercased; dataset- and variable-level attributes
 185          converted to snake_case (except GRIB section attributes which are
 186          normalized to ``grib...``).
 187        * CF-adjacent metadata populated: ``standard_name`` and
 188          ``cell_methods`` are set via the shortname→CF lookup table.
 189        * Percent units normalized from ``"%"`` to ``"percent"`` on coordinates.
 190        * For precipitation type (``PTYPE``) thresholds, numeric codes are
 191          decoded to strings (GRIB2 Table 4.201) in relevant attrs/coords.
 192
 193    Notes
 194    -----
 195    - Precipitation type decoding uses GRIB2 Table 4.201 via
 196      ``tables.get_value_from_table(code, "4.201")`` and returns a NumPy
 197      array with ``np.dtypes.StringDType``.
 198    - CF-related lookups are performed using
 199      ``tables.get_table("shortname_to_cf")``.
 200    - Vertical coordinate surface names are validated against
 201      ``VERTICAL_COORDINATE_SURFACES`` before promotion to coordinates.
 202
 203    Warnings
 204    --------
 205    This function assumes the presence of certain GRIB-derived attributes on the
 206    first data variable (e.g., ``typeOfFirstFixedSurface``,
 207    ``typeOfSecondFixedSurface``, and possibly ``typeOfProbability``).
 208    If these are absent or malformed, errors (e.g., ``KeyError``) may occur.
 209
 210    Examples
 211    --------
 212    >>> ds2 = parse_data_model(ds, "nws-viz")
 213    >>> list(ds2.coords)
 214    ['forecast_reference_time', 'lead_time', 'time', 'percentile', ...]
 215    """
 216
 217    def _decode_ptype(values):
 218        """
 219        Decode precipitation type values into human-readable strings.
 220
 221        Uses GRIB2 Table 4.201 to map numeric codes to precipitation type descriptions.
 222
 223        Parameters
 224        ----------
 225        values : array_like
 226            Array of numeric precipitation type codes (e.g., integers or floats).
 227            Each value corresponds to a GRIB2 Table 4.201 precipitation type code.
 228
 229        Returns
 230        -------
 231        numpy.ndarray
 232            Array of decoded precipitation type strings with
 233            NumPy’s flexible string data type (`np.dtypes.StringDType`).
 234        """
 235        results = []
 236        for val in values:
 237            # Convert each numeric code to string and look up in Table 4.201
 238            results.append(str(tables.get_value_from_table(str(int(val)), '4.201')))
 239
 240        # Return array of strings using numpy's string data type
 241        return np.array(results, dtype=np.dtypes.StringDType)
 242
 243    # convert coordinates and attributes to CF if requested
 244    if data_model == 'nws-viz':
 245
 246        # define regex to convert to snake case
 247        pattern = re.compile(r'(?<!^)(?=[A-Z])')
 248
 249        # check for coordinates and rename
 250        for coord in ds.coords:
 251            if coord == 'refDate':
 252                ds = ds.rename({'refDate': 'forecast_reference_time'})
 253
 254            elif coord == 'leadTime':
 255                ds = ds.rename({'leadTime': 'lead_time'})
 256
 257            elif coord == 'validDate':
 258                ds = ds.rename({'validDate': 'time'})
 259
 260            elif coord == 'percentileValue':
 261                ds = ds.rename({'percentileValue': 'percentile'})
 262
 263            elif coord == 'thresholdLowerLimit':
 264                ds = ds.rename({'thresholdLowerLimit': 'threshold_lower_limit'})
 265                ds['threshold_lower_limit'].attrs['long_name'] = 'Threshold Lower Limit'
 266                ds['threshold_lower_limit'].attrs['units'] = ds[list(ds.data_vars.keys())[0]].attrs['units']
 267
 268                if 'PTYPE' in ds.data_vars:
 269                    ds['threshold_lower_limit'] = xr.apply_ufunc(_decode_ptype, ds['threshold_lower_limit'])
 270
 271                # check if thresholdLowerLimit should be a dimension coordinate
 272                if 'threshold' in ds.dims:
 273                    var_key = list(ds.data_vars.keys())[0]
 274                    prob_types = [
 275                        'Probability of event below lower limit',
 276                        'Probability of event above lower limit',
 277                        'Probability of event equal to lower limit',
 278                        'Probability of event between upper and lower limits (the range includes lower limit but not the upper limit)'
 279                    ]
 280                    if ds[var_key].attrs['typeOfProbability'] in prob_types:
 281                        ds = ds.swap_dims({'threshold': 'threshold_lower_limit'})
 282
 283            elif coord == 'thresholdUpperLimit':
 284                ds = ds.rename({'thresholdUpperLimit': 'threshold_upper_limit'})
 285                ds['threshold_upper_limit'].attrs['long_name'] = 'Threshold Upper Limit'
 286                ds['threshold_upper_limit'].attrs['units'] = ds[list(ds.data_vars.keys())[0]].attrs['units']
 287
 288                if 'PTYPE' in ds.data_vars:
 289                    ds['threshold_upper_limit'] = xr.apply_ufunc(_decode_ptype, ds['threshold_upper_limit'])
 290
 291                if 'threshold' in ds.dims:
 292                    var_key = list(ds.data_vars.keys())[0]
 293                    prob_types = [
 294                        'Probability of event below upper limit',
 295                        'Probability of event above upper limit'
 296                    ]
 297                    if ds[var_key].attrs['typeOfProbability'] in prob_types:
 298                        ds = ds.swap_dims({'threshold': 'threshold_upper_limit'})
 299
 300            # If the dataset has valueOfFirstFixedSurface as a coordinate
 301            elif coord == 'valueOfFirstFixedSurface':
 302                # Get the valueOfFirstFixedSurface coordinate
 303                da = ds.valueOfFirstFixedSurface
 304
 305                # Get the definition and units from typeOfFirstFixedSurface
 306                var_key = list(ds.data_vars.keys())[0]
 307                definition, units = ds[var_key].attrs['typeOfFirstFixedSurface']
 308
 309                if definition in VERTICAL_COORDINATE_SURFACES:
 310                    # Convert definition to lowercase and replace spaces with underscores
 311                    key = definition.lower().replace(' ', '_')
 312
 313                    # remove special characters
 314                    key = re.sub(r'[^a-z0-9_]', '', key)
 315
 316                    # Add units and grib_name attributes
 317                    da.attrs['units'] = units
 318                    da.attrs['grib_name'] = ['valueOfFirstFixedSurface', 'typeOfFirstFixedSurface']
 319
 320                    # Assign the coordinate with the new key name
 321                    ds = ds.assign_coords({key: da})
 322
 323                    # If valueOfFirstFixedSurface is a dimension, swap it with the new key
 324                    if 'level' in ds.dims:
 325                        ds = ds.swap_dims({"level": key})
 326
 327                # Remove the original coordinates
 328                del ds['valueOfFirstFixedSurface']
 329
 330            # If the dataset has valueOfSecondFixedSurface as a coordinate
 331            elif coord == 'valueOfSecondFixedSurface':
 332                # Get the valueOfSecondFixedSurface coordinate
 333                da = ds.valueOfSecondFixedSurface
 334
 335                # Get the definition and units from typeOfSecondFixedSurface
 336                var_key = list(ds.data_vars.keys())[0]
 337                definition, units = ds[var_key].attrs['typeOfSecondFixedSurface']
 338
 339                if definition in VERTICAL_COORDINATE_SURFACES:
 340                    # Convert definition to lowercase and replace spaces with underscores
 341                    key = definition.lower().replace(' ', '_')
 342
 343                    # remove special characters
 344                    key = re.sub(r'[^a-z0-9_]', '', key)
 345
 346                    # check if key is already in coords
 347                    if key in ds.coords:
 348                        key = key + '_2'
 349
 350                    # Add units and grib_name attributes
 351                    da.attrs['units'] = units
 352                    da.attrs['grib_name'] = ['valueOfSecondFixedSurface', 'typeOfSecondFixedSurface']
 353
 354                    # Assign the coordinate with the new key name
 355                    ds = ds.assign_coords({key: da})
 356
 357                # Remove the original coordinates
 358                del ds['valueOfSecondFixedSurface']
 359            else:
 360                # change coord name to snake case
 361                new_coord_name = pattern.sub('_', coord).lower()
 362                ds = ds.rename({coord: new_coord_name})
 363
 364        # convert all attributes and variable names to snake case
 365        for var in ds.data_vars:
 366            da = ds[var]
 367            record = tables.get_table('shortname_to_cf').get(da.name)
 368            da.attrs['standard_name'] = 'unknown' if record is None else record['cf_standard_name']
 369            da.attrs['cell_methods'] = 'unknown' if record is None else record['cf_cell_methods']
 370
 371            ds[var] = da
 372
 373            # rename variable
 374            new_var_name = var.lower()
 375            ds = ds.rename({var: new_var_name})
 376
 377            # remove attr for typeOfFirstFixedSurface (applied as coordinate above)
 378            if 'typeOfFirstFixedSurface' in ds[new_var_name].attrs:
 379                definition, units = ds[new_var_name].attrs['typeOfFirstFixedSurface']
 380                ds[new_var_name].attrs['typeOfFirstFixedSurface'] = f'{definition} ({units})'
 381
 382            if 'typeOfSecondFixedSurface' in ds[new_var_name].attrs:
 383                definition, units = ds[new_var_name].attrs['typeOfSecondFixedSurface']
 384                ds[new_var_name].attrs['typeOfSecondFixedSurface'] = f'{definition} ({units})'
 385
 386            ds[new_var_name].attrs.pop('percentileValue', None)
 387
 388            if 'threshold_lower_limit' in ds.coords:
 389                ds[new_var_name].attrs.pop('thresholdLowerLimit', None)
 390
 391            if 'threshold_upper_limit' in ds.coords:
 392                ds[new_var_name].attrs.pop('thresholdUpperLimit', None)
 393
 394            for attr in list(ds[new_var_name].attrs.keys()):
 395                # skip grib section attrs
 396                if 'GRIB2IO_section' in attr:
 397                    # replace GRIB2IO with grib in attr
 398                    new_attr_name = attr.replace('GRIB2IO', 'grib')
 399                else:
 400                    # change attr name to snake case
 401                    new_attr_name = pattern.sub('_', attr).lower()
 402
 403                # update new attr name for specific CF names
 404                if new_attr_name == 'full_name':
 405                    new_attr_name = 'long_name'
 406
 407                # change % to percent
 408                if attr == 'units' and ds[new_var_name].attrs[attr] == '%':
 409                    ds[new_var_name].attrs[attr] = 'percent'
 410
 411                if new_var_name == 'ptype' and 'threshold' in new_attr_name:
 412                    value = ds[new_var_name].attrs.pop(attr)
 413                    ds[new_var_name].attrs[attr] = _decode_ptype(value)
 414                else:
 415                    # change attr name in attrs
 416                    ds[new_var_name].attrs[new_attr_name] = ds[new_var_name].attrs.pop(attr)
 417
 418
 419        # change dataset attrs to snake case
 420        for attr in list(ds.attrs.keys()):
 421            # change attr name to snake case
 422            new_attr_name = pattern.sub('_', attr).lower()
 423
 424            # change attr name in attrs
 425            ds.attrs[new_attr_name] = ds.attrs.pop(attr)
 426
 427        # change % to percent
 428        for coord in ds.coords:
 429            if 'units' in ds[coord].attrs and ds[coord].attrs['units'] == '%':
 430                ds[coord].attrs['units'] = 'percent'
 431
 432    return ds
 433
 434
 435class GribBackendEntrypoint(BackendEntrypoint):
 436    """
 437    xarray backend engine entrypoint for opening and decoding grib2 files.
 438
 439    .. warning::
 440
 441       This backend is experimental and the API/behavior may change without
 442       backward compatibility.
 443    """
 444
 445    def open_dataset(
 446        self,
 447        filename,
 448        *,
 449        drop_variables=None,
 450        filters: typing.Mapping[str, typing.Any] = dict(),
 451        data_model=None
 452    ):
 453        """
 454        Read and parse metadata from grib file.
 455
 456        Parameters
 457        ----------
 458        filename
 459            GRIB2 file to be opened.
 460        filters
 461            Filter GRIB2 messages to single hypercube. Dict keys can be any
 462            GRIB2 metadata attribute name.
 463        data_model
 464            Parse GRIB metadata following a defined data model comvention.
 465
 466        Returns
 467        -------
 468        open_dataset
 469            Xarray dataset of grib2 messages.
 470        """
 471        with grib2io.open(filename, _xarray_backend=True) as f:
 472            file_index = pd.DataFrame(f._index)
 473            file_index = file_index.assign(msg=msgs_from_index(f._index))
 474
 475        # parse grib2io _index to dataframe and acquire non-geo possible dims
 476        # (scalar coord when not dim due to squeeze) parse_grib_index applies
 477        # filters to index and expands metadata based on product definition
 478        # template number
 479        file_index, dim_coords, attrs, coord_attrs = parse_grib_index(file_index, filters)
 480
 481        # Divide up records by variable
 482        frames, cube, extra_geo = make_variables(file_index, filename, dim_coords)  # have this return var_attrs
 483
 484        # return empty dataset if no data
 485        if frames is None:
 486            return xr.Dataset()
 487
 488        # create dataframe and add datarrays without any coords
 489        ds = xr.Dataset()
 490        for var_df in frames:
 491            da = build_da_without_coords(var_df, cube, filename, attrs)
 492            ds[da.name] = da
 493
 494        # add coords and dataset meta
 495        ds = assign_xr_meta(ds, frames, cube, dim_coords, extra_geo, coord_attrs)
 496
 497        if data_model is not None:
 498            ds = parse_data_model(ds, data_model)
 499
 500        # assign attributes
 501        ds.attrs['engine'] = 'grib2io'
 502
 503        return ds
 504
 505    def open_datatree(
 506        self,
 507        filename,
 508        *,
 509        drop_variables=None,
 510        filters: typing.Mapping[str, typing.Any] = None,
 511        stack_vertical: bool = False,
 512    ):
 513        """
 514        Open a GRIB2 file as an xarray DataTree.
 515
 516        Parameters
 517        ----------
 518        filename : str
 519            Path to the GRIB2 file.
 520        drop_variables : list, optional
 521            List of variables to exclude.
 522        filters : dict, optional
 523            Filter criteria for GRIB2 messages.
 524        stack_vertical : bool, optional
 525            If True, organize the tree with vertical layers stacked in a single dataset.
 526
 527        Returns
 528        -------
 529        xarray.DataTree
 530            A hierarchical DataTree representation of the GRIB2 data.
 531        """
 532        if not _HAS_DATATREE:
 533            raise ImportError("xarray version does not support DataTree functionality.")
 534
 535        if filters is None:
 536            filters = {}
 537
 538        # Open the file without any filters first to get all messages
 539        with grib2io.open(filename, _xarray_backend=True) as f:
 540            file_index = pd.DataFrame(f._index)
 541            file_index = file_index.assign(msg=msgs_from_index(f._index))
 542
 543        # Build tree structure from GRIB messages with specified options
 544        tree = build_datatree_from_grib(filename, file_index, filters, stack_vertical=stack_vertical)
 545
 546        # Put warning here so it is the last message from likely other Xarray warnings.
 547        warnings.warn(
 548            "grib2io’s xarray backend DataTree support is experimental. "
 549            "The DataTree structure or attributes may change in future releases.",
 550        UserWarning,
 551        stacklevel=2,
 552        )
 553
 554        return tree
 555
 556
 557class GribBackendArray(BackendArray):
 558
 559    def __init__(self, array, lock):
 560        self.array = array
 561        self.shape = array.shape
 562        self.dtype = np.dtype(array.dtype)
 563        self.lock = lock
 564
 565    def __getitem__(self, key: xr.core.indexing.ExplicitIndexer) -> np.typing.ArrayLike:
 566        return xr.core.indexing.explicit_indexing_adapter(
 567            key,
 568            self.shape,
 569            indexing.IndexingSupport.BASIC,
 570            self._raw_getitem,
 571        )
 572
 573    def _raw_getitem(self, key: tuple):
 574        """Implement thread safe access to data on disk."""
 575        with self.lock:
 576            return self.array[key]
 577
 578
 579def exclusive_slice_to_inclusive(item: slice):
 580    """
 581    Convert a slice with exclusive stop to an inclusive slice.
 582
 583    If the slice has a step, the stop is reduced by the step, so that both
 584    interpretations would yield the same result.
 585
 586    The means that [start, stop) is converted to [start, stop - step].
 587
 588    Parameters
 589    ----------
 590    item
 591        The slice to convert.
 592
 593    Returns
 594    -------
 595    slice
 596        The converted slice.
 597    """
 598    # return the None slice
 599    if item.start is None and item.stop is None and item.step is None:
 600        return item
 601    if not isinstance(item, slice):
 602        raise ValueError(f'item must be a slice; it was of type {type(item)}')
 603    # if step is None, it's one
 604    step = 1 if item.step is None else item.step
 605    if item.stop < item.start or step < 1:
 606        raise ValueError(f'slice {item} not accounted for')
 607    # handle case where slice has one item
 608    if abs(item.stop - item.start) == step:
 609        return [item.start]
 610    # other cases require reducing the stop by the step
 611    s = slice(item.start, item.stop - step, step)
 612    return s
 613
 614
 615class Validator:
 616    def __set_name__(self, owner, name):
 617        self.private_name = f'_{name}'
 618        self.name = name
 619
 620    def __get__(self, obj, objtype=None):
 621        try:
 622            value = getattr(obj, self.private_name)
 623        except AttributeError:
 624            value = None
 625        return value
 626
 627
 628class PdIndex(Validator):
 629
 630    def __set__(self, obj, value):
 631        try:
 632            value = pd.Index(value)
 633        except TypeError:
 634            value = pd.Index([value])
 635        setattr(obj, self.private_name, value)
 636
 637
 638def _asarray_tuplesafe(values):
 639    """
 640    Convert values to a numpy array of at most 1-dimension and preserve tuples.
 641
 642    Adapted from pandas.core.common._asarray_tuplesafe
 643    """
 644    if isinstance(values, tuple):
 645        result = np.empty(1, dtype=object)
 646        result[0] = values
 647    else:
 648        result = np.asarray(values)
 649        if result.ndim == 2:
 650            result = np.empty(len(values), dtype=object)
 651            result[:] = values
 652
 653    return result
 654
 655
 656def array_safe_eq(a, b) -> bool:
 657    """Check if a and b are equal, even if they are numpy arrays."""
 658    if a is b:
 659        return True
 660    if hasattr(a, 'equals'):
 661        return a.equals(b)
 662    if hasattr(a, 'all') and hasattr(b, 'all'):
 663        return a.shape == b.shape and (a == b).all()
 664    if hasattr(a, 'all') or hasattr(b, 'all'):
 665        return False
 666    try:
 667        return a == b
 668    except TypeError:
 669        return NotImplementedError
 670
 671
 672def dc_eq(dc1, dc2) -> bool:
 673    """Check if two dataclasses which hold numpy arrays are equal."""
 674    if dc1 is dc2:
 675        return True
 676    if dc1.__class__ is not dc2.__class__:
 677        return NotImplementedError
 678    t1 = astuple(dc1)
 679    t2 = astuple(dc2)
 680    return all(array_safe_eq(a1, a2) for a1, a2 in zip(t1, t2))
 681
 682
 683def coords_from_cube(cube) -> typing.Dict[str, xr.Variable]:
 684    keys = list(cube.keys())
 685    keys.remove('x')
 686    keys.remove('y')
 687    coords = dict()
 688    for k in keys:
 689        if k is not None:
 690            if len(cube[k]) > 1:
 691                coords[k] = xr.Variable(dims=k, data=cube[k], attrs=dict(grib_name=k))
 692            elif len(cube[k]) == 1:
 693                coords[k] = xr.Variable(dims=tuple(), data=cube[k][0], attrs=dict(grib_name=k))
 694    return coords
 695
 696
 697@dataclass
 698class OnDiskArray:
 699    file_name: str
 700    index: pd.DataFrame = field(repr=False)
 701    cube: dict = field(repr=False)
 702    shape: typing.Tuple[int, ...] = field(init=False)
 703    ndim: int = field(init=False)
 704    geo_ndim: int = field(init=False)
 705    dtype = 'float32'
 706
 707    def __post_init__(self):
 708        # multiple grids not allowed so can just use first
 709        geo_shape = (self.index.iloc[0].ny, self.index.iloc[0].nx)
 710
 711        self.geo_shape = geo_shape
 712        self.geo_ndim = len(geo_shape)
 713
 714        if len(self.index) == 1:
 715            self.shape = geo_shape
 716        else:
 717            if self.index.index.nlevels == 1:
 718                self.shape = tuple([len(self.index.index)]) + geo_shape
 719            else:
 720                self.shape = tuple([len(i) for i in self.index.index.levels]) + geo_shape
 721        self.ndim = len(self.shape)
 722
 723        cols = ['msg', 'sectionOffset']
 724        self.index = self.index[cols]
 725
 726    def __getitem__(self, item) -> np.array:
 727        # dimensions not in index are internal to tdlpack records; 2 dims for
 728        # grids; 1 dim for stations
 729
 730        index_slicer = item[:-self.geo_ndim]
 731        # maintain all multindex levels
 732        index_slicer = tuple([[i] if isinstance(i, int) else i for i in index_slicer])
 733
 734        # pandas loc slicing is inclusive, therefore convert slices into
 735        # explicit lists
 736        index_slicer_inclusive = tuple([exclusive_slice_to_inclusive(
 737            i) if isinstance(i, slice) else i for i in index_slicer])
 738
 739        # get records selected by item in new index dataframe
 740        if len(index_slicer_inclusive) == 1:
 741            index = self.index.loc[index_slicer_inclusive]
 742        elif len(index_slicer_inclusive) > 1:
 743            index = self.index.loc[index_slicer_inclusive, :]
 744        else:
 745            index = self.index
 746        index = index.set_index(index.index)
 747
 748        # set miloc to new relative locations in sub array
 749        index['miloc'] = list(
 750            zip(*[index.index.unique(level=dim).get_indexer(index.index.get_level_values(dim)) for dim in index.index.names]))
 751
 752        if len(index_slicer_inclusive) == 1:
 753            array_field_shape = tuple([len(index.index)]) + self.geo_shape
 754        elif len(index_slicer_inclusive) > 1:
 755            array_field_shape = index.index.levshape + self.geo_shape
 756        else:
 757            array_field_shape = self.geo_shape
 758
 759        array_field = np.full(array_field_shape, fill_value=np.nan, dtype="float32")
 760
 761        with open(self.file_name, mode='rb') as filehandle:
 762            for key, row in index.iterrows():
 763
 764                bitmap_offset = None if pd.isna(row['sectionOffset'][6]) else int(row['sectionOffset'][6])
 765                values = _data(filehandle, row.msg, bitmap_offset, row['sectionOffset'][7])
 766
 767                if len(index_slicer_inclusive) >= 1:
 768                    array_field[row.miloc] = values
 769                else:
 770                    array_field = values
 771
 772        # handle geo dim slicing
 773        array_field = array_field[(Ellipsis,) + item[-self.geo_ndim:]]
 774
 775        # squeeze array dimensions expressed as integer
 776        for i, it in reversed(list(enumerate(item[: -self.geo_ndim]))):
 777            if isinstance(it, int):
 778                array_field = array_field[(slice(None, None, None),) * i + (0,)]
 779
 780        return array_field
 781
 782
 783def dims_to_shape(d) -> tuple:
 784    if 'nx' in d:
 785        t = (d['ny'], d['nx'])
 786    else:
 787        t = (d['nsta'],)
 788    return t
 789
 790
 791def filter_index(index, k, v):
 792    if isinstance(v, slice):
 793        index = index.set_index(k)
 794        index = index.loc[v]
 795        index = index.reset_index()
 796    else:
 797        label = (
 798            v
 799            if getattr(v, "ndim", 1) > 1  # vectorized-indexing
 800            else _asarray_tuplesafe(v)
 801        )
 802        if label.ndim == 0:
 803            # see https://github.com/pydata/xarray/pull/4292 for details
 804            label_value = label[()] if label.dtype.kind in "mM" else label.item()
 805            try:
 806                indexer = pd.Index(index[k]).get_loc(label_value)
 807                if isinstance(indexer, int):
 808                    index = index.iloc[[indexer]]
 809                else:
 810                    index = index.iloc[indexer]
 811            except KeyError:
 812                index = index.iloc[[]]
 813        else:
 814            indexer = pd.Index(index[k]).get_indexer_for(np.ravel(v))
 815            index = index.iloc[indexer[indexer >= 0]]
 816
 817    return index
 818
 819
 820def parse_grib_index(
 821    index: pd.DataFrame,
 822    filters: typing.Mapping[str, typing.Any] = dict(),
 823):
 824    """
 825    Apply filters.
 826
 827    Evaluate remaining dimensions based on pdtn and parse each out.
 828
 829    Parameters
 830    ----------
 831    index
 832        Pandas DataFrame containing the GRIB2 message index.
 833    filters
 834        Filter GRIB2 messages to single hypercube. Dict keys can be any
 835        GRIB2 metadata attribute name.
 836
 837    Returns
 838    -------
 839    index
 840        Modified Pandas DataFrame with added GRIB2 metadata columns.
 841    dim_coords
 842        List of GRIB2 attributes that will be used for coordinates and/or dimensions.
 843    attrs
 844        Dict of metadata attributes (non-coordinates, non-geo)
 845    """
 846
 847    # make a copy of filters, remove filters as they are applied
 848    filters = copy(filters)
 849
 850    for k, v in filters.items():
 851        if k not in index.columns:
 852            kwarg = {k: index.msg.apply(lambda msg: getattr(msg, k))}
 853            index = index.assign(**kwarg)
 854        # adopt parts of xarray's sel logic  so that filters behave similarly
 855        # allowed to filter to nothing to make empty dataset
 856        index = filter_index(index, k, v)
 857
 858    if len(index) == 0:
 859        return index, list()
 860
 861    dim_coords = dict()  # key=name of dim, value=list of coord names
 862    attrs = dict()
 863    coord_attrs = dict()
 864
 865    # expand index
 866    index = index.assign(shortName=index.msg.apply(lambda msg: msg.shortName))
 867    index = index.assign(nx=index.msg.apply(lambda msg: msg.nx))
 868    index = index.assign(ny=index.msg.apply(lambda msg: msg.ny))
 869    index = index.astype({'ny': 'int', 'nx': 'int'})
 870
 871    # apply common filters(to all definition templates) to reduce dataset to
 872    # single cube
 873    # ensure only one of each of the below exists after filters applied
 874    required_uniques = [
 875        "productDefinitionTemplateNumber",
 876        "typeOfGeneratingProcess",
 877        "typeOfFirstFixedSurface",
 878        "typeOfSecondFixedSurface",
 879    ]
 880
 881    def meta_check(index, attrs, meta):
 882        """
 883        add meta to the datframe index
 884        check that there is a single type
 885        add the type to attrs
 886
 887        returns index, attrs
 888        """
 889        index = index.assign(**{meta: index.msg.apply(lambda msg: getattr(msg, meta))})
 890
 891        unique = index[meta].unique()
 892        if len(index[meta].unique()) > 1:
 893            raise ValueError(f'filter to a single {meta}; found: {[str(i) for i in unique]}')
 894        value = unique.item()
 895        if type(value) == grib2io.templates.Grib2Metadata:
 896            value = value.definition
 897
 898        # None is returned if no value found,
 899        # check and change to string None
 900        if value is None:
 901            value = 'None'
 902
 903        attrs[meta] = value
 904        return index, attrs
 905
 906    for meta in required_uniques:
 907        index, attrs = meta_check(index, attrs, meta)
 908
 909    pdtn = index.productDefinitionTemplateNumber.iloc[0].value
 910
 911    # determine which non geo dimensions can be created from data by this point
 912    # the index is filtered down to a single type for all required_uniques
 913
 914    # Dim Name     # matching dim_name for using this data as index coordinate
 915    dim_coords["refDate"] = ["refDate"]
 916    coord_attrs["refDate"] = dict(standard_name="forecast_reference_time")
 917#   dim_coords["refDate"] = ["refDate", "hour"] # non dim name matching items in list are used as non-index coordinates
 918
 919    dim_coords["leadTime"] = ["leadTime"]
 920    coord_attrs["leadTime"] = dict(standard_name="forecast_period")
 921
 922    if 'valueOfFirstFixedSurface' not in index.columns:
 923        index = index.assign(valueOfFirstFixedSurface=index.msg.apply(lambda msg: msg.valueOfFirstFixedSurface))
 924    if 'valueOfsecondFixedSurface' not in index.columns:
 925        index = index.assign(valueOfSecondFixedSurface=index.msg.apply(lambda msg: msg.valueOfSecondFixedSurface))
 926
 927    # dim name api change, user could run ds = ds.swap_dims(fixedSurface="valueOfFirstFixedSurface")
 928    index = index.assign(level=list(zip(index['valueOfFirstFixedSurface'], index['valueOfSecondFixedSurface'])))
 929#   index = index.assign(level=index.msg.apply(lambda msg: msg.level))
 930    # lack of "level" indeicates don't create extra index coordinate "level"
 931    dim_coords["level"] = ["valueOfFirstFixedSurface", "valueOfSecondFixedSurface"]
 932
 933    # logic for parsing possible dims from specific product definition section
 934
 935    if pdtn in {5, 9}:
 936
 937        # Probability forecasts at a horizontal level or in a horizontal layer
 938        # in a continuous or non-continuous time interval.  (see Template
 939        # 4.9)
 940        #       AVAILABLE_THRESHOLD = {
 941        #           0: {'has_lower': True, 'has_upper': False},
 942        #           1: {'has_lower': False, 'has_upper': True},
 943        #           2: {'has_lower': True, 'has_upper': True},
 944        #           3: {'has_lower': True, 'has_upper': False},
 945        #           4: {'has_lower': False, 'has_upper': True},
 946        #           5: {'has_lower': True, 'has_upper': False},
 947        #       }
 948
 949        index, attrs = meta_check(index, attrs, "typeOfProbability")
 950        if 'thresholdLowerLimit' not in index.columns:
 951            index = index.assign(thresholdLowerLimit=index.msg.apply(lambda msg: msg.thresholdLowerLimit))
 952        if 'thresholdUpperLimit' not in index.columns:
 953            index = index.assign(thresholdUpperLimit=index.msg.apply(lambda msg: msg.thresholdUpperLimit))
 954        if 'threshold' not in index.columns:
 955            # using composite of lower and upper, but could use threshold string from grib2io as long as that is unique and based on lower and upper
 956            index = index.assign(threshold=list(zip(index['thresholdLowerLimit'], index['thresholdUpperLimit'])))
 957#           index = index.assign(threshold = index.msg.apply(lambda msg: msg.threshold))
 958
 959        # ommiting threshold results in no index being assigned for this possible dim
 960        dim_coords["threshold"] = ["thresholdLowerLimit", "thresholdUpperLimit"]
 961
 962    if pdtn in {6, 10}:
 963
 964        # Percentile forecasts at a horizontal level or in a horizontal layer
 965        # in a continuous or non-continuous time interval.  (see Template
 966        # 4.10)
 967        dim_coords["percentileValue"] = ["percentileValue"]
 968        coord_attrs["percentileValue"] = dict(long_name='percentile', units='percent')
 969
 970    if pdtn in {8, 9, 10, 11, 12, 13, 14, 42, 43, 45, 46, 47, 61, 62, 63, 67, 68, 72, 73, 78, 79, 82, 83, 84, 85, 87, 91}:
 971        dim_coords["duration"] = ["duration"]
 972
 973    if pdtn in {1, 11, 33, 34, 41, 43, 45, 47, 49, 54, 56, 58, 59, 63, 68, 77, 79, 81, 83, 84, 85, 92}:
 974        dim_coords["perturbationNumber"] = ["perturbationNumber"]
 975
 976    if pdtn in {2,3,4,12,13,14}:
 977        index, attrs = meta_check(index, attrs, 'typeOfDerivedForecast')
 978
 979    if pdtn in {8,15,42,46,62,67,72,78,82,1001,1002,1100,1101}:
 980        index, attrs = meta_check(index, attrs, 'statisticalProcess')
 981
 982    # Finish logic by pdtn
 983
 984    for k, v in dim_coords.items():
 985        for meta in v:
 986            if meta not in index.columns:
 987                index = index.assign(**{meta: index.msg.apply(lambda msg: getattr(msg, meta))})
 988
 989    return index, dim_coords, attrs, coord_attrs
 990
 991
 992# Custom open_datatree function to open grib files as DataTree
 993def open_datatree(filename, *, filters: typing.Mapping[str, typing.Any] = None, engine="grib2io"):
 994    """
 995    Open a GRIB2 file as an xarray DataTree.
 996
 997    Parameters
 998    ----------
 999    filename : str
1000        Path to the GRIB2 file.
1001    filters : dict, optional
1002        Filter criteria for GRIB2 messages.
1003    engine : str, optional
1004        Engine to use for opening the file, defaults to "grib2io".
1005
1006    Returns
1007    -------
1008    xarray.DataTree
1009        A hierarchical DataTree representation of the GRIB2 data.
1010    """
1011    if not _HAS_DATATREE:
1012        raise ImportError("xarray version does not support DataTree functionality.")
1013
1014    if filters is None:
1015        filters = {}
1016
1017    # Open the file without any filters first to get all messages
1018    with grib2io.open(filename, _xarray_backend=True) as f:
1019        file_index = pd.DataFrame(f._index)
1020
1021    # Create a DataTree root
1022    tree = xr.DataTree()
1023
1024    # Build tree structure from GRIB messages
1025    return build_datatree_from_grib(filename, file_index, filters)
1026
1027
1028def build_da_without_coords(index, cube, filename, attrs) -> xr.DataArray:
1029    """
1030    Build a DataArray without coordinates from a cube of grib2 messages.
1031
1032    Parameters
1033    ----------
1034    index
1035        Index of cube.
1036    cube
1037        Cube of grib2 messages.
1038    filename
1039        Filename of grib2 file
1040    add_grib_section_attrs
1041        Include grib section arrays as dataArray attributes
1042
1043    Returns
1044    -------
1045    DataArray
1046        DataArray without coordinates
1047    """
1048
1049    dim_names = [k for k in cube.keys() if cube[k] is not None and len(cube[k]) > 1]
1050    constant_meta_names = [k for k in cube.keys() if cube[k] is None]
1051    dims = {k: len(cube[k]) for k in dim_names}
1052
1053    # guard against bad datarrays being formed
1054    dims_total = 1
1055    dims_to_filter = []
1056    for dim_name, dim_len, in dims.items():
1057        if dim_name not in {'x', 'y', 'station'}:
1058            dims_total *= dim_len
1059            dims_to_filter.append(dim_name)
1060
1061    # Check number of GRIB2 message indexed compared to non-X/Y
1062    # dimensions.
1063    if dims_total != len(index):
1064        raise ValueError(
1065            f"DataArray dimensions are not compatible with number of GRIB2 messages; DataArray has {dims_total} "
1066            f"and GRIB2 index has {len(index)}. Consider applying a filter for dimensions: {dims_to_filter}"
1067        )
1068
1069    data = OnDiskArray(filename, index, cube)
1070    lock = _LOCK
1071    data = GribBackendArray(data, lock)
1072    data = indexing.LazilyIndexedArray(data)
1073    if len(dim_names) != len(data.shape):
1074        raise ValueError(
1075            "different number of dimensions on data "
1076            f"and dims: {len(data.shape)} vs {len(dim_names)}\n"
1077            "Grib2 messages could not be formed into a data cube; "
1078            "It's possible extra messages exist along a non-accounted for dimension based on PDTN\n"
1079            "It might be possible to get around this by applying a filter on the non-accounted for dimension"
1080        )
1081    da = xr.DataArray(data, dims=dim_names)
1082
1083    da.encoding['original_shape'] = data.shape
1084
1085    da.encoding['preferred_chunks'] = {'y': -1, 'x': -1}
1086    msg1 = index.msg.iloc[0]
1087
1088    # plain language metadata is minimized
1089    # add grib section metadata
1090    da.attrs['GRIB2IO_section0'] = msg1.section0
1091    da.attrs['GRIB2IO_section1'] = msg1.section1
1092    da.attrs['GRIB2IO_section2'] = msg1.section2 if msg1.section2 else []
1093    da.attrs['GRIB2IO_section3'] = msg1.section3
1094    da.attrs['GRIB2IO_section4'] = msg1.section4
1095    da.attrs['GRIB2IO_section5'] = msg1.section5
1096    da.attrs['fullName'] = str(msg1.fullName)
1097    da.attrs['shortName'] = str(msg1.shortName)
1098    da.attrs['units'] = str(msg1.units)
1099    da.attrs['originatingCenter'] = str(msg1.originatingCenter.definition)
1100    da.attrs['originatingSubCenter'] = str(msg1.originatingSubCenter.definition)
1101
1102    # add master table
1103    da.attrs['masterTableInfo'] = str(msg1.masterTableInfo.definition)
1104
1105    da.name = index.shortName.iloc[0]
1106    for meta_name in constant_meta_names:
1107        if meta_name in index.columns:
1108            da.attrs[meta_name] = index[meta_name].iloc[0]
1109
1110    da.attrs.update(attrs)
1111
1112    return da
1113
1114
1115def assign_xr_meta(ds, frames, cube, non_geo_dims, extra_geo, coord_attrs):
1116
1117    # assign coords from the cube; the cube prevents datarrays with
1118    # different shapes
1119    ds = ds.assign_coords(coords_from_cube(cube))
1120    # assign extra index associated coords
1121    df = frames[0]  # use first variable as they all have same shape and index metadata
1122    for dim_name, coord_names in non_geo_dims.items():
1123        retain_index_coord = False
1124        for name in coord_names:
1125            if name == dim_name:
1126                retain_index_coord = True
1127            else:
1128                if ds[dim_name].size == 1:
1129                    # for assigning scalar coords
1130                    coord_data = [df[name].unique().item()]
1131                    ds = ds.assign_coords({name: coord_data}).squeeze()
1132                else:
1133                    # "ValueError: can only convert an array of size 1 to a Python scalar" indicates the coord is not compatible with the index
1134                    coord_data = [df[df.index.get_level_values(f'{dim_name}_ix') == val][name].unique(
1135                    ).item() for val in range(ds[dim_name].size)]
1136                    coord = pd.Index(coord_data, name=dim_name)
1137                    ds = ds.assign_coords({name: coord})
1138        if not retain_index_coord:
1139            ds = ds.drop_vars(dim_name)
1140
1141    # assign extra geo coords
1142    ds = ds.assign_coords(extra_geo)
1143    # add crs data from first grib message to each data variable and the dataset
1144    geo_attrs = {
1145        'crs_wkt': CRS.from_dict(df.msg.iloc[0].projParameters).to_wkt(),
1146        'gridlengthXDirection': df.msg.iloc[0].gridlengthXDirection,
1147        'gridlengthYDirection': df.msg.iloc[0].gridlengthYDirection,
1148        'latitudeFirstGridpoint': df.msg.iloc[0].latitudeFirstGridpoint,
1149        'longitudeFirstGridpoint': df.msg.iloc[0].longitudeFirstGridpoint,
1150    }
1151    for data_var in ds.data_vars:
1152        ds[data_var].attrs.update(geo_attrs)
1153    ds.attrs.update(geo_attrs)
1154
1155    # add coordinate specific attributes
1156    for coord, attrs in coord_attrs.items():
1157        ds[coord].attrs.update(attrs)
1158
1159    # assign valid date coords
1160    ds = ds.assign_coords(dict(validDate=ds.coords['refDate']+ds.coords['leadTime']))
1161    ds.validDate.attrs['standard_name'] = 'time'
1162    ds.validDate.attrs['long_name'] = 'time'
1163
1164    # assign attributes
1165    ds.attrs['engine'] = 'grib2io'
1166
1167    return ds
1168
1169
1170def make_variables(index, f, non_geo_dims, allow_uneven_dims=False):
1171    """
1172    Create an individual dataframe index and cube for each variable.
1173
1174    Parameters
1175    ----------
1176    index
1177        Index of cube.
1178    f
1179        ?
1180    non_geo_dims
1181        Dimensions not associated with the x,y grid
1182    allow_uneven_dims
1183        If True, allows uneven dimensions (used for DataTree creation)
1184
1185    Returns
1186    -------
1187    ordered_frames
1188        List of dataframes, one for each variable.
1189    cube
1190        Cube of grib2 messages.
1191    extra_geo
1192        Extra geographic coordinates.
1193    """
1194    # let shortName determine the variables
1195
1196    # set the index to the name
1197    index = index.set_index('shortName').sort_index()
1198    # return nothing if no data
1199    if index.empty:
1200        return None, None, None
1201
1202    # define the DimCube
1203    dims = copy(non_geo_dims)
1204
1205    ordered_meta = list(non_geo_dims.keys())
1206    cube = None
1207    ordered_frames = list()
1208    for key in index.index.unique():
1209        frame = index.loc[[key]]
1210        frame = frame.reset_index()
1211        # frame is a dataframe with all records for one variable
1212        c = dict()
1213        # for colname in frame.columns:
1214        for colname in ordered_meta:
1215            uniques = pd.Index(frame[colname]).unique()
1216            if len(uniques) > 1:
1217                c[colname] = uniques.sort_values()
1218            else:
1219                c[colname] = [uniques[0]]
1220
1221        dims = [k for k in ordered_meta if len(c[k]) > 1]
1222
1223        for dim in dims:
1224            if frame[dim].value_counts().nunique() > 1 and not allow_uneven_dims:
1225                raise ValueError(
1226                    f'uneven number of grib msgs associated with dimension: {dim}\n unique values for {dim}: {frame[dim].unique()} ')
1227
1228        if len(dims) >= 1:  # dims may be empty if no extra dims on top of x,y
1229            frame = frame.sort_values(dims)
1230            frame = frame.set_index(dims)
1231
1232        if cube:
1233            if cube != c and not allow_uneven_dims:
1234                raise ValueError(f'{cube},\n {c};\n cubes are not the same; filter to a single cube')
1235        else:
1236            cube = c
1237
1238        # miloc is multi-index integer location of msg in nd DataArray
1239        miloc = list(zip(*[frame.index.unique(level=dim).get_indexer(frame.index.get_level_values(dim))
1240                     for dim in dims]))
1241
1242        # set frame multi index
1243        if len(miloc) >= 1:  # miloc will be empty when no extra dims, thus no multiindex
1244            dim_ix = tuple([n+'_ix' for n in dims])
1245            frame = frame.set_index(pd.MultiIndex.from_tuples(miloc, names=dim_ix))
1246
1247        ordered_frames.append(frame)
1248
1249    # no variables
1250    if cube is None:
1251        cube = dict()
1252
1253    # check geography of data and assign to cube
1254    if len(index.ny.unique()) > 1 or len(index.nx.unique()) > 1:
1255        raise ValueError('multiple grids not accommodated')
1256    cube["y"] = range(int(index.ny.iloc[0]))
1257    cube["x"] = range(int(index.nx.iloc[0]))
1258
1259    extra_geo = None
1260    msg = index.msg.iloc[0]
1261
1262    # we want the lat lons; make them via accessing a record; we are assuming
1263    # all records are the same grid because they have the same shape;
1264    # may want a unique grid identifier from grib2io to avoid assuming this
1265    latitude, longitude = msg.latlons()
1266    latitude = xr.DataArray(latitude, dims=['y', 'x'])
1267    latitude.attrs['standard_name'] = 'latitude'
1268    latitude.attrs['units'] = 'degrees_north'
1269    longitude = xr.DataArray(longitude, dims=['y', 'x'])
1270    longitude.attrs['standard_name'] = 'longitude'
1271    longitude.attrs['units'] = 'degrees_east'
1272    extra_geo = dict(latitude=latitude, longitude=longitude)
1273
1274    return ordered_frames, cube, extra_geo
1275
1276
1277def interp_nd(a, *, method, grid_def_in, grid_def_out, method_options=None, num_threads=1):
1278    front_shape = a.shape[:-2]
1279    a = a.reshape(-1, a.shape[-2], a.shape[-1])
1280    a = grib2io.interpolate(a, method, grid_def_in, grid_def_out, method_options=method_options,
1281                            num_threads=num_threads)
1282    a = a.reshape(front_shape + (a.shape[-2], a.shape[-1]))
1283    return a
1284
1285
1286def interp_nd_stations(a, *, method, grid_def_in, lats, lons, method_options=None, num_threads=1):
1287    front_shape = a.shape[:-2]
1288    a = a.reshape(-1, a.shape[-2], a.shape[-1])
1289    a = grib2io.interpolate_to_stations(a, method, grid_def_in, lats, lons, method_options=method_options,
1290                                        num_threads=num_threads)
1291    a = a.reshape(front_shape + (len(lats),))
1292    return a
1293
1294
1295@xr.register_dataset_accessor("grib2io")
1296class Grib2ioDataSet:
1297
1298    def __init__(self, xarray_obj):
1299        self._obj = xarray_obj
1300
1301    def griddef(self):
1302        return Grib2GridDef.from_section3(self._obj[list(self._obj.data_vars)[0]].attrs['GRIB2IO_section3'])
1303
1304    def interp(self, method, grid_def_out, method_options=None, num_threads=1) -> xr.Dataset:
1305        # see interp method of class Grib2ioDataArray
1306        da = self._obj.to_array()
1307        da.attrs['GRIB2IO_section3'] = self._obj[list(self._obj.data_vars)[0]].attrs['GRIB2IO_section3']
1308        da = da.grib2io.interp(method, grid_def_out, method_options=method_options,
1309                               num_threads=num_threads)
1310        ds = da.to_dataset(dim='variable')
1311        return ds
1312
1313    def interp_to_stations(self, method, calls, lats, lons, method_options=None, num_threads=1) -> xr.Dataset:
1314        # see interp_to_stations method of class Grib2ioDataArray
1315        da = self._obj.to_array()
1316        da.attrs['GRIB2IO_section3'] = self._obj[list(self._obj.data_vars)[0]].attrs['GRIB2IO_section3']
1317        da = da.grib2io.interp_to_stations(method, calls, lats, lons, method_options=method_options,
1318                                           num_threads=num_threads)
1319        ds = da.to_dataset(dim='variable')
1320        return ds
1321
1322    def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
1323        """
1324        Write a DataSet to a grib2 file.
1325
1326        Parameters
1327        ----------
1328        filename
1329            Name of the grib2 file to write to.
1330        mode: {"x", "w", "a"}, optional, default="x"
1331            Persistence mode
1332
1333            | mode | Description                       |
1334            | :---:| :---:                             |
1335            | 'x'  | create (fail if exists)           |
1336            | 'w'  | create (overwrite if exists)      |
1337            | 'a'  | append (create if does not exist) |
1338
1339        """
1340        ds = self._obj
1341
1342        for shortName in sorted(ds):
1343            # make a DataArray from the "Data Variables" in the DataSet
1344            da = ds[shortName]
1345
1346            da.grib2io.to_grib2(filename, mode=mode)
1347            mode = "a"
1348
1349    def update_attrs(self, **kwargs):
1350        """
1351        Raises an error because Datasets don't have a .attrs attribute.
1352
1353        Parameters
1354        ----------
1355        attrs
1356            Attributes to update.
1357        """
1358        raise ValueError(
1359            f"Datasets do not have a .attrs attribute; use .grib2io.update_attrs({kwargs}) on a DataArray instead."
1360        )
1361
1362    def subset(self, lats, lons) -> xr.Dataset:
1363        """
1364        Subset the DataSet to a region defined by latitudes and longitudes.
1365
1366        Parameters
1367        ----------
1368        lats
1369            Latitude bounds of the region.
1370        lons
1371            Longitude bounds of the region.
1372
1373        Returns
1374        -------
1375        subset
1376            DataSet subset to the region.
1377        """
1378        ds = self._obj
1379
1380        newds = xr.Dataset()
1381        for shortName in ds:
1382            newds[shortName] = ds[shortName].grib2io.subset(lats, lons).copy()
1383
1384        return newds
1385
1386
1387@xr.register_dataarray_accessor("grib2io")
1388class Grib2ioDataArray:
1389
1390    def __init__(self, xarray_obj):
1391        self._obj = xarray_obj
1392
1393    def griddef(self):
1394        return Grib2GridDef.from_section3(self._obj.attrs['GRIB2IO_section3'])
1395
1396    def interp(self, method, grid_def_out, method_options=None, num_threads=1) -> xr.DataArray:
1397        """
1398        Perform grid spatial interpolation.
1399
1400        Uses the [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip).
1401
1402        Parameters
1403        ----------
1404        method
1405            Interpolate method to use. This can either be an integer or string
1406            using the following mapping:
1407
1408            | Interpolate Scheme | Integer Value |
1409            | :---:              | :---:         |
1410            | 'bilinear'         | 0             |
1411            | 'bicubic'          | 1             |
1412            | 'neighbor'         | 2             |
1413            | 'budget'           | 3             |
1414            | 'spectral'         | 4             |
1415            | 'neighbor-budget'  | 6             |
1416        grid_def_out
1417            Grib2GridDef object of the output grid.
1418        method_options : list of ints, optional
1419            Interpolation options. See the NCEPLIBS-ip documentation for
1420            more information on how these are used.
1421        num_threads : int, optional
1422            Number of OpenMP threads to use for interpolation. The default
1423            value is 1. If grib2io_interp was not built with OpenMP, then
1424            this keyword argument and value will have no impact.
1425
1426        Returns
1427        -------
1428        interp
1429            DataSet interpolated to new grid definition.  The attribute
1430            GRIB2IO_section3 is replaced with the section3 array from the new
1431            grid definition.
1432        """
1433        da = self._obj
1434        # ensure that y, x are rightmost dims; they should be if opening with
1435        # grib2io engine
1436
1437        # gdtn and gdt is not the entirety of the new s3
1438        npoints = grid_def_out.npoints
1439        s3_new = np.array([0, npoints, 0, 0, grid_def_out.gdtn] + list(grid_def_out.gdt))
1440
1441        # make new lat lons
1442        lats, lons = Grib2Message(section3=s3_new, pdtn=0, drtn=0).grid()
1443        latitude = xr.DataArray(lats, dims=['y', 'x'])
1444        longitude = xr.DataArray(lons, dims=['y', 'x'])
1445
1446        # create new coords
1447        new_coords = dict(da.coords)
1448        del new_coords['latitude']
1449        del new_coords['longitude']
1450        new_coords['longitude'] = longitude
1451        new_coords['latitude'] = latitude
1452
1453        # make grid def in from section3 on da.attrs
1454        grid_def_in = self.griddef()
1455
1456        if da.chunks is None:
1457            data = interp_nd(da.data, method=method, grid_def_in=grid_def_in,
1458                             grid_def_out=grid_def_out,
1459                             method_options=method_options, num_threads=num_threads)
1460        else:
1461            import dask
1462            front_shape = da.shape[:-2]
1463            data = da.data.map_blocks(interp_nd, method=method, grid_def_in=grid_def_in,
1464                                      grid_def_out=grid_def_out, method_options=method_options,
1465                                      chunks=da.chunks[:-2]+latitude.shape, dtype=da.dtype)
1466
1467        new_da = xr.DataArray(data, dims=da.dims, coords=new_coords, attrs=da.attrs)
1468
1469        new_da.attrs['GRIB2IO_section3'] = s3_new
1470        new_da.name = da.name
1471        return new_da
1472
1473    def interp_to_stations(self, method, calls, lats, lons, method_options=None, num_threads=1) -> xr.DataArray:
1474        """
1475        Perform spatial interpolation to station points.
1476
1477        Parameters
1478        ----------
1479        method
1480            Interpolate method to use. This can either be an integer or string
1481            using the following mapping:
1482
1483            | Interpolate Scheme | Integer Value |
1484            | :---:              | :---:         |
1485            | 'bilinear'         | 0             |
1486            | 'bicubic'          | 1             |
1487            | 'neighbor'         | 2             |
1488            | 'budget'           | 3             |
1489            | 'spectral'         | 4             |
1490            | 'neighbor-budget'  | 6             |
1491
1492        calls
1493            Station calls used for labeling new station index coordinate
1494        lats
1495            Latitudes of the station points.
1496        lons
1497            Longitudes of the station points.
1498
1499        Returns
1500        -------
1501        interp_to_stations
1502            DataArray interpolated to lat and lon locations and labeled with
1503            dimension and coordinate 'station'. (..., y, x) -> (..., station)
1504        """
1505        da = self._obj
1506        # TODO ensure that y, x are rightmost dims; they should be if opening
1507        # with grib2io engine
1508
1509        calls = np.asarray(calls)
1510        lats = np.asarray(lats)
1511        lons = np.asarray(lons)
1512        latitude = xr.DataArray(lats, dims=['station'])
1513        longitude = xr.DataArray(lons, dims=['station'])
1514
1515        # create new coords
1516        new_coords = dict(da.coords)
1517        del new_coords['latitude']
1518        del new_coords['longitude']
1519        new_coords['longitude'] = longitude
1520        new_coords['latitude'] = latitude
1521        new_coords['station'] = calls
1522
1523        new_dims = da.dims[:-2] + ('station',)
1524
1525        # make grid def in from section3 on da attrs
1526        grid_def_in = self.griddef()
1527
1528        if da.chunks is None:
1529            data = interp_nd_stations(da.data, method=method, grid_def_in=grid_def_in, lats=lats,
1530                                      lons=lons, method_options=method_options, num_threads=num_threads)
1531        else:
1532            import dask
1533            front_shape = da.shape[:-1]
1534            data = da.data.map_blocks(interp_nd_stations, method=method, grid_def_in=grid_def_in,
1535                                      lats=lats, lons=lons, method_options=method_options,
1536                                      drop_axis=-1, chunks=da.chunks[:-2]+latitude.shape,
1537                                      dtype=da.dtype)
1538
1539        new_da = xr.DataArray(data, dims=new_dims, coords=new_coords, attrs=da.attrs)
1540
1541        new_da.name = da.name
1542        return new_da
1543
1544    def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
1545        """
1546        Write a DataArray to a grib2 file.
1547
1548        Parameters
1549        ----------
1550        filename
1551            Name of the grib2 file to write to.
1552        mode: {"x", "w", "a"}, optional, default="x"
1553            Persistence mode
1554
1555            +------+-----------------------------------+
1556            | mode | Description                       |
1557            +======+===================================+
1558            | x    | create (fail if exists)           |
1559            +------+-----------------------------------+
1560            | w    | create (overwrite if exists)      |
1561            +------+-----------------------------------+
1562            | a    | append (create if does not exist) |
1563            +------+-----------------------------------+
1564
1565        """
1566        da = self._obj.copy(deep=True)
1567
1568        coords_keys = sorted(da.coords.keys())
1569        coords_keys = [k for k in coords_keys if k in AVAILABLE_NON_GEO_COORDS]
1570
1571        # If there are dimension coordinates, the DataArray is a hypercube of
1572        # grib2 messages.
1573
1574        # Create `indexes` which is a list of lists of dictionaries for all
1575        # dimension coordinates. Each dictionary key is the dimension
1576        # coordinate name and the value is a list of the dimension coordinate
1577        # values.  This allows for easy iteration over all possible grib2
1578        # messages in the DataArray by using itertools.product.
1579        #
1580        # For example:
1581        # indexes = [
1582        #     [
1583        #         {"leadTime": 9},
1584        #         {"leadTime": 12},
1585        #     ],
1586        #     [
1587        #         {"valueOfFirstFixedSurface": 900},
1588        #         {"valueOfFirstFixedSurface": 925},
1589        #         {"valueOfFirstFixedSurface": 950},
1590        #     ],
1591        # ]
1592
1593        # assign loc indexes to dimensions without indexes for uniform selection by name
1594        loc_indexes = list()
1595        for dim in da.dims:
1596            if dim not in da.indexes:
1597                da = da.assign_coords({dim: range(da[dim].size)})
1598                loc_indexes.append(dim)
1599
1600        indexes = []
1601        for index in [i for i in AVAILABLE_NON_GEO_DIMS if i in da.dims]:
1602            values = da.coords[index].values
1603            if len(values) != len(set(values)):
1604                raise ValueError(
1605                    f"Dimension coordinate '{index}' has duplicate values, but to_grib2 requires unique values to find each GRIB2 message in the DataArray."
1606                )
1607            listeach = [{index: value} for value in sorted(values)]
1608            indexes.append(listeach)
1609
1610        # If `dim_coords` is [], then the DataArray is a single grib2 message and
1611        # itertools.product(*dim_coords) will run once with `selectors = ()`.
1612        for selectors in itertools.product(*indexes):
1613            # Need to find the correct data in the DataArray based on the
1614            # dimension coordinates.
1615            filters = {k: v for d in selectors for k, v in d.items()}
1616
1617            # If `filters` is {}, then the DataArray is a single grib2 message
1618            # and da.sel(indexers={}) returns the DataArray.
1619            selected = da.sel(indexers=filters)
1620
1621            newmsg = Grib2Message(
1622                selected.attrs["GRIB2IO_section0"],
1623                selected.attrs["GRIB2IO_section1"],
1624                selected.attrs["GRIB2IO_section2"],
1625                selected.attrs["GRIB2IO_section3"],
1626                selected.attrs["GRIB2IO_section4"],
1627                selected.attrs["GRIB2IO_section5"],
1628            )
1629            newmsg.data = np.array(selected.data)
1630
1631            # For dimension coordinates, set the grib2 message metadata to the
1632            # dimension coordinate value.
1633            for index, value in filters.items():
1634                if index not in loc_indexes:
1635                    setattr(newmsg, index, value)
1636
1637            # For non-dimension coordinates, set the grib2 message metadata to
1638            # the DataArray coordinate value.
1639            for index in [i for i in coords_keys if i not in da.dims]:
1640                setattr(newmsg, index, selected.coords[index].values)
1641
1642            # Set section 5 attributes to the da.encoding dictionary.
1643            for key, value in selected.encoding.items():
1644                if key in ["dtype", "chunks", "original_shape"]:
1645                    continue
1646                setattr(newmsg, key, value)
1647
1648            # write the message to file
1649            with grib2io.open(filename, mode=mode) as f:
1650                f.write(newmsg)
1651            mode = "a"
1652
1653    def update_attrs(self, **kwargs):
1654        """
1655        Update many of the attributes of the DataArray.
1656
1657        Parameters
1658        ----------
1659        **kwargs
1660            Attributes to update.  This can include many of the GRIB2IO message
1661            attributes that you can find when you print a GRIB2IO message. For
1662            conflicting updates, the last keyword will be used.
1663
1664            +-----------------------+------------------------------------------+
1665            | kwargs                | Description                              |
1666            +=======================+==========================================+
1667            | shortName="VTMP"      | Set shortName to "VTMP", along with      |
1668            |                       | appropriate discipline,                  |
1669            |                       | parameterCategory, parameterNumber,      |
1670            |                       | fullName and units.                      |
1671            +-----------------------+------------------------------------------+
1672            | discipline=0,         | Set shortName, discipline,               |
1673            | parameterCategory=0,  | parameterCategory, parameterNumber,      |
1674            | parameterNumber=1     | fullName and units appropriate for       |
1675            |                       | "Virtual Temperature".                   |
1676            +-----------------------+------------------------------------------+
1677            | discipline=0,         | Conflicting keywords but                 |
1678            | parameterCategory=0,  | 'shortName="TMP"' wins.  Set shortName,  |
1679            | parameterNumber=1,    | discipline, parameterCategory,           |
1680            | shortName="TMP"       | parameterNumber, fullName and units      |
1681            |                       | appropriate for "Temperature".           |
1682            +-----------------------+------------------------------------------+
1683
1684        Returns
1685        -------
1686        DataArray
1687            DataArray with updated attributes.
1688        """
1689        da = self._obj.copy(deep=True)
1690
1691        newmsg = Grib2Message(
1692            da.attrs["GRIB2IO_section0"],
1693            da.attrs["GRIB2IO_section1"],
1694            da.attrs["GRIB2IO_section2"],
1695            da.attrs["GRIB2IO_section3"],
1696            da.attrs["GRIB2IO_section4"],
1697            da.attrs["GRIB2IO_section5"],
1698        )
1699
1700        coords_keys = [
1701            k
1702            for k in da.coords.keys()
1703            if k in AVAILABLE_NON_GEO_COORDS
1704        ]
1705
1706        for grib2_name, value in kwargs.items():
1707            if grib2_name == "gridDefinitionTemplateNumber":
1708                raise ValueError(
1709                    "The gridDefinitionTemplateNumber attribute cannot be updated.  The best way to change to a different grid is to interpolate the data to a new grid using the grib2io interpolate functions."
1710                )
1711            if grib2_name == "productDefinitionTemplateNumber":
1712                raise ValueError(
1713                    "The productDefinitionTemplateNumber attribute cannot be updated."
1714                )
1715            if grib2_name == "dataRepresentationTemplateNumber":
1716                raise ValueError(
1717                    "The dataRepresentationTemplateNumber attribute cannot be updated."
1718                )
1719            if grib2_name in coords_keys:
1720                warnings.warn(
1721                    f"Skipping attribute '{grib2_name}' because it is a coordinate. Use da.assign_coords() to change coordinate values."
1722                )
1723                continue
1724            if hasattr(newmsg, grib2_name):
1725                setattr(newmsg, grib2_name, value)
1726            else:
1727                warnings.warn(
1728                    f"Skipping attribute '{grib2_name}' because it is not a valid GRIB2 attribute for this message and cannot be updated."
1729                )
1730                continue
1731
1732        da.attrs["GRIB2IO_section0"] = newmsg.section0
1733        da.attrs["GRIB2IO_section1"] = newmsg.section1
1734        da.attrs["GRIB2IO_section2"] = newmsg.section2 or []
1735        da.attrs["GRIB2IO_section3"] = newmsg.section3
1736        da.attrs["GRIB2IO_section4"] = newmsg.section4
1737        da.attrs["GRIB2IO_section5"] = newmsg.section5
1738        da.attrs["fullName"] = newmsg.fullName
1739        da.attrs["shortName"] = newmsg.shortName
1740        da.attrs["units"] = newmsg.units
1741
1742        return da
1743
1744    def subset(self, lats, lons) -> xr.DataArray:
1745        """
1746        Subset the DataArray to a region defined by latitudes and longitudes.
1747
1748        Parameters
1749        ----------
1750        lats
1751            Latitude bounds of the region.
1752        lons
1753            Longitude bounds of the region.
1754
1755        Returns
1756        -------
1757        subset
1758            DataArray subset to the region.
1759        """
1760        da = self._obj.copy(deep=True)
1761
1762        newmsg = Grib2Message(
1763            da.attrs["GRIB2IO_section0"],
1764            da.attrs["GRIB2IO_section1"],
1765            da.attrs["GRIB2IO_section2"],
1766            da.attrs["GRIB2IO_section3"],
1767            da.attrs["GRIB2IO_section4"],
1768            da.attrs["GRIB2IO_section5"],
1769        )
1770
1771        newmsg.data = np.zeros((newmsg.ny, newmsg.nx), dtype=np.float32)
1772
1773        newmsg = newmsg.subset(lats, lons)
1774
1775        da.attrs["GRIB2IO_section3"] = newmsg.section3
1776
1777        mask_lat = (da.latitude >= newmsg.latitudeLastGridpoint) & (
1778            da.latitude <= newmsg.latitudeFirstGridpoint
1779        )
1780        mask_lon = (da.longitude >= newmsg.longitudeFirstGridpoint) & (
1781            da.longitude <= newmsg.longitudeLastGridpoint
1782        )
1783
1784        del newmsg
1785
1786        return da.where((mask_lon & mask_lat).compute(), drop=True)
1787
1788
1789def build_datatree_from_grib(filename, file_index, filters=None, stack_vertical=False):
1790    """
1791    Build a DataTree from GRIB2 messages.
1792
1793    Parameters
1794    ----------
1795    filename : str
1796        Path to the GRIB2 file.
1797    file_index : pandas.DataFrame
1798        DataFrame of GRIB2 message index.
1799    filters : dict, optional
1800        Filter criteria for GRIB2 messages.
1801    stack_vertical : bool, optional
1802        If True, vertical levels will be stacked in a single dataset
1803        instead of being organized in separate tree nodes.
1804
1805    Returns
1806    -------
1807    xarray.DataTree
1808        A hierarchical DataTree representation of the GRIB2 data.
1809    """
1810    if filters is None:
1811        filters = {}
1812
1813    # Apply any filters from user
1814    for k, v in filters.items():
1815        if k not in file_index.columns:
1816            file_index = file_index.copy()
1817            file_index[k] = file_index.msg.apply(lambda msg: getattr(msg, k, None))
1818        file_index = filter_index(file_index, k, v)
1819
1820    # Make a copy to avoid the SettingWithCopyWarning
1821    file_index = file_index.copy()
1822
1823    # Extract metadata needed for tree organization
1824    # Use a safer approach to handle missing attributes
1825    def safe_getattr(obj, name):
1826        try:
1827            attr = getattr(obj, name)
1828            # Need to test if the attribute is Grib2Metadata. If so,
1829            # then get the value attribute.
1830            if isinstance(attr, grib2io.templates.Grib2Metadata):
1831                attr = attr.value
1832            return attr
1833        except (AttributeError, KeyError):
1834            return None
1835
1836    for attr in _TREE_HIERARCHY_LEVELS:
1837        if (attr not in file_index.columns) and (attr != 'valueOfFirstFixedSurface'):
1838            file_index[attr] = file_index.msg.apply(lambda msg: safe_getattr(msg, attr))
1839
1840    # Also extract shortName for variable naming
1841    if 'shortName' not in file_index.columns:
1842        file_index = file_index.assign(shortName=file_index.msg.apply(lambda msg: getattr(msg, 'shortName', None)))
1843        file_index = file_index.assign(nx=file_index.msg.apply(lambda msg: getattr(msg, 'nx', None)))
1844        file_index = file_index.assign(ny=file_index.msg.apply(lambda msg: getattr(msg, 'ny', None)))
1845
1846    # Create root DataTree
1847    root = xr.DataTree()
1848
1849    # Adjust hierarchy levels if we're stacking vertical levels
1850    hierarchy_levels = list(_TREE_HIERARCHY_LEVELS) # This makes a copy
1851    if stack_vertical and "valueOfFirstFixedSurface" in hierarchy_levels:
1852        hierarchy_levels.remove("valueOfFirstFixedSurface")
1853
1854    # First group by level type
1855    level_groups = {}
1856
1857    # Create a dictionary to group data by level type
1858    for level_type in file_index['typeOfFirstFixedSurface'].unique():
1859        if pd.notna(level_type):  # Skip None/NaN values
1860            level_info = _LEVEL_NAME_MAPPING.get(level_type, f"level_{level_type}")
1861            level_name = level_info[0]
1862            level_source = level_info[1]
1863            # Get all rows for this level type
1864            level_data = file_index[file_index['typeOfFirstFixedSurface'] == level_type]
1865            level_groups[level_type] = {'name': level_name, 'data': level_data}
1866
1867    # Process each level group
1868    for level_type, group_info in level_groups.items():
1869        level_name = group_info['name']
1870        level_df = group_info['data']
1871
1872        # Create a branch for this level type
1873        level_tree = xr.DataTree()
1874
1875        # Process this branch based on PDTN, perturbation number, etc.
1876        process_level_branch(level_tree, level_df, filename)
1877
1878        # Add this branch to the main tree
1879        root[level_name] = level_tree
1880
1881    return root
1882
1883
1884def process_level_branch(level_tree, df, filename):
1885    """
1886    Process a level type branch of the data tree, organizing by PDTN and other attributes.
1887
1888    Parameters
1889    ----------
1890    level_tree : xarray.DataTree
1891        The DataTree node for this level type
1892    df : pandas.DataFrame
1893        DataFrame of messages for this level type
1894    filename : str
1895        Path to the GRIB2 file
1896    """
1897    # Group by PDTN
1898    pdtn_groups = {}
1899
1900    # Group data by PDTN first
1901    for pdtn_value in df['productDefinitionTemplateNumber'].unique():
1902        if pd.notna(pdtn_value):
1903            pdtn_df = df[df['productDefinitionTemplateNumber'] == pdtn_value]
1904            pdtn_groups[pdtn_value] = pdtn_df
1905
1906    # If there's only one PDTN value, skip creating PDTN branch level
1907    if len(pdtn_groups) == 1:
1908        pdtn, pdtn_df = next(iter(pdtn_groups.items()))
1909
1910        pdtn_name = f"pdtn_{int(pdtn)}"
1911
1912        # Check if we need to further subdivide by perturbation number
1913        has_perturbations = ('perturbationNumber' in pdtn_df.columns and
1914                             len(pdtn_df['perturbationNumber'].dropna().unique()) > 1)
1915
1916        # Check if we need to further subdivide by probabilities unique for each variable.
1917        has_probabilities = ('typeOfProbability' in pdtn_df.columns and
1918                             len(pdtn_df['typeOfProbability'].dropna().unique()) > 1)
1919
1920        if has_perturbations:
1921            # Process perturbations directly on the level tree
1922            process_perturbation_groups(level_tree, pdtn_df, filename)
1923        elif has_probabilities:
1924            # Process probability groups
1925            process_probability_groups(level_tree, pdtn_df, filename)
1926        else:
1927            # Try to create dataset directly on level
1928            try:
1929                dss = create_datasets_from_df(pdtn_df, filename)
1930                if dss is not None:
1931                    dt = xr.DataTree()
1932                    if len(dss) == 1:
1933                        dt.ds = dss[0]
1934                    else:
1935                        for ds in dss:
1936                            varname = list(ds.data_vars)[0]
1937                            dt[f"var_{varname}"] = ds
1938                    level_tree[pdtn_name] = dt
1939            except Exception as e:
1940                print(f"Error creating dataset for level with pdtn {int(pdtn)}: {e}")
1941
1942                # Try to separate by variable name as a fallback
1943                try_process_by_variables(level_tree, pdtn_df, filename)
1944    else:
1945        # Multiple PDTN values, process each group with PDTN branch nodes
1946        for pdtn, pdtn_df in pdtn_groups.items():
1947            # Use a simple node name that's easy to use in code
1948            pdtn_name = f"pdtn_{int(pdtn)}"
1949
1950            # Check if we need to further subdivide by perturbation number
1951            has_perturbations = ('perturbationNumber' in pdtn_df.columns and
1952                                 len(pdtn_df['perturbationNumber'].dropna().unique()) > 1)
1953
1954            # Check if we need to further subdivide by probabilities unique for each variable.
1955            has_probabilities = ('typeOfProbability' in pdtn_df.columns and
1956                                 len(pdtn_df['typeOfProbability'].dropna().unique()) > 1)
1957
1958            if has_perturbations:
1959                # Create a branch for this PDTN
1960                pdtn_tree = xr.DataTree()
1961
1962                # Process perturbation groups
1963                process_perturbation_groups(pdtn_tree, pdtn_df, filename)
1964
1965                # Only add the PDTN branch if it has children
1966                if len(pdtn_tree.children) > 0 or pdtn_tree.ds is not None:
1967                    level_tree[pdtn_name] = pdtn_tree
1968            elif has_probabilities:
1969                # Create a branch for this PDTN
1970                pdtn_tree = xr.DataTree()
1971
1972                # Process probability groups
1973                process_probability_groups(pdtn_tree, pdtn_df, filename)
1974
1975                # Only add the PDTN branch if it has children
1976                if len(pdtn_tree.children) > 0 or pdtn_tree.ds is not None:
1977                    level_tree[pdtn_name] = pdtn_tree
1978            else:
1979                # Create a subtree for this PDTN
1980                pdtn_tree = xr.DataTree()
1981
1982                # Try to create dataset directly on level
1983                try:
1984                    dss = create_datasets_from_df(pdtn_df, filename)
1985                    if dss is not None:
1986                        if len(dss) == 1:
1987                            pdtn_tree.ds = dss[0]
1988                        else:
1989                            for ds in dss:
1990                                varname = list(ds.data_vars)[0]
1991                                pdtn_tree[f"var_{varname}"] = ds
1992                        level_tree[pdtn_name] = pdtn_tree
1993                except Exception as e:
1994                    print(f"Error creating dataset for level with pdtn {int(pdtn)}: {e}")
1995
1996                    # Try to separate by variable name as a fallback
1997                    try_process_by_variables(pdtn_tree, pdtn_df, filename)
1998                    level_tree[pdtn_name] = pdtn_tree
1999
2000
2001def process_probability_groups(target_tree, pdtn_df, filename):
2002    """
2003    """
2004    success = False
2005    # Group by type of probability
2006    prob_groups = {}
2007    for prob_value in pdtn_df['typeOfProbability'].unique():
2008        if pd.notna(prob_value):
2009            prob_df = pdtn_df[pdtn_df['typeOfProbability'] == prob_value]
2010            prob_groups[prob_value] = prob_df
2011
2012    # Process each probability group
2013    prob_dict = {}
2014    for prob_num, prob_df in prob_groups.items():
2015        prob_name = f"prob_{int(prob_num)}"
2016
2017        # Try to create dataset for this probability group
2018        try:
2019            dss = create_datasets_from_df(prob_df, filename)
2020            dt = xr.DataTree()
2021            if len(dss) == 1:
2022                dt.ds = dss[0]
2023                target_tree[prob_name] = dt
2024            elif len(dss) > 1:
2025                for ds in dss:
2026                    dt[f"var_{ds.data_vars[0]}"] = ds
2027            target_tree[prob_name] = dt
2028        except Exception as e:
2029            # Log error but continue processing other groups
2030            print(f"Error creating dataset for type of probability {prob_name}: {e}")
2031
2032    return success
2033
2034
2035def process_perturbation_groups(target_tree, pdtn_df, filename):
2036    """
2037    Process perturbation groups and add them to the target tree.
2038
2039    Parameters
2040    ----------
2041    target_tree : xarray.DataTree
2042        The tree node to add perturbation groups to
2043    pdtn_df : pandas.DataFrame
2044        DataFrame of messages for a specific PDTN
2045    filename : str
2046        Path to the GRIB2 file
2047
2048    Returns
2049    -------
2050    bool
2051        True if at least one perturbation was successfully processed
2052    """
2053    success = False
2054    # Group by perturbation number
2055    pert_groups = {}
2056    for pert_value in pdtn_df['perturbationNumber'].unique():
2057        if pd.notna(pert_value):
2058            pert_df = pdtn_df[pdtn_df['perturbationNumber'] == pert_value]
2059            pert_groups[pert_value] = pert_df
2060
2061    # Process each perturbation group
2062    for pert_num, pert_df in pert_groups.items():
2063        pert_name = f"pert_{int(pert_num)}"
2064
2065        ## Try to create dataset for this perturbation group
2066        #try:
2067        #    dss = create_datasets_from_df(pert_df, filename)
2068        #    if dss is not None:
2069        #        if len(dss) == 1:
2070        #            target_tree.ds = dss[0]
2071        #        else:
2072        #            dss_dict = {f"ds_{i}": ds for i, ds in enumerate(dss)}
2073        #            atree = xr.DataTree(dss_dict)
2074        #            target_tree[prob_name] = atree
2075        #        success = True
2076        #except Exception as e:
2077        #    # Log error but continue processing other groups
2078        #    print(f"Error creating dataset for perturbation {pert_name}: {e}")
2079
2080        # Try to create dataset for this perturbation group
2081        try:
2082            dss = create_datasets_from_df(pert_df, filename)
2083            dt = xr.DataTree()
2084            if len(dss) == 1:
2085                dt.ds = dss[0]
2086                target_tree[pert_name] = dt
2087            elif len(dss) > 1:
2088                for ds in dss:
2089                    dt[f"pert{ds.data_vars[0]}"] = ds
2090            target_tree[pert_name] = dt
2091        except Exception as e:
2092            # Log error but continue processing other groups
2093            print(f"Error creating dataset for perturbation {pert_name}: {e}")
2094
2095    return success
2096
2097
2098def try_process_by_variables(target_tree, df, filename):
2099    """
2100    Try to separate data by variable names and create datasets.
2101
2102    Parameters
2103    ----------
2104    target_tree : xarray.DataTree
2105        The tree node to add variable datasets to
2106    df : pandas.DataFrame
2107        DataFrame of messages
2108    filename : str
2109        Path to the GRIB2 file
2110
2111    Returns
2112    -------
2113    bool
2114        True if at least one variable was successfully processed
2115    """
2116    success = False
2117
2118    try:
2119        for var_name in df['shortName'].unique():
2120            if pd.notna(var_name):
2121                var_df = df[df['shortName'] == var_name]
2122                try:
2123                    var_ds = create_datasets_from_df(var_df, filename)
2124                    if var_ds is not None:
2125                        target_tree[f"var_{var_name}"] = var_ds[0]
2126                        success = True
2127                except Exception as var_e:
2128                    print(f"Error creating dataset for variable {var_name}: {var_e}")
2129    except Exception as nested_e:
2130        print(f"Failed to process variables: {nested_e}")
2131
2132    return success
2133
2134
2135def create_datasets_from_df(
2136    df,
2137    filename,
2138    verbose=False
2139) -> typing.Optional[typing.List[xr.Dataset]]:
2140    """
2141    Create a list of xarray Datasets from a DataFrame of messages.
2142
2143    Parameters
2144    ----------
2145    df : pandas.DataFrame
2146        DataFrame of GRIB messages
2147    filename : str
2148        Path to the GRIB2 file
2149    verbose : bool, optional
2150        If True, prints detailed debugging information
2151
2152    Returns
2153    -------
2154    dss
2155        List of Datasets, or None if creation failed
2156    """
2157    try:
2158        if verbose:
2159            print(f"\n==== VERBOSE DEBUG INFO ====")
2160            print(f"Creating dataset from DataFrame with {len(df)} messages")
2161            print(f"DataFrame columns: {df.columns.tolist()}")
2162
2163            if 'shortName' in df.columns:
2164                print(f"Variables in group: {df['shortName'].unique().tolist()}")
2165
2166            if 'valueOfFirstFixedSurface' in df.columns:
2167                print(f"Vertical levels: {df['valueOfFirstFixedSurface'].unique().tolist()}")
2168
2169        # Process by variables
2170        datasets = {}
2171
2172        # Process each variable separately, regardless of whether there are vertical levels
2173        for var_name, var_df in df.groupby('shortName'):
2174            if verbose:
2175                print(
2176                    f"\n  Processing variable: {var_name} with {len(var_df)} messages, with pdtn(s) = {var_df['productDefinitionTemplateNumber'].unique()}")
2177
2178            # Process vertical levels if present
2179            if 'valueOfFirstFixedSurface' in var_df.columns and len(var_df['valueOfFirstFixedSurface'].unique()) > 1:
2180                if verbose:
2181                    print(f"  Variable {var_name} has multiple vertical levels")
2182                # Process each level separately
2183                level_das = []
2184
2185                for level, level_df in var_df.groupby('valueOfFirstFixedSurface'):
2186                    if verbose:
2187                        print(f"    Processing level {level} with {len(level_df)} messages")
2188                    try:
2189                        # Parse the index and get dimensions for this level
2190                        file_index, non_geo_dims, attrs, coord_attrs = parse_grib_index(level_df, {})
2191                        # Remove valueOfFirstFixedSurface from dimensions since we're handling it separately
2192                        non_geo_dims = [d for d in non_geo_dims if d.__name__ != "ValueOfFirstFixedSurfaceDim"]
2193
2194                        frames, cube, extra_geo = make_variables(
2195                            file_index, filename, non_geo_dims, allow_uneven_dims=True)
2196
2197                        if frames is not None and len(frames) == 1:
2198                            level_da = build_da_without_coords(frames[0], cube, filename, attrs)
2199                            # Add this level to the list with its level value as coord
2200                            level_da = level_da.assign_coords(valueOfFirstFixedSurface=level)
2201                            level_das.append(level_da)
2202                    except Exception as e:
2203                        if verbose:
2204                            print(f"    Error processing level {level} for {var_name}: {e}")
2205
2206                if level_das:
2207                    # Combine all levels into a single DataArray along the valueOfFirstFixedSurface dimension
2208                    if verbose:
2209                        print(f"    Combining {len(level_das)} levels for {var_name}")
2210                    try:
2211                        combined_da = xr.concat(level_das, dim='valueOfFirstFixedSurface')
2212                        # Create a simple dataset with just this variable
2213                        var_ds = xr.Dataset({var_name: combined_da})
2214                        # Assign the coords from the first level's cube
2215                        var_ds = assign_xr_meta(var_ds, frames, cube, non_geo_dims, extra_geo, coord_attrs)
2216                       # TODO: is the below code all now in assign_xr_meta? was there instances where refDate and leadTime were not coords?
2217                       # var_ds = var_ds.assign_coords(coords_from_cube(cube))
2218                       # Add extra geo coords
2219                       # if extra_geo:
2220                       #    var_ds = var_ds.assign_coords(extra_geo)
2221                       # Add valid date coords if available
2222                       # if 'refDate' in var_ds.coords and 'leadTime' in var_ds.coords:
2223                       #    var_ds = var_ds.assign_coords(dict(validDate=var_ds.coords['refDate']+var_ds.coords['leadTime']))
2224
2225                        # Store this variable's dataset
2226                        datasets[var_name] = var_ds
2227                        if verbose:
2228                            print(f"    Created dataset for {var_name} with levels")
2229                    except Exception as e:
2230                        if verbose:
2231                            print(f"    Error combining levels for {var_name}: {e}")
2232            else:
2233                # Single level or no vertical levels
2234                if verbose:
2235                    print(f"  Variable {var_name} is a single level or has no vertical dimension")
2236                try:
2237                    # Parse the index and get dimensions
2238                    file_index, non_geo_dims, attrs, coord_attrs = parse_grib_index(var_df, {})
2239                    frames, cube, extra_geo = make_variables(file_index, filename, non_geo_dims, allow_uneven_dims=True)
2240
2241                    if frames is not None and len(frames) == 1:
2242                        # Create dataset with this variable
2243                        var_ds = xr.Dataset()
2244                        da = build_da_without_coords(frames[0], cube, filename, attrs)
2245                        var_ds[da.name] = da
2246
2247                        # Assign coords
2248                        var_ds = assign_xr_meta(var_ds, frames, cube, non_geo_dims, extra_geo, coord_attrs)
2249                       # TODO: is the below code all now in assign_xr_meta? was there instances where refDate and leadTime were not coords?
2250                       # var_ds = var_ds.assign_coords(coords_from_cube(cube))
2251                       # if extra_geo:
2252                       #    var_ds = var_ds.assign_coords(extra_geo)
2253                       # if 'refDate' in var_ds.coords and 'leadTime' in var_ds.coords:
2254                       #    var_ds = var_ds.assign_coords(dict(validDate=var_ds.coords['refDate']+var_ds.coords['leadTime']))
2255
2256                        # Store this variable's dataset
2257                        datasets[var_name] = var_ds
2258                        if verbose:
2259                            print(f"  Created dataset for {var_name}")
2260                    elif frames is not None and len(frames) > 1:
2261                        if verbose:
2262                            print(f"  Variable {var_name} has multiple frames, possibly different parameters")
2263                        # Just use the first frame for now (simplified approach)
2264                        var_ds = xr.Dataset()
2265                        da = build_da_without_coords(frames[0], cube, filename, attrs)
2266                        var_ds[da.name] = da
2267
2268                        # Assign coords
2269                        var_ds = assign_xr_meta(var_ds, frames, cube, non_geo_dims, extra_geo, coord_attrs)
2270                       # TODO: is the below code all now in assign_xr_meta? was there instances where refDate and leadTime were not coords?
2271                       # var_ds = var_ds.assign_coords(coords_from_cube(cube))
2272                       # if extra_geo:
2273                       #    var_ds = var_ds.assign_coords(extra_geo)
2274                       # if 'refDate' in var_ds.coords and 'leadTime' in var_ds.coords:
2275                       #    var_ds = var_ds.assign_coords(dict(validDate=var_ds.coords['refDate']+var_ds.coords['leadTime']))
2276
2277                        datasets[var_name] = var_ds
2278                        if verbose:
2279                            print(f"  Created dataset with first frame for {var_name}")
2280                except Exception as e:
2281                    if verbose:
2282                        print(f"  Error processing variable {var_name}: {e}")
2283
2284        # Attempt to merge all the variable datasets
2285        if datasets:
2286            try:
2287                if verbose:
2288                    print(f"\nMerging {len(datasets)} datasets...")
2289                # Get the list of datasets to merge
2290                ds_list = list(datasets.values())
2291
2292                # Try merging them all at once
2293                try:
2294                    combined_ds = xr.merge(ds_list)
2295                    if verbose:
2296                        print(f"Successfully merged all datasets into one.")
2297                        print(f"Final dataset has variables: {list(combined_ds.data_vars)}")
2298                        print(f"==== END VERBOSE DEBUG INFO ====\n")
2299                    return [combined_ds]
2300                except Exception as merge_error:
2301                    if verbose:
2302                        print(f"Error merging all datasets: {merge_error}")
2303                    return ds_list
2304            except Exception as e:
2305                if verbose:
2306                    print(f"Error in final merge process: {e}")
2307                    print(f"==== END VERBOSE DEBUG INFO ====\n")
2308                return None
2309        else:
2310            if verbose:
2311                print(f"No datasets were created for any variables")
2312                print(f"==== END VERBOSE DEBUG INFO ====\n")
2313            return None
2314
2315    except Exception as e:
2316        # If there's an error, log it and return None
2317        if verbose:
2318            print(f"Error creating dataset: {e}")
2319            import traceback
2320            traceback.print_exc()
2321            print(f"==== END VERBOSE DEBUG INFO ====\n")
2322        return None
2323
2324
2325# Only register the DataTree accessor if DataTree is supported
2326if _HAS_DATATREE:
2327    @xr.register_datatree_accessor("grib2io")
2328    class Grib2ioDataTree:
2329        """
2330        DataTree accessor for GRIB2 files.
2331
2332        This accessor provides methods for working with GRIB2 data organized
2333        in a hierarchical tree structure.
2334        """
2335
2336        def __init__(self, datatree_obj):
2337            self._obj = datatree_obj
2338
2339        def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
2340            """
2341            Write all datasets in the DataTree to a GRIB2 file.
2342
2343            Parameters
2344            ----------
2345            filename : str
2346                Name of the GRIB2 file to write to.
2347            mode : {"x", "w", "a"}, optional
2348                Persistence mode, default is "x" (create, fail if exists)
2349            """
2350            # Start with the specified mode
2351            current_mode = mode
2352
2353            # Function to recursively process the tree
2354            def process_tree(node):
2355                nonlocal current_mode
2356
2357                # If this is a Dataset node with data variables
2358                if node.ds is not None and node.ds.data_vars:
2359                    # Write dataset to GRIB2 file
2360                    node.ds.grib2io.to_grib2(filename, mode=current_mode)
2361                    # Switch to append mode after first write
2362                    current_mode = "a"
2363
2364                # Process children
2365                for child_name, child_node in node.children.items():
2366                    process_tree(child_node)
2367
2368            # Start processing from the root
2369            process_tree(self._obj)
2370
2371        def griddef(self):
2372            """
2373            Get the grid definition from the first dataset in the tree that has one.
2374
2375            Returns
2376            -------
2377            grib2io.Grib2GridDef
2378                Grid definition object
2379            """
2380            # Function to find first dataset with GRIB2IO_section3
2381            def find_griddef(node):
2382                if node.ds is not None and node.ds.data_vars:
2383                    for var_name in node.ds.data_vars:
2384                        if 'GRIB2IO_section3' in node.ds[var_name].attrs:
2385                            return Grib2GridDef.from_section3(node.ds[var_name].attrs['GRIB2IO_section3'])
2386
2387                # Check children
2388                for child_name, child_node in node.children.items():
2389                    griddef = find_griddef(child_node)
2390                    if griddef is not None:
2391                        return griddef
2392
2393                return None
2394
2395            return find_griddef(self._obj)
2396
2397        def interp(self, method, grid_def_out, method_options=None, num_threads=1):
2398            """
2399            Interpolate all datasets in the tree to a new grid.
2400
2401            Parameters
2402            ----------
2403            method : str or int
2404                Interpolation method to use
2405            grid_def_out : grib2io.Grib2GridDef
2406                Target grid definition
2407            method_options : list, optional
2408                Options for interpolation method
2409            num_threads : int, optional
2410                Number of threads to use for interpolation
2411
2412            Returns
2413            -------
2414            xarray.DataTree
2415                New DataTree with interpolated data
2416            """
2417            new_tree = xr.DataTree()
2418
2419            # Function to recursively process the tree
2420            def process_tree(node, new_parent):
2421                # If this is a Dataset node with data variables
2422                if node.ds is not None and node.ds.data_vars:
2423                    # Interpolate dataset
2424                    interp_ds = node.ds.grib2io.interp(method, grid_def_out,
2425                                                       method_options=method_options,
2426                                                       num_threads=num_threads)
2427
2428                    # Add to new tree at the same path
2429                    if node == self._obj:  # Root node
2430                        new_parent.ds = interp_ds
2431                    else:
2432                        new_parent.ds = interp_ds
2433
2434                # Process children
2435                for child_name, child_node in node.children.items():
2436                    # Create same child in new tree
2437                    new_child = xr.DataTree()
2438                    new_parent[child_name] = new_child
2439                    process_tree(child_node, new_child)
2440
2441            # Start processing from the root
2442            process_tree(self._obj, new_tree)
2443
2444            return new_tree
2445
2446        def subset(self, lats, lons):
2447            """
2448            Subset all datasets in the tree to a region.
2449
2450            Parameters
2451            ----------
2452            lats : list or tuple
2453                Latitude bounds [min_lat, max_lat]
2454            lons : list or tuple
2455                Longitude bounds [min_lon, max_lon]
2456
2457            Returns
2458            -------
2459            xarray.DataTree
2460                New DataTree with subset data
2461            """
2462            new_tree = xr.DataTree()
2463
2464            # Function to recursively process the tree
2465            def process_tree(node, new_parent):
2466                # If this is a Dataset node with data variables
2467                if node.ds is not None and node.ds.data_vars:
2468                    # Subset dataset
2469                    subset_ds = node.ds.grib2io.subset(lats, lons)
2470
2471                    # Add to new tree at the same path
2472                    if node == self._obj:  # Root node
2473                        new_parent.ds = subset_ds
2474                    else:
2475                        new_parent.ds = subset_ds
2476
2477                # Process children
2478                for child_name, child_node in node.children.items():
2479                    # Create same child in new tree
2480                    new_child = xr.DataTree()
2481                    new_parent[child_name] = new_child
2482                    process_tree(child_node, new_child)
2483
2484            # Start processing from the root
2485            process_tree(self._obj, new_tree)
2486
2487            return new_tree