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"""
  18
  19from grib2io._grib2io import _data
  20from grib2io import Grib2Message, Grib2GridDef, msgs_from_index, tables, templates
  21from grib2io.utils.spatial import snap_to_nearest_cell_center
  22import grib2io
  23from xarray.backends.locks import SerializableLock
  24from xarray.core import indexing
  25from xarray.backends import (
  26    BackendArray,
  27    BackendEntrypoint,
  28)
  29from copy import copy
  30from dataclasses import dataclass, field, astuple
  31import datetime
  32import glob
  33import itertools
  34import json
  35import logging
  36import os
  37import re
  38import textwrap
  39import typing
  40from numpy.typing import NDArray
  41import warnings
  42
  43from . import tables
  44
  45import numpy as np
  46import pandas as pd
  47import xarray as xr
  48from pathlib import Path
  49from pyproj import CRS
  50
  51# Check if xarray version supports DataTree
  52_HAS_DATATREE = hasattr(xr, "DataTree")
  53
  54# Check for NumPy 2.0+ StringDType
  55_HAS_STRINGDTYPE = hasattr(np, "dtypes") and hasattr(np.dtypes, "StringDType")
  56
  57_logger = logging.getLogger(__name__)
  58
  59_DEFAULT_REMOTE_STORAGE_OPTIONS = {
  60    "config_kwargs": {
  61        "connect_timeout": 30,
  62        "read_timeout": 120,
  63        "retries": {"max_attempts": 10, "mode": "adaptive"},
  64    }
  65}
  66
  67
  68def _merge_default_storage_options(storage_options):
  69    """Merge user storage options over conservative remote defaults."""
  70    merged = {
  71        "config_kwargs": {
  72            "connect_timeout": _DEFAULT_REMOTE_STORAGE_OPTIONS["config_kwargs"]["connect_timeout"],
  73            "read_timeout": _DEFAULT_REMOTE_STORAGE_OPTIONS["config_kwargs"]["read_timeout"],
  74            "retries": dict(_DEFAULT_REMOTE_STORAGE_OPTIONS["config_kwargs"]["retries"]),
  75        }
  76    }
  77    user = storage_options or {}
  78    if not isinstance(user, dict):
  79        return merged
  80
  81    for key, value in user.items():
  82        if key == "config_kwargs" and isinstance(value, dict):
  83            merged_cfg = merged.setdefault("config_kwargs", {})
  84            for cfg_key, cfg_val in value.items():
  85                if cfg_key == "retries" and isinstance(cfg_val, dict):
  86                    retries = merged_cfg.setdefault("retries", {})
  87                    retries.update(cfg_val)
  88                else:
  89                    merged_cfg[cfg_key] = cfg_val
  90        else:
  91            merged[key] = value
  92    return merged
  93
  94
  95_LOCK = SerializableLock()
  96
  97_LEVEL_NAME_MAPPING = grib2io.tables.get_table("4.5.grib2io.level.name")
  98
  99_TREE_HIERARCHY_LEVELS = [
 100    "typeOfFirstFixedSurface",
 101    "valueOfFirstFixedSurface",
 102    "productDefinitionTemplateNumber",
 103    "perturbationNumber",
 104    "leadTime",
 105    "duration",
 106    "percentileValue",
 107    "typeOfProbability",
 108    "thresholdLowerLimit",
 109    "thresholdUpperLimit",
 110]
 111
 112
 113def _decode_ptype(values: np.ndarray) -> np.ndarray:
 114    """
 115    Decode numeric precipitation type codes into human-readable strings.
 116
 117    Parameters
 118    ----------
 119    values : np.ndarray
 120        Array of numeric precipitation type codes.
 121
 122    Returns
 123    -------
 124    np.ndarray
 125        Array of decoded precipitation type strings.
 126    """
 127    return _decode_code(values, "4.201")
 128
 129
 130def _decode_code(values: np.ndarray, table: str) -> np.ndarray:
 131    """
 132    Decode numeric codes into human-readable strings using a GRIB2 table.
 133
 134    Parameters
 135    ----------
 136    values : np.ndarray
 137    Array of numeric codes.
 138    table : str
 139    The GRIB2 table to use for decoding (e.g., "4.201").
 140
 141    Returns
 142    -------
 143    np.ndarray
 144    Array of decoded string definitions.
 145    """
 146
 147    def _lookup(val):
 148        res = tables.get_value_from_table(str(int(val)), table)
 149        if isinstance(res, list):
 150            return str(res[0])
 151        return str(res)
 152
 153    # Pass otypes to avoid string truncation based on the first element
 154    # Use StringDType if available (NumPy 2.0+), otherwise fallback to object
 155    if _HAS_STRINGDTYPE:
 156        vlookup = np.vectorize(_lookup, otypes=[np.dtypes.StringDType])
 157    else:
 158        vlookup = np.vectorize(_lookup, otypes=[object])
 159    return vlookup(values)
 160
 161
 162AVAILABLE_NON_GEO_COORDS = [
 163    "duration",
 164    "leadTime",
 165    "percentileValue",
 166    "perturbationNumber",
 167    "refDate",
 168    "thresholdLowerLimit",
 169    "thresholdUpperLimit",
 170    "valueOfFirstFixedSurface",
 171    "valueOfSecondFixedSurface",
 172    "typeOfAerosol",
 173    "constituentType",
 174    "sourceSinkIndicator",
 175    "firstWavelength",
 176    "secondWavelength",
 177    "firstSizeOfAerosol",
 178    "secondSizeOfAerosol",
 179    "scaledValueOfFirstWavelength",
 180    "scaledValueOfSecondWavelength",
 181    "scaledValueOfCentralWaveNumber",
 182    "scaledValueOfFirstSize",
 183    "scaledValueOfSecondSize",
 184]
 185"""Available non-geographic coordinate names."""
 186
 187AVAILABLE_NON_GEO_DIMS = [
 188    "duration",
 189    "leadTime",
 190    "percentileValue",
 191    "perturbationNumber",
 192    "refDate",
 193    "threshold",
 194    "level",
 195    "typeOfAerosol",
 196    "constituentType",
 197    "sourceSinkIndicator",
 198    "firstWavelength",
 199    "secondWavelength",
 200    "firstSizeOfAerosol",
 201    "secondSizeOfAerosol",
 202]
 203"""Available non-geographic dimension names."""
 204
 205# Lookup table to define surface types that should be parsed as vertical coordinates
 206VERTICAL_COORDINATE_SURFACES = [
 207    "Ground or Water Surface",
 208    "Isothermal Level",
 209    "Specified radius from the centre of the Sun",
 210    "Isobaric Surface",
 211    "Mean Sea Level",
 212    "Specific Altitude Above Mean Sea Level",
 213    "Specified Height Level Above Ground",
 214    "Sigma Level",
 215    "Hybrid Level",
 216    "Depth Below Land Surface",
 217    "Isentropic (theta) Level",
 218    "Level at Specified Pressure Difference from Ground to Level",
 219    "Potential Vorticity Surface",
 220    "Eta Level",
 221    "Logarithmic Hybrid Level",
 222    "Sigma height level",
 223    "Hybrid Height Level",
 224    "Hybrid Pressure Level",
 225    "Soil level",
 226    "Sea-ice level",
 227    "Depth Below Sea Level",
 228    "Depth Below Water Surface",
 229    "Ocean Model Level",
 230    "Ocean level defined by water density (sigma-theta) difference from near-surface to level",
 231    "Ocean level defined by water potential temperature difference from near-surface to level",
 232    "Ocean level defined by vertical eddy diffusivity difference from near-surface to level",
 233    "Ocean level defined by water density (rho) difference from near-surface to level",
 234]
 235"""
 236Lookup table to define surface types that should be parsed as vertical coordinates
 237when `data_model="nws-viz"`.
 238"""
 239
 240
 241def parse_data_model(ds: xr.Dataset, data_model: str) -> xr.Dataset:
 242    """
 243    Normalize a GRIB2-derived Dataset to a target data model (currently ``"nws-viz"``).
 244
 245    When ``data_model == "nws-viz"``, this function converts coordinate and
 246    variable names to snake_case, derives CF-like metadata, promotes select
 247    GRIB-derived quantities to coordinates, optionally swaps dimensions, and
 248    standardizes units/attributes. If ``data_model`` is anything else, the
 249    input dataset is returned unchanged.
 250
 251    Parameters
 252    ----------
 253    ds : xarray.Dataset
 254        GRIB2-derived dataset whose variables and attributes follow the
 255        conventions emitted by ``grib2io``. Expected to contain GRIB-related
 256        attributes such as ``typeOfFirstFixedSurface``,
 257        ``typeOfSecondFixedSurface``, and (for probabilistic variables)
 258        ``typeOfProbability``.
 259    data_model : str
 260        Target data model name. Only the value ``"nws-viz"`` triggers
 261        transformations.
 262
 263    Returns
 264    -------
 265    xarray.Dataset
 266        A new dataset with:
 267        * Selected coordinates renamed:
 268          ``refDate -> forecast_reference_time``,
 269          ``leadTime -> lead_time``,
 270          ``validDate -> time``,
 271          ``percentileValue -> percentile``,
 272          ``thresholdLowerLimit -> threshold_lower_limit``,
 273          ``thresholdUpperLimit -> threshold_upper_limit``.
 274        * Vertical coordinates derived from
 275          ``valueOfFirstFixedSurface`` / ``valueOfSecondFixedSurface`` and their
 276          corresponding ``typeOf*FixedSurface`` definitions. New coordinate
 277          names are generated from the surface definition (lowercased, spaces
 278          to underscores, punctuation removed). If the name already exists, a
 279          ``"_2"`` suffix is appended.
 280        * Possible dimension swaps:
 281          ``level -> <derived_vertical_coord>`` when present; and for
 282          probabilistic variables, ``threshold -> threshold_lower_limit`` or
 283          ``threshold -> threshold_upper_limit`` when
 284          ``typeOfProbability`` indicates the appropriate semantics.
 285        * Variable names lowercased; dataset- and variable-level attributes
 286          converted to snake_case (except GRIB section attributes which are
 287          normalized to ``grib...``).
 288        * CF-adjacent metadata populated: ``standard_name`` and
 289          ``cell_methods`` are set via the shortname→CF lookup table.
 290        * Percent units normalized from ``"%"`` to ``"percent"`` on coordinates.
 291        * For precipitation type (``PTYPE``) thresholds, numeric codes are
 292          decoded to strings (GRIB2 Table 4.201) in relevant attrs/coords.
 293
 294    Notes
 295    -----
 296    - Precipitation type decoding uses GRIB2 Table 4.201 via
 297      ``tables.get_value_from_table(code, "4.201")`` and returns a NumPy
 298      array with ``np.dtypes.StringDType``.
 299    - CF-related lookups are performed using
 300      ``tables.get_table("shortname_to_cf")``.
 301    - Vertical coordinate surface names are validated against
 302      ``VERTICAL_COORDINATE_SURFACES`` before promotion to coordinates.
 303
 304    Warnings
 305    --------
 306    This function assumes the presence of certain GRIB-derived attributes on the
 307    first data variable (e.g., ``typeOfFirstFixedSurface``,
 308    ``typeOfSecondFixedSurface``, and possibly ``typeOfProbability``).
 309    If these are absent or malformed, errors (e.g., ``KeyError``) may occur.
 310
 311    Examples
 312    --------
 313    >>> ds2 = parse_data_model(ds, "nws-viz")
 314    >>> list(ds2.coords)
 315    ['forecast_reference_time', 'lead_time', 'time', 'percentile', ...]
 316    """
 317    # convert coordinates and attributes to CF if requested
 318    if data_model == "nws-viz":
 319        # define regex to convert to snake case
 320        pattern = re.compile(r"(?<!^)(?=[A-Z])")
 321
 322        # check for coordinates and rename
 323        for coord in ds.coords:
 324            if coord == "refDate":
 325                ds = ds.rename({"refDate": "forecast_reference_time"})
 326
 327            elif coord == "leadTime":
 328                ds = ds.rename({"leadTime": "lead_time"})
 329
 330            elif coord == "validDate":
 331                ds = ds.rename({"validDate": "time"})
 332
 333            elif coord == "percentileValue":
 334                ds = ds.rename({"percentileValue": "percentile"})
 335
 336            elif coord == "perturbationNumber":
 337                ds = ds.rename({"perturbationNumber": "perturbation"})
 338                ds["perturbation"].attrs["long_name"] = "Ensemble Perturbation Number"
 339
 340            elif coord == "thresholdLowerLimit":
 341                ds = ds.rename({"thresholdLowerLimit": "threshold_lower_limit"})
 342                ds["threshold_lower_limit"].attrs["long_name"] = "Threshold Lower Limit"
 343                ds["threshold_lower_limit"].attrs["units"] = ds[list(ds.data_vars.keys())[0]].attrs["units"]
 344
 345                if "PTYPE" in ds.data_vars:
 346                    ds["threshold_lower_limit"] = xr.apply_ufunc(
 347                        _decode_ptype,
 348                        ds["threshold_lower_limit"],
 349                        dask="parallelized",
 350                        output_dtypes=[np.dtypes.StringDType] if _HAS_STRINGDTYPE else [object],
 351                    )
 352
 353                # check if thresholdLowerLimit should be a dimension coordinate
 354                if "threshold" in ds.dims:
 355                    var_key = list(ds.data_vars.keys())[0]
 356                    prob_types = [
 357                        "Probability of event below lower limit",
 358                        "Probability of event above lower limit",
 359                        "Probability of event equal to lower limit",
 360                        "Probability of event between upper and lower limits (the range includes lower limit but not the upper limit)",
 361                    ]
 362                    if ds[var_key].attrs["typeOfProbability"] in prob_types:
 363                        ds = ds.swap_dims({"threshold": "threshold_lower_limit"})
 364
 365            elif coord == "thresholdUpperLimit":
 366                ds = ds.rename({"thresholdUpperLimit": "threshold_upper_limit"})
 367                ds["threshold_upper_limit"].attrs["long_name"] = "Threshold Upper Limit"
 368                ds["threshold_upper_limit"].attrs["units"] = ds[list(ds.data_vars.keys())[0]].attrs["units"]
 369
 370                if "PTYPE" in ds.data_vars:
 371                    ds["threshold_upper_limit"] = xr.apply_ufunc(
 372                        _decode_ptype,
 373                        ds["threshold_upper_limit"],
 374                        dask="parallelized",
 375                        output_dtypes=[np.dtypes.StringDType] if _HAS_STRINGDTYPE else [object],
 376                    )
 377
 378                if "threshold" in ds.dims:
 379                    var_key = list(ds.data_vars.keys())[0]
 380                    prob_types = [
 381                        "Probability of event below upper limit",
 382                        "Probability of event above upper limit",
 383                    ]
 384                    if ds[var_key].attrs["typeOfProbability"] in prob_types:
 385                        ds = ds.swap_dims({"threshold": "threshold_upper_limit"})
 386
 387            elif coord == "typeOfAerosol":
 388                ds = ds.rename({"typeOfAerosol": "aerosol_type"})
 389                ds["aerosol_type"].attrs["long_name"] = "Aerosol Type"
 390                ds["aerosol_type"] = xr.apply_ufunc(
 391                    _decode_code,
 392                    ds["aerosol_type"],
 393                    "4.233",
 394                    dask="parallelized",
 395                    output_dtypes=[np.dtypes.StringDType] if _HAS_STRINGDTYPE else [object],
 396                )
 397
 398            elif coord == "constituentType":
 399                ds = ds.rename({"constituentType": "constituent_type"})
 400                ds["constituent_type"].attrs["long_name"] = "Chemical Constituent Type"
 401                ds["constituent_type"] = xr.apply_ufunc(
 402                    _decode_code,
 403                    ds["constituent_type"],
 404                    "4.230",
 405                    dask="parallelized",
 406                    output_dtypes=[np.dtypes.StringDType] if _HAS_STRINGDTYPE else [object],
 407                )
 408
 409            elif coord == "sourceSinkIndicator":
 410                ds = ds.rename({"sourceSinkIndicator": "source_sink_indicator"})
 411                ds["source_sink_indicator"].attrs["long_name"] = "Source/Sink Indicator"
 412                ds["source_sink_indicator"] = xr.apply_ufunc(
 413                    _decode_code,
 414                    ds["source_sink_indicator"],
 415                    "4.238",
 416                    dask="parallelized",
 417                    output_dtypes=[np.dtypes.StringDType] if _HAS_STRINGDTYPE else [object],
 418                )
 419
 420            elif coord == "firstWavelength":
 421                ds = ds.rename({"firstWavelength": "first_wavelength"})
 422                ds["first_wavelength"].attrs["long_name"] = "First Wavelength"
 423                ds["first_wavelength"].attrs["units"] = "m"
 424
 425            elif coord == "secondWavelength":
 426                ds = ds.rename({"secondWavelength": "second_wavelength"})
 427                ds["second_wavelength"].attrs["long_name"] = "Second Wavelength"
 428                ds["second_wavelength"].attrs["units"] = "m"
 429
 430            elif coord == "firstSizeOfAerosol":
 431                ds = ds.rename({"firstSizeOfAerosol": "first_size_of_aerosol"})
 432                ds["first_size_of_aerosol"].attrs["long_name"] = "First Size of Aerosol"
 433                ds["first_size_of_aerosol"].attrs["units"] = "m"
 434
 435            elif coord == "secondSizeOfAerosol":
 436                ds = ds.rename({"secondSizeOfAerosol": "second_size_of_aerosol"})
 437                ds["second_size_of_aerosol"].attrs["long_name"] = "Second Size of Aerosol"
 438                ds["second_size_of_aerosol"].attrs["units"] = "m"
 439
 440            elif coord == "scaledValueOfFirstWavelength":
 441                ds = ds.rename({"scaledValueOfFirstWavelength": "scaled_first_wavelength"})
 442                ds["scaled_first_wavelength"].attrs["long_name"] = "Scaled Value of First Wavelength"
 443
 444            elif coord == "scaledValueOfSecondWavelength":
 445                ds = ds.rename({"scaledValueOfSecondWavelength": "scaled_second_wavelength"})
 446                ds["scaled_second_wavelength"].attrs["long_name"] = "Scaled Value of Second Wavelength"
 447
 448            elif coord == "scaledValueOfCentralWaveNumber":
 449                ds = ds.rename({"scaledValueOfCentralWaveNumber": "scaled_central_wave_number"})
 450                ds["scaled_central_wave_number"].attrs["long_name"] = "Scaled Value of Central Wave Number"
 451
 452            elif coord == "scaledValueOfFirstSize":
 453                ds = ds.rename({"scaledValueOfFirstSize": "scaled_first_size"})
 454                ds["scaled_first_size"].attrs["long_name"] = "Scaled Value of First Size"
 455
 456            elif coord == "scaledValueOfSecondSize":
 457                ds = ds.rename({"scaledValueOfSecondSize": "scaled_second_size"})
 458                ds["scaled_second_size"].attrs["long_name"] = "Scaled Value of Second Size"
 459
 460            # If the dataset has valueOfFirstFixedSurface as a coordinate
 461            elif coord == "valueOfFirstFixedSurface":
 462                # Get the valueOfFirstFixedSurface coordinate
 463                da = ds.valueOfFirstFixedSurface
 464
 465                # Get the definition and units from typeOfFirstFixedSurface
 466                var_key = list(ds.data_vars.keys())[0]
 467                definition, units = ds[var_key].attrs["typeOfFirstFixedSurface"]
 468
 469                if definition in VERTICAL_COORDINATE_SURFACES:
 470                    # Convert definition to lowercase and replace spaces with underscores
 471                    key = definition.lower().replace(" ", "_")
 472
 473                    # remove special characters
 474                    key = re.sub(r"[^a-z0-9_]", "", key)
 475
 476                    # Add units and grib_name attributes
 477                    da.attrs["units"] = units
 478                    da.attrs["grib_name"] = [
 479                        "valueOfFirstFixedSurface",
 480                        "typeOfFirstFixedSurface",
 481                    ]
 482
 483                    # Assign the coordinate with the new key name
 484                    ds = ds.assign_coords({key: da})
 485
 486                    # If valueOfFirstFixedSurface is a dimension, swap it with the new key
 487                    if "level" in ds.dims:
 488                        ds = ds.swap_dims({"level": key})
 489
 490                # Remove the original coordinates
 491                del ds["valueOfFirstFixedSurface"]
 492
 493            # If the dataset has valueOfSecondFixedSurface as a coordinate
 494            elif coord == "valueOfSecondFixedSurface":
 495                # Get the valueOfSecondFixedSurface coordinate
 496                da = ds.valueOfSecondFixedSurface
 497
 498                # Get the definition and units from typeOfSecondFixedSurface
 499                var_key = list(ds.data_vars.keys())[0]
 500                definition, units = ds[var_key].attrs["typeOfSecondFixedSurface"]
 501
 502                if definition in VERTICAL_COORDINATE_SURFACES:
 503                    # Convert definition to lowercase and replace spaces with underscores
 504                    key = definition.lower().replace(" ", "_")
 505
 506                    # remove special characters
 507                    key = re.sub(r"[^a-z0-9_]", "", key)
 508
 509                    # check if key is already in coords
 510                    if key in ds.coords:
 511                        key = key + "_2"
 512
 513                    # Add units and grib_name attributes
 514                    da.attrs["units"] = units
 515                    da.attrs["grib_name"] = [
 516                        "valueOfSecondFixedSurface",
 517                        "typeOfSecondFixedSurface",
 518                    ]
 519
 520                    # Assign the coordinate with the new key name
 521                    ds = ds.assign_coords({key: da})
 522
 523                # Remove the original coordinates
 524                del ds["valueOfSecondFixedSurface"]
 525            else:
 526                # change coord name to snake case
 527                new_coord_name = pattern.sub("_", coord).lower()
 528                ds = ds.rename({coord: new_coord_name})
 529
 530        # convert all attributes and variable names to snake case
 531        for var in ds.data_vars:
 532            da = ds[var]
 533            record = tables.get_table("shortname_to_cf").get(da.name)
 534            da.attrs["standard_name"] = "unknown" if record is None else record["cf_standard_name"]
 535            da.attrs["cell_methods"] = "unknown" if record is None else record["cf_cell_methods"]
 536
 537            ds[var] = da
 538
 539            # rename variable
 540            new_var_name = var.lower()
 541            ds = ds.rename({var: new_var_name})
 542
 543            # remove attr for typeOfFirstFixedSurface (applied as coordinate above)
 544            if "typeOfFirstFixedSurface" in ds[new_var_name].attrs:
 545                definition, units = ds[new_var_name].attrs["typeOfFirstFixedSurface"]
 546                ds[new_var_name].attrs["typeOfFirstFixedSurface"] = f"{definition} ({units})"
 547
 548            if "typeOfSecondFixedSurface" in ds[new_var_name].attrs:
 549                definition, units = ds[new_var_name].attrs["typeOfSecondFixedSurface"]
 550                ds[new_var_name].attrs["typeOfSecondFixedSurface"] = f"{definition} ({units})"
 551
 552            ds[new_var_name].attrs.pop("percentileValue", None)
 553
 554            if "threshold_lower_limit" in ds.coords:
 555                ds[new_var_name].attrs.pop("thresholdLowerLimit", None)
 556
 557            if "threshold_upper_limit" in ds.coords:
 558                ds[new_var_name].attrs.pop("thresholdUpperLimit", None)
 559
 560            for attr in list(ds[new_var_name].attrs.keys()):
 561                # skip grib section attrs
 562                if "GRIB2IO_section" in attr:
 563                    # replace GRIB2IO with grib in attr
 564                    new_attr_name = attr.replace("GRIB2IO", "grib")
 565                else:
 566                    # change attr name to snake case
 567                    new_attr_name = pattern.sub("_", attr).lower()
 568
 569                # update new attr name for specific CF names
 570                if new_attr_name == "full_name":
 571                    new_attr_name = "long_name"
 572
 573                # change % to percent
 574                if attr == "units" and ds[new_var_name].attrs[attr] == "%":
 575                    ds[new_var_name].attrs[attr] = "percent"
 576
 577                if new_var_name == "ptype" and "threshold" in new_attr_name:
 578                    value = ds[new_var_name].attrs.pop(attr)
 579                    ds[new_var_name].attrs[attr] = _decode_ptype(value)
 580                else:
 581                    # change attr name in attrs
 582                    ds[new_var_name].attrs[new_attr_name] = ds[new_var_name].attrs.pop(attr)
 583
 584            try:
 585                new_cell_methods = section4_to_cell_methods(ds[new_var_name].attrs["grib_section4"])
 586            except KeyError:
 587                pass
 588            else:
 589                if new_cell_methods is not None:
 590                    if ds[new_var_name].attrs["cell_methods"] is None:
 591                        ds[new_var_name].attrs["cell_methods"] = new_cell_methods
 592                    else:
 593                        ds[new_var_name].attrs["cell_methods"] = " ".join(ds[new_var_name].attrs["cell_methods"], new_cell_methods)
 594
 595        # change dataset attrs to snake case
 596        for attr in list(ds.attrs.keys()):
 597            # change attr name to snake case
 598            new_attr_name = pattern.sub("_", attr).lower()
 599
 600            # change attr name in attrs
 601            ds.attrs[new_attr_name] = ds.attrs.pop(attr)
 602
 603        # change % to percent
 604        for coord in ds.coords:
 605            if "units" in ds[coord].attrs and ds[coord].attrs["units"] == "%":
 606                ds[coord].attrs["units"] = "percent"
 607
 608        # Update history for provenance
 609        history = ds.attrs.get("history", "")
 610        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
 611        ds.attrs["history"] = f"{now}: Parsed to data model {data_model}\n{history}"
 612
 613    # Update history for provenance
 614    history = ds.attrs.get("history", "")
 615    now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
 616    ds.attrs["history"] = f"{now}: Normalized to {data_model} data model\n{history}"
 617
 618    return ds
 619
 620
 621def section4_to_cell_methods(section4_array: np.ndarray) -> typing.Optional[str]:
 622    cell_methods = None
 623    if section4_array[1] == 0:
 624        cell_methods = f"{cell_methods:s} lead_time: point"
 625    elif section4_array[1] == 8:
 626        to_join = []
 627        # interval_end = datetime.datetime(*section4_array[17:23]
 628        time_unit_table = tables.get_table("4.4")
 629        for i in reversed(range(section4_array[23])):
 630            offset = 6 * i
 631            method = tables.get_table("4.10")[str(section4_array[25 + offset])]
 632            method = method.replace("Average", "mean").lower()
 633            if section4_array[26 + offset] == 1:
 634                dim = "forecast_reference_time"
 635            elif section4_array[26 + offset] == 2:
 636                dim = "lead_time"
 637            # duration_unit = time_unit_table[str(section4_array[27 + offset])].lower()
 638            # duration_value = section4_array[28 + offset]
 639            input_interval_units = time_unit_table[str(section4_array[29 + offset])].lower()
 640            input_interval_value = section4_array[30 + offset]
 641            to_join.append(f"{dim:s}: {method:s} (interval: {input_interval_value:d} {input_interval_units:s})")
 642            # comment: duration {duration_value:d} {duration_unit:s} ending {interval_end:%Y-%m-%dT%H:%M:%s}
 643        cell_methods = " ".join(to_join)
 644    return cell_methods
 645
 646
 647# ---------------------------------------------------------------------------
 648# Lazy import guards for optional dependencies
 649# ---------------------------------------------------------------------------
 650
 651
 652class GribBackendEntrypoint(BackendEntrypoint):
 653    """
 654    xarray backend engine entrypoint for opening and decoding grib2 files.
 655
 656    .. warning::
 657
 658       This backend is experimental and the API/behavior may change without
 659       backward compatibility.
 660    """
 661
 662    def open_dataset(
 663        self,
 664        filename_or_obj,
 665        drop_variables=None,
 666        save_index=True,
 667        filters=None,
 668        data_model=None,
 669        chunks=None,
 670        storage_options=None,
 671    ) -> xr.Dataset:
 672        """
 673        Read and parse metadata from a GRIB2 file.
 674
 675        Parameters
 676        ----------
 677        filename_or_obj : str or file-like
 678            GRIB2 file to be opened. Can be a local path or a remote URI.
 679        drop_variables : list of str, optional
 680            List of variables to exclude from the dataset.
 681        save_index : bool, optional
 682            Whether to save the GRIB2 index to a file (default is True).
 683        filters : dict, optional
 684            Filter GRIB2 messages to a single hypercube. Dictionary keys can
 685            be any GRIB2 metadata attribute name.
 686        data_model : str, optional
 687            Parse GRIB metadata following a defined data model convention
 688            (e.g., "nws-viz").
 689        chunks : int, dict or 'auto', optional
 690            If chunks is provided, it is used to load the dataset into a
 691            dask-backed dataset.
 692        storage_options : dict, optional
 693            Extra options passed to the storage backend.
 694
 695        Returns
 696        -------
 697        xarray.Dataset
 698            Xarray dataset of GRIB2 messages.
 699        """
 700        if filters is None:
 701            filters = {}
 702
 703        with grib2io.open(
 704            filename_or_obj,
 705            save_index=save_index,
 706            _xarray_backend=True,
 707            **(storage_options or {}),
 708        ) as f:
 709            file_index = pd.DataFrame(f._index)
 710            file_index = file_index.assign(msg=list(f))
 711
 712        ds = _open_dataset_from_index(
 713            file_index,
 714            filename_or_obj,
 715            filters,
 716            data_model,
 717            drop_variables=drop_variables,
 718            chunks=chunks,
 719        )
 720
 721        # Update history for provenance
 722        history = ds.attrs.get("history", "")
 723        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
 724        ds.attrs["history"] = f"{now}: Initialized via grib2io.open_dataset from {filename_or_obj}\n{history}"
 725
 726        return ds
 727
 728    def open_datatree(
 729        self,
 730        filename_or_obj,
 731        drop_variables=None,
 732        save_index=True,
 733        filters=None,
 734        stack_vertical=False,
 735        chunks=None,
 736    ) -> typing.Any:
 737        """
 738        Open a GRIB2 file as an xarray DataTree.
 739
 740        Parameters
 741        ----------
 742        filename : str
 743            Path to the GRIB2 file.
 744        drop_variables : list, optional
 745            List of variables to exclude.
 746        filters : dict, optional
 747            Filter criteria for GRIB2 messages.
 748        stack_vertical : bool, optional
 749            If True, organize the tree with vertical layers stacked in a single dataset.
 750        chunks : int, dict or 'auto', optional
 751            If chunks is provided, it is used to load the dataset into a
 752            dask-backed dataset.
 753
 754        Returns
 755        -------
 756        xarray.DataTree
 757            A hierarchical DataTree representation of the GRIB2 data.
 758        """
 759        if not _HAS_DATATREE:
 760            raise ImportError("xarray version does not support DataTree functionality.")
 761
 762        if filters is None:
 763            filters = {}
 764
 765        # Open the file without any filters first to get all messages
 766        with grib2io.open(filename_or_obj, save_index=save_index, _xarray_backend=True) as f:
 767            file_index = pd.DataFrame(f._index)
 768            file_index = file_index.assign(msg=list(f))
 769
 770        # Build tree structure from GRIB messages with specified options
 771        tree = build_datatree_from_grib(
 772            filename_or_obj,
 773            file_index,
 774            filters,
 775            stack_vertical=stack_vertical,
 776            drop_variables=drop_variables,
 777            chunks=chunks,
 778        )
 779
 780        # Update history for provenance
 781        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
 782        history = f"{now}: Initialized via grib2io.open_datatree\n"
 783
 784        def _add_history(node):
 785            if node.ds is not None:
 786                node.ds.attrs["history"] = history + node.ds.attrs.get("history", "")
 787            for child in node.children.values():
 788                _add_history(child)
 789
 790        _add_history(tree)
 791
 792        # Put warning here so it is the last message from likely other Xarray warnings.
 793        warnings.warn(
 794            "grib2io’s xarray backend DataTree support is experimental. The DataTree structure or attributes may change in future releases.",
 795            UserWarning,
 796            stacklevel=2,
 797        )
 798
 799        return tree
 800
 801
 802class GribBackendArray(BackendArray):
 803    """
 804    BackendArray implementation for GRIB2 data.
 805    """
 806
 807    def __init__(self, array: "OnDiskArray", lock: SerializableLock):
 808        """
 809        Initialize the GribBackendArray.
 810
 811        Parameters
 812        ----------
 813        array : OnDiskArray
 814            The on-disk array object.
 815        lock : SerializableLock
 816            The lock to use for thread-safe access.
 817        """
 818        self.array = array
 819        self.shape = array.shape
 820        self.dtype = np.dtype(array.dtype)
 821        self.lock = lock
 822
 823    def __getitem__(self, key: xr.core.indexing.ExplicitIndexer) -> np.typing.ArrayLike:
 824        return xr.core.indexing.explicit_indexing_adapter(
 825            key,
 826            self.shape,
 827            indexing.IndexingSupport.BASIC,
 828            self._raw_getitem,
 829        )
 830
 831    def _raw_getitem(self, key: tuple) -> np.ndarray:
 832        """
 833        Implement thread-safe access to data on disk.
 834
 835        Parameters
 836        ----------
 837        key : tuple
 838            The indexing key.
 839
 840        Returns
 841        -------
 842        np.ndarray
 843            The indexed array.
 844        """
 845        with self.lock:
 846            return self.array[key]
 847
 848
 849class Grid:
 850    def __new__(cls, section3):
 851        gdtn = section3[4]
 852        Gdt = templates.gdt_class_by_gdtn(gdtn)
 853
 854        @dataclass
 855        class _Grid(Gdt):
 856            section3: NDArray = field(init=True, repr=True)
 857            # Section 3 looked up common attributes.  Other looked up attributes are available according
 858            # to the Grid Definition Template.
 859            gridDefinitionSection: NDArray = field(init=False, repr=False, default=templates.GridDefinitionSection())
 860            sourceOfGridDefinition: int = field(init=False, repr=False, default=templates.SourceOfGridDefinition())
 861            numberOfDataPoints: int = field(init=False, repr=False, default=templates.NumberOfDataPoints())
 862            interpretationOfListOfNumbers: templates.Grib2Metadata = field(
 863                init=False,
 864                repr=False,
 865                default=templates.InterpretationOfListOfNumbers(),
 866            )
 867            gridDefinitionTemplateNumber: templates.Grib2Metadata = field(init=False, repr=False, default=templates.GridDefinitionTemplateNumber())
 868            gridDefinitionTemplate: list = field(init=False, repr=False, default=templates.GridDefinitionTemplate())
 869            _earthparams: dict = field(init=False, repr=False, default=templates.EarthParams())
 870            _dxsign: float = field(init=False, repr=False, default=templates.DxSign())
 871            _dysign: float = field(init=False, repr=False, default=templates.DySign())
 872            _llscalefactor: float = field(init=False, repr=False, default=templates.LLScaleFactor())
 873            _lldivisor: float = field(init=False, repr=False, default=templates.LLDivisor())
 874            _xydivisor: float = field(init=False, repr=False, default=templates.XYDivisor())
 875            shapeOfEarth: templates.Grib2Metadata = field(init=False, repr=False, default=templates.ShapeOfEarth())
 876            earthShape: str = field(init=False, repr=False, default=templates.EarthShape())
 877            earthRadius: float = field(init=False, repr=False, default=templates.EarthRadius())
 878            earthMajorAxis: float = field(init=False, repr=False, default=templates.EarthMajorAxis())
 879            earthMinorAxis: float = field(init=False, repr=False, default=templates.EarthMinorAxis())
 880            resolutionAndComponentFlags: list = field(init=False, repr=False, default=templates.ResolutionAndComponentFlags())
 881            ny: int = field(init=False, repr=False, default=templates.Ny())
 882            nx: int = field(init=False, repr=False, default=templates.Nx())
 883            scanModeFlags: list = field(init=False, repr=False, default=templates.ScanModeFlags())
 884            projParameters: dict = field(init=False, repr=False, default=templates.ProjParameters())
 885
 886            def __post_init__(self):
 887                self.gdtn = self.section3[4]
 888
 889        grid = _Grid(section3)
 890        return grid
 891
 892
 893def exclusive_slice_to_inclusive(item: slice):
 894    """
 895    Convert a slice with exclusive stop to an inclusive slice.
 896
 897    If the slice has a step, the stop is reduced by the step, so that both
 898    interpretations would yield the same result.
 899
 900    The means that [start, stop) is converted to [start, stop - step].
 901
 902    Parameters
 903    ----------
 904    item
 905        The slice to convert.
 906
 907    Returns
 908    -------
 909    slice
 910        The converted slice.
 911    """
 912    # return the None slice
 913    if item.start is None and item.stop is None and item.step is None:
 914        return item
 915    if not isinstance(item, slice):
 916        raise ValueError(f"item must be a slice; it was of type {type(item)}")
 917    # if step is None, it's one
 918    step = 1 if item.step is None else item.step
 919    if item.stop < item.start or step < 1:
 920        raise ValueError(f"slice {item} not accounted for")
 921    # handle case where slice has one item
 922    if abs(item.stop - item.start) == step:
 923        return [item.start]
 924    # other cases require reducing the stop by the step
 925    s = slice(item.start, item.stop - step, step)
 926    return s
 927
 928
 929class Validator:
 930    def __set_name__(self, owner, name):
 931        self.private_name = f"_{name}"
 932        self.name = name
 933
 934    def __get__(self, obj, objtype=None):
 935        try:
 936            value = getattr(obj, self.private_name)
 937        except AttributeError:
 938            value = None
 939        return value
 940
 941
 942class PdIndex(Validator):
 943    def __set__(self, obj, value):
 944        try:
 945            value = pd.Index(value)
 946        except TypeError:
 947            value = pd.Index([value])
 948        setattr(obj, self.private_name, value)
 949
 950
 951def _asarray_tuplesafe(values: typing.Any) -> np.ndarray:
 952    """
 953    Convert values to a numpy array of at most 1-dimension and preserve tuples.
 954
 955    Adapted from ``pandas.core.common._asarray_tuplesafe``.
 956
 957    Parameters
 958    ----------
 959    values : any
 960        The values to convert.
 961
 962    Returns
 963    -------
 964    np.ndarray
 965        The converted numpy array.
 966    """
 967    if isinstance(values, tuple):
 968        result = np.empty(1, dtype=object)
 969        result[0] = values
 970    else:
 971        result = np.asarray(values)
 972        if result.ndim == 2:
 973            result = np.empty(len(values), dtype=object)
 974            result[:] = values
 975
 976    return result
 977
 978
 979def array_safe_eq(a: typing.Any, b: typing.Any) -> bool:
 980    """
 981    Check if a and b are equal, even if they are numpy arrays.
 982
 983    Parameters
 984    ----------
 985    a : any
 986        First object to compare.
 987    b : any
 988        Second object to compare.
 989
 990    Returns
 991    -------
 992    bool
 993        True if equal, False otherwise.
 994    """
 995    if a is b:
 996        return True
 997    if hasattr(a, "equals"):
 998        return a.equals(b)
 999    if hasattr(a, "all") and hasattr(b, "all"):
1000        return a.shape == b.shape and (a == b).all()
1001    if hasattr(a, "all") or hasattr(b, "all"):
1002        return False
1003    try:
1004        return a == b
1005    except TypeError:
1006        return NotImplementedError
1007
1008
1009def dc_eq(dc1: typing.Any, dc2: typing.Any) -> bool:
1010    """
1011    Check if two dataclasses which hold numpy arrays are equal.
1012
1013    Parameters
1014    ----------
1015    dc1 : any
1016        First dataclass to compare.
1017    dc2 : any
1018        Second dataclass to compare.
1019
1020    Returns
1021    -------
1022    bool
1023        True if equal, False otherwise.
1024    """
1025    if dc1 is dc2:
1026        return True
1027    if dc1.__class__ is not dc2.__class__:
1028        return NotImplementedError
1029    t1 = astuple(dc1)
1030    t2 = astuple(dc2)
1031    return all(array_safe_eq(a1, a2) for a1, a2 in zip(t1, t2))
1032
1033
1034def coords_from_cube(cube: dict) -> typing.Dict[str, xr.Variable]:
1035    """
1036    Create a dictionary of xarray Variables from a cube definition.
1037
1038    Parameters
1039    ----------
1040    cube : dict
1041        Dimension cube definition.
1042
1043    Returns
1044    -------
1045    dict of str to xarray.Variable
1046        Coordinates for the Dataset/DataArray.
1047    """
1048    keys = list(cube.keys())
1049    keys.remove("x")
1050    keys.remove("y")
1051    coords = dict()
1052    for k in keys:
1053        if k is not None:
1054            if len(cube[k]) > 1:
1055                coords[k] = xr.Variable(dims=k, data=cube[k], attrs=dict(grib_name=k))
1056            elif len(cube[k]) == 1:
1057                coords[k] = xr.Variable(dims=tuple(), data=cube[k][0], attrs=dict(grib_name=k))
1058    return coords
1059
1060
1061@dataclass
1062class OnDiskArray:
1063    """
1064    On-disk array representation for GRIB2 messages.
1065    """
1066
1067    file_name: typing.Union[str, typing.List[str]]
1068    index: pd.DataFrame = field(repr=False)
1069    cube: dict = field(repr=False)
1070    shape: typing.Tuple[int, ...] = field(init=False)
1071    ndim: int = field(init=False)
1072    geo_ndim: int = field(init=False)
1073    dtype: str = "float32"
1074
1075    def __post_init__(self):
1076        # multiple grids not allowed so can just use first
1077        geo_shape = (self.index.iloc[0].ny, self.index.iloc[0].nx)
1078
1079        self.geo_shape = geo_shape
1080        self.geo_ndim = len(geo_shape)
1081
1082        if len(self.index) == 1:
1083            self.shape = geo_shape
1084        else:
1085            if self.index.index.nlevels == 1:
1086                self.shape = tuple([len(self.index.index)]) + geo_shape
1087            else:
1088                self.shape = tuple([len(i) for i in self.index.index.levels]) + geo_shape
1089        self.ndim = len(self.shape)
1090
1091        cols = ["msg", "sectionOffset"]
1092        if "file_index" in self.index.columns:
1093            cols.append("file_index")
1094        self.index = self.index[cols]
1095
1096    def __getitem__(self, item: tuple) -> np.ndarray:
1097        """
1098        Retrieve data from disk for the specified slices.
1099
1100        Parameters
1101        ----------
1102        item : tuple
1103            The slicing tuple.
1104
1105        Returns
1106        -------
1107        np.ndarray
1108            The retrieved data array.
1109        """
1110        # dimensions not in index are internal to tdlpack records; 2 dims for
1111        # grids; 1 dim for stations
1112
1113        index_slicer = item[: -self.geo_ndim]
1114        # maintain all multindex levels
1115        index_slicer = tuple([[i] if isinstance(i, int) else i for i in index_slicer])
1116
1117        # pandas loc slicing is inclusive, therefore convert slices into
1118        # explicit lists
1119        index_slicer_inclusive = tuple([exclusive_slice_to_inclusive(i) if isinstance(i, slice) else i for i in index_slicer])
1120
1121        # get records selected by item in new index dataframe
1122        if len(index_slicer_inclusive) == 1:
1123            index = self.index.loc[index_slicer_inclusive]
1124        elif len(index_slicer_inclusive) > 1:
1125            index = self.index.loc[index_slicer_inclusive, :]
1126        else:
1127            index = self.index
1128        index = index.set_index(index.index)
1129
1130        # set miloc to new relative locations in sub array
1131        index["miloc"] = list(zip(*[index.index.unique(level=dim).get_indexer(index.index.get_level_values(dim)) for dim in index.index.names]))
1132
1133        if len(index_slicer_inclusive) == 1:
1134            array_field_shape = tuple([len(index.index)]) + self.geo_shape
1135        elif len(index_slicer_inclusive) > 1:
1136            array_field_shape = index.index.levshape + self.geo_shape
1137        else:
1138            array_field_shape = self.geo_shape
1139
1140        array_field = np.full(array_field_shape, fill_value=np.nan, dtype="float32")
1141
1142        if "file_index" in index.columns:
1143            for file_idx, group in index.groupby("file_index"):
1144                filename = self.file_name[file_idx] if isinstance(self.file_name, list) else self.file_name
1145                with open(filename, mode="rb") as filehandle:
1146                    for key, row in group.iterrows():
1147                        bitmap_offset = None if pd.isna(row["sectionOffset"][6]) else int(row["sectionOffset"][6])
1148                        values = _data(filehandle, row.msg, bitmap_offset, row["sectionOffset"][7])
1149
1150                        if len(index_slicer_inclusive) >= 1:
1151                            array_field[row.miloc] = values
1152                        else:
1153                            array_field = values
1154        else:
1155            with open(self.file_name, mode="rb") as filehandle:
1156                for key, row in index.iterrows():
1157                    bitmap_offset = None if pd.isna(row["sectionOffset"][6]) else int(row["sectionOffset"][6])
1158                    values = _data(filehandle, row.msg, bitmap_offset, row["sectionOffset"][7])
1159
1160                    if len(index_slicer_inclusive) >= 1:
1161                        array_field[row.miloc] = values
1162                    else:
1163                        array_field = values
1164
1165        # handle geo dim slicing
1166        array_field = array_field[(Ellipsis,) + item[-self.geo_ndim :]]
1167
1168        # squeeze array dimensions expressed as integer
1169        for i, it in reversed(list(enumerate(item[: -self.geo_ndim]))):
1170            if isinstance(it, int):
1171                array_field = array_field[(slice(None, None, None),) * i + (0,)]
1172
1173        return array_field
1174
1175
1176def dims_to_shape(d: dict) -> tuple:
1177    """
1178    Convert dimension metadata to a shape tuple.
1179
1180    Parameters
1181    ----------
1182    d : dict
1183        Dimension metadata dictionary.
1184
1185    Returns
1186    -------
1187    tuple
1188        Shape tuple.
1189    """
1190    if "nx" in d:
1191        t = (d["ny"], d["nx"])
1192    else:
1193        t = (d["nsta"],)
1194    return t
1195
1196
1197def filter_index(index: pd.DataFrame, k: str, v: typing.Any) -> pd.DataFrame:
1198    """
1199    Filter a GRIB2 index DataFrame by a key-value pair.
1200
1201    Supports slice and vectorized-indexing similar to xarray's ``sel``.
1202
1203    Parameters
1204    ----------
1205    index : pandas.DataFrame
1206        The GRIB2 index DataFrame to filter.
1207    k : str
1208        Column name to filter by.
1209    v : any
1210        Value(s) or slice to filter for.
1211
1212    Returns
1213    -------
1214    pandas.DataFrame
1215        Filtered index.
1216    """
1217    if isinstance(v, slice):
1218        index = index.set_index(k)
1219        index = index.loc[v]
1220        index = index.reset_index()
1221    else:
1222        label = (
1223            v
1224            if getattr(v, "ndim", 1) > 1  # vectorized-indexing
1225            else _asarray_tuplesafe(v)
1226        )
1227        if label.ndim == 0:
1228            # see https://github.com/pydata/xarray/pull/4292 for details
1229            label_value = label[()] if label.dtype.kind in "mM" else label.item()
1230            try:
1231                indexer = pd.Index(index[k]).get_loc(label_value)
1232                if isinstance(indexer, int):
1233                    index = index.iloc[[indexer]]
1234                else:
1235                    index = index.iloc[indexer]
1236            except KeyError:
1237                index = index.iloc[[]]
1238        else:
1239            indexer = pd.Index(index[k]).get_indexer_for(np.ravel(v))
1240            index = index.iloc[indexer[indexer >= 0]]
1241
1242    return index
1243
1244
1245def parse_grib_index(
1246    index: pd.DataFrame,
1247    filters: typing.Mapping[str, typing.Any] = dict(),
1248) -> typing.Tuple[pd.DataFrame, typing.Dict[str, typing.List[str]], dict, typing.Dict[str, dict]]:
1249    """
1250    Apply filters.
1251
1252    Evaluate remaining dimensions based on pdtn and parse each out.
1253
1254    Parameters
1255    ----------
1256    index
1257        Pandas DataFrame containing the GRIB2 message index.
1258    filters
1259        Filter GRIB2 messages to single hypercube. Dict keys can be any
1260        GRIB2 metadata attribute name.
1261
1262    Returns
1263    -------
1264    index
1265        Modified Pandas DataFrame with added GRIB2 metadata columns.
1266    dim_coords
1267        List of GRIB2 attributes that will be used for coordinates and/or dimensions.
1268    attrs
1269        Dict of metadata attributes (non-coordinates, non-geo)
1270    """
1271
1272    # make a copy of filters, remove filters as they are applied
1273    filters = copy(filters)
1274
1275    for k, v in filters.items():
1276        if k not in index.columns:
1277            kwarg = {k: index.msg.apply(lambda msg: getattr(msg, k))}
1278            index = index.assign(**kwarg)
1279        # adopt parts of xarray's sel logic  so that filters behave similarly
1280        # allowed to filter to nothing to make empty dataset
1281        index = filter_index(index, k, v)
1282
1283    if len(index) == 0:
1284        return index, list(), dict(), dict()
1285
1286    dim_coords = dict()  # key=name of dim, value=list of coord names
1287    attrs = dict()
1288    coord_attrs = dict()
1289
1290    # expand index
1291    index = index.assign(shortName=index.msg.apply(lambda msg: msg.shortName))
1292    index = index.assign(nx=index.msg.apply(lambda msg: msg.nx))
1293    index = index.assign(ny=index.msg.apply(lambda msg: msg.ny))
1294    index = index.astype({"ny": "int", "nx": "int"})
1295
1296    # apply common filters(to all definition templates) to reduce dataset to
1297    # single cube
1298    # ensure only one of each of the below exists after filters applied
1299    required_uniques = [
1300        "productDefinitionTemplateNumber",
1301        "typeOfGeneratingProcess",
1302        "typeOfFirstFixedSurface",
1303        "typeOfSecondFixedSurface",
1304    ]
1305
1306    def meta_check(index, attrs, meta):
1307        """
1308        add meta to the datframe index
1309        check that there is a single type
1310        add the type to attrs
1311
1312        returns index, attrs
1313        """
1314        index = index.assign(**{meta: index.msg.apply(lambda msg: getattr(msg, meta))})
1315
1316        unique = index[meta].unique()
1317        if len(index[meta].unique()) > 1:
1318            raise ValueError(f"filter to a single {meta}; found: {[str(i) for i in unique]}")
1319        value = unique.item()
1320        if isinstance(value, grib2io.templates.Grib2Metadata):
1321            value = value.definition
1322
1323        # None is returned if no value found,
1324        # check and change to string None
1325        if value is None:
1326            value = "None"
1327
1328        attrs[meta] = value
1329        return index, attrs
1330
1331    for meta in required_uniques:
1332        index, attrs = meta_check(index, attrs, meta)
1333
1334    pdtn = index.productDefinitionTemplateNumber.iloc[0].value
1335
1336    # determine which non geo dimensions can be created from data by this point
1337    # the index is filtered down to a single type for all required_uniques
1338
1339    # Dim Name     # matching dim_name for using this data as index coordinate
1340    dim_coords["refDate"] = ["refDate"]
1341    coord_attrs["refDate"] = dict(standard_name="forecast_reference_time")
1342    #   dim_coords["refDate"] = ["refDate", "hour"] # non dim name matching items in list are used as non-index coordinates
1343
1344    dim_coords["leadTime"] = ["leadTime"]
1345    coord_attrs["leadTime"] = dict(standard_name="forecast_period")
1346
1347    if "valueOfFirstFixedSurface" not in index.columns:
1348        index = index.assign(valueOfFirstFixedSurface=index.msg.apply(lambda msg: msg.valueOfFirstFixedSurface))
1349    if "valueOfsecondFixedSurface" not in index.columns:
1350        index = index.assign(valueOfSecondFixedSurface=index.msg.apply(lambda msg: msg.valueOfSecondFixedSurface))
1351
1352    # dim name api change, user could run ds = ds.swap_dims(fixedSurface="valueOfFirstFixedSurface")
1353    index = index.assign(level=list(zip(index["valueOfFirstFixedSurface"], index["valueOfSecondFixedSurface"])))
1354    #   index = index.assign(level=index.msg.apply(lambda msg: msg.level))
1355    # lack of "level" indeicates don't create extra index coordinate "level"
1356    dim_coords["level"] = ["valueOfFirstFixedSurface", "valueOfSecondFixedSurface"]
1357
1358    # logic for parsing possible dims from specific product definition section
1359
1360    if pdtn in {5, 9}:
1361        # Probability forecasts at a horizontal level or in a horizontal layer
1362        # in a continuous or non-continuous time interval.  (see Template
1363        # 4.9)
1364        #       AVAILABLE_THRESHOLD = {
1365        #           0: {'has_lower': True, 'has_upper': False},
1366        #           1: {'has_lower': False, 'has_upper': True},
1367        #           2: {'has_lower': True, 'has_upper': True},
1368        #           3: {'has_lower': True, 'has_upper': False},
1369        #           4: {'has_lower': False, 'has_upper': True},
1370        #           5: {'has_lower': True, 'has_upper': False},
1371        #       }
1372
1373        index, attrs = meta_check(index, attrs, "typeOfProbability")
1374        if "thresholdLowerLimit" not in index.columns:
1375            index = index.assign(thresholdLowerLimit=index.msg.apply(lambda msg: msg.thresholdLowerLimit))
1376        if "thresholdUpperLimit" not in index.columns:
1377            index = index.assign(thresholdUpperLimit=index.msg.apply(lambda msg: msg.thresholdUpperLimit))
1378        if "threshold" not in index.columns:
1379            # 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
1380            index = index.assign(threshold=list(zip(index["thresholdLowerLimit"], index["thresholdUpperLimit"])))
1381        #           index = index.assign(threshold = index.msg.apply(lambda msg: msg.threshold))
1382
1383        # ommiting threshold results in no index being assigned for this possible dim
1384        dim_coords["threshold"] = ["thresholdLowerLimit", "thresholdUpperLimit"]
1385
1386    if pdtn in {6, 10}:
1387        # Percentile forecasts at a horizontal level or in a horizontal layer
1388        # in a continuous or non-continuous time interval.  (see Template
1389        # 4.10)
1390        dim_coords["percentileValue"] = ["percentileValue"]
1391        coord_attrs["percentileValue"] = dict(long_name="percentile", units="percent")
1392
1393    if pdtn in {
1394        8,
1395        9,
1396        10,
1397        11,
1398        12,
1399        13,
1400        14,
1401        42,
1402        43,
1403        45,
1404        46,
1405        47,
1406        61,
1407        62,
1408        63,
1409        67,
1410        68,
1411        72,
1412        73,
1413        78,
1414        79,
1415        82,
1416        83,
1417        84,
1418        85,
1419        87,
1420        91,
1421    }:
1422        dim_coords["duration"] = ["duration"]
1423
1424    if pdtn in {
1425        1,
1426        11,
1427        33,
1428        34,
1429        41,
1430        43,
1431        45,
1432        47,
1433        49,
1434        54,
1435        56,
1436        58,
1437        59,
1438        63,
1439        68,
1440        77,
1441        79,
1442        81,
1443        83,
1444        84,
1445        85,
1446        92,
1447    }:
1448        dim_coords["perturbationNumber"] = ["perturbationNumber"]
1449
1450    if pdtn in {2, 3, 4, 12, 13, 14}:
1451        index, attrs = meta_check(index, attrs, "typeOfDerivedForecast")
1452
1453    if pdtn in {5, 9}:
1454        dim_coords["typeOfProbability"] = ["typeOfProbability"]
1455
1456    if pdtn in {6, 10}:
1457        dim_coords["percentileValue"] = ["percentileValue"]
1458
1459    if pdtn in {8, 15, 42, 46, 62, 67, 72, 78, 82, 1001, 1002, 1100, 1101}:
1460        index, attrs = meta_check(index, attrs, "statisticalProcess")
1461
1462    # Logic for Trace Gas and Aerosol dimensions
1463    if pdtn in {40, 41, 42, 43, 76, 77, 78, 79}:
1464        dim_coords["constituentType"] = ["constituentType"]
1465
1466    if pdtn in {76, 77, 78, 79}:
1467        dim_coords["sourceSinkIndicator"] = ["sourceSinkIndicator"]
1468
1469    if pdtn in {44, 45, 46, 47, 48, 49, 50, 80, 81, 82, 83, 84, 85}:
1470        dim_coords["typeOfAerosol"] = ["typeOfAerosol"]
1471
1472    if pdtn in {80, 81, 82, 83, 84}:
1473        dim_coords["sourceSinkIndicator"] = ["sourceSinkIndicator"]
1474
1475    if pdtn in {48, 49, 80, 81}:
1476        dim_coords["firstWavelength"] = ["firstWavelength"]
1477        dim_coords["secondWavelength"] = ["secondWavelength"]
1478        dim_coords["firstSizeOfAerosol"] = ["firstSizeOfAerosol"]
1479        dim_coords["secondSizeOfAerosol"] = ["secondSizeOfAerosol"]
1480
1481    # Finish logic by pdtn
1482
1483    for k, v in dim_coords.items():
1484        for meta in v:
1485            if meta not in index.columns:
1486                index = index.assign(**{meta: index.msg.apply(lambda msg: getattr(msg, meta))})
1487
1488    return index, dim_coords, attrs, coord_attrs
1489
1490
1491# Custom open_datatree function to open grib files as DataTree
1492def open_datatree(
1493    filename: str,
1494    *,
1495    drop_variables: typing.Optional[typing.List[str]] = None,
1496    filters: typing.Optional[typing.Mapping[str, typing.Any]] = None,
1497    engine: str = "grib2io",
1498    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
1499    **kwargs,
1500) -> typing.Any:
1501    """
1502    Open a GRIB2 file as an xarray DataTree.
1503
1504    Parameters
1505    ----------
1506    filename : str
1507        Path to the GRIB2 file.
1508    drop_variables : list, optional
1509        List of variables to exclude.
1510    filters : dict, optional
1511        Filter criteria for GRIB2 messages.
1512    engine : str, optional
1513        Engine to use for opening the file, defaults to "grib2io".
1514    chunks : int, dict or 'auto', optional
1515        If chunks is provided, it is used to load the dataset into a
1516        dask-backed dataset.
1517    **kwargs : optional
1518        Additional keyword arguments passed to the xarray backend.
1519
1520    Returns
1521    -------
1522    xarray.DataTree
1523        A hierarchical DataTree representation of the GRIB2 data.
1524    """
1525    if not _HAS_DATATREE:
1526        raise ImportError("xarray version does not support DataTree functionality.")
1527
1528    if filters is None:
1529        filters = {}
1530
1531    # Open the file without any filters first to get all messages
1532    with grib2io.open(filename, _xarray_backend=True) as f:
1533        file_index = pd.DataFrame(f._index)
1534        file_index = file_index.assign(msg=msgs_from_index(f._index))
1535
1536    # Build tree structure from GRIB messages
1537    root = build_datatree_from_grib(
1538        filename,
1539        file_index,
1540        filters,
1541        drop_variables=drop_variables,
1542        chunks=chunks,
1543    )
1544
1545    # Update history for provenance
1546    now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
1547    existing_history = root.attrs.get("history", "") if hasattr(root, "attrs") else ""
1548    history = f"{now}: Initialized via grib2io.open_datatree from {filename}\n{existing_history}"
1549    if hasattr(root, "attrs"):
1550        root.attrs["history"] = history
1551    # Also add to all datasets in the tree
1552    for node in root.subtree:
1553        if node.ds is not None:
1554            node.ds.attrs["history"] = history + node.ds.attrs.get("history", "")
1555
1556    return root
1557
1558    return root
1559
1560
1561def build_da_without_coords(index: pd.DataFrame, cube: dict, filename: str, attrs: dict) -> xr.DataArray:
1562    """
1563    Build a DataArray without coordinates from a cube of grib2 messages.
1564
1565    Parameters
1566    ----------
1567    index : pd.DataFrame
1568        Index of cube.
1569    cube : dict
1570        Cube of grib2 messages.
1571    filename : str
1572        Filename of grib2 file.
1573    attrs : dict
1574        Attributes for the DataArray.
1575
1576    Returns
1577    -------
1578    xr.DataArray
1579        DataArray without coordinates.
1580    """
1581
1582    dim_names = [k for k in cube.keys() if cube[k] is not None and len(cube[k]) > 1]
1583    constant_meta_names = [k for k in cube.keys() if cube[k] is None]
1584    dims = {k: len(cube[k]) for k in dim_names}
1585
1586    # guard against bad datarrays being formed
1587    dims_total = 1
1588    dims_to_filter = []
1589    for (
1590        dim_name,
1591        dim_len,
1592    ) in dims.items():
1593        if dim_name not in {"x", "y", "station"}:
1594            dims_total *= dim_len
1595            dims_to_filter.append(dim_name)
1596
1597    # Check number of GRIB2 message indexed compared to non-X/Y
1598    # dimensions.
1599    if dims_total != len(index):
1600        raise ValueError(
1601            f"DataArray dimensions are not compatible with number of GRIB2 messages; DataArray has {dims_total} "
1602            f"and GRIB2 index has {len(index)}. Consider applying a filter for dimensions: {dims_to_filter}"
1603        )
1604
1605    data = OnDiskArray(filename, index, cube)
1606    lock = _LOCK
1607    data = GribBackendArray(data, lock)
1608    data = indexing.LazilyIndexedArray(data)
1609    if len(dim_names) != len(data.shape):
1610        raise ValueError(
1611            "different number of dimensions on data "
1612            f"and dims: {len(data.shape)} vs {len(dim_names)}\n"
1613            "Grib2 messages could not be formed into a data cube; "
1614            "It's possible extra messages exist along a non-accounted for dimension based on PDTN\n"
1615            "It might be possible to get around this by applying a filter on the non-accounted for dimension"
1616        )
1617    da = xr.DataArray(data, dims=dim_names)
1618
1619    da.encoding["original_shape"] = data.shape
1620
1621    da.encoding["preferred_chunks"] = {"y": -1, "x": -1}
1622    msg1 = index.msg.iloc[0]
1623
1624    # plain language metadata is minimized
1625    # add grib section metadata
1626    da.attrs["GRIB2IO_section0"] = msg1.section0
1627    da.attrs["GRIB2IO_section1"] = msg1.section1
1628    da.attrs["GRIB2IO_section2"] = msg1.section2 if msg1.section2 else []
1629    da.attrs["GRIB2IO_section3"] = msg1.section3
1630    da.attrs["GRIB2IO_section4"] = msg1.section4
1631    da.attrs["GRIB2IO_section5"] = msg1.section5
1632    da.attrs["fullName"] = str(msg1.fullName)
1633    da.attrs["shortName"] = str(msg1.shortName)
1634    da.attrs["units"] = str(msg1.units)
1635    da.attrs["originatingCenter"] = str(msg1.originatingCenter.definition)
1636    da.attrs["originatingSubCenter"] = str(msg1.originatingSubCenter.definition)
1637
1638    # add master table
1639    da.attrs["masterTableInfo"] = str(msg1.masterTableInfo.definition)
1640
1641    da.name = index.shortName.iloc[0]
1642    for meta_name in constant_meta_names:
1643        if meta_name in index.columns:
1644            da.attrs[meta_name] = index[meta_name].iloc[0]
1645
1646    da.attrs.update(attrs)
1647
1648    return da
1649
1650
1651def assign_xr_meta(
1652    ds: xr.Dataset,
1653    frames: typing.List[pd.DataFrame],
1654    cube: dict,
1655    non_geo_dims: typing.Dict[str, typing.List[str]],
1656    extra_geo: dict,
1657    coord_attrs: typing.Dict[str, dict],
1658) -> xr.Dataset:
1659    """
1660    Assign coordinates and attributes to the dataset.
1661
1662    Parameters
1663    ----------
1664    ds : xr.Dataset
1665        The dataset to update.
1666    frames : list of pd.DataFrame
1667        The dataframes for each variable.
1668    cube : dict
1669        The dimensions cube.
1670    non_geo_dims : dict
1671        The non-geographic dimensions.
1672    extra_geo : dict
1673        Extra geographic coordinates.
1674    coord_attrs : dict
1675        Attributes for coordinates.
1676
1677    Returns
1678    -------
1679    xr.Dataset
1680        The updated dataset.
1681    """
1682    df = frames[0]
1683
1684    # assign extra geo coords
1685    ds = ds.assign_coords(extra_geo)
1686    # add crs data from first grib message to each data variable and the dataset
1687    geo_attrs = {
1688        "crs_wkt": CRS.from_dict(df.msg.iloc[0].projParameters).to_wkt(),
1689        "gridlengthXDirection": df.msg.iloc[0].gridlengthXDirection,
1690        "gridlengthYDirection": df.msg.iloc[0].gridlengthYDirection,
1691        "latitudeFirstGridpoint": df.msg.iloc[0].latitudeFirstGridpoint,
1692        "longitudeFirstGridpoint": df.msg.iloc[0].longitudeFirstGridpoint,
1693    }
1694    for data_var in ds.data_vars:
1695        ds[data_var].attrs.update(geo_attrs)
1696    ds.attrs.update(geo_attrs)
1697
1698    # add coordinate specific attributes
1699    for coord, attrs in coord_attrs.items():
1700        ds[coord].attrs.update(attrs)
1701
1702    # assign valid date coords
1703    try:
1704        ds = ds.assign_coords(dict(validDate=ds.coords["refDate"] + ds.coords["leadTime"]))
1705        ds.validDate.attrs["standard_name"] = "time"
1706        ds.validDate.attrs["long_name"] = "time"
1707    except Exception as e:
1708        warnings.warn(f"could not parse validTime: {e}")
1709
1710    # assign attributes
1711    ds.attrs["engine"] = "grib2io"
1712
1713    return ds
1714
1715
1716def make_variables(
1717    index: pd.DataFrame,
1718    f: str,
1719    non_geo_dims: typing.Dict[str, typing.List[str]],
1720    allow_uneven_dims: bool = False,
1721) -> typing.Tuple[
1722    typing.Optional[typing.List[pd.DataFrame]],
1723    typing.Optional[typing.List[dict]],
1724    typing.Optional[dict],
1725]:
1726    """
1727    Create an individual dataframe index and cube for each variable.
1728
1729    Parameters
1730    ----------
1731    index : pd.DataFrame
1732        Index of messages.
1733    f : str
1734        Filename.
1735    non_geo_dims : dict
1736        Dimensions not associated with the x,y grid.
1737    allow_uneven_dims : bool, optional
1738        If True, allows uneven dimensions (used for DataTree creation).
1739
1740    Returns
1741    -------
1742    ordered_frames : list of pd.DataFrame, optional
1743        List of dataframes, one for each variable.
1744    cubes : list of dict, optional
1745        List of cubes, one for each variable.
1746    extra_geo : dict, optional
1747        Extra geographic coordinates.
1748    """
1749    # let shortName determine the variables
1750
1751    # set the index to the name
1752    index = index.set_index("shortName").sort_index()
1753    # return nothing if no data
1754    if index.empty:
1755        return None, None, None
1756
1757    # define the DimCube
1758    dims = copy(non_geo_dims)
1759
1760    ordered_meta = list(non_geo_dims.keys())
1761    cubes = list()
1762    ordered_frames = list()
1763    for key in index.index.unique():
1764        frame = index.loc[[key]]
1765        frame = frame.reset_index()
1766        # frame is a dataframe with all records for one variable
1767        c = dict()
1768        # for colname in frame.columns:
1769        for colname in ordered_meta:
1770            uniques = pd.Index(frame[colname]).unique()
1771            if len(uniques) > 1:
1772                c[colname] = uniques.sort_values()
1773            else:
1774                c[colname] = [uniques[0]]
1775
1776        dims = [k for k in ordered_meta if len(c[k]) > 1]
1777
1778        for dim in dims:
1779            if frame[dim].value_counts().nunique() > 1 and not allow_uneven_dims:
1780                raise ValueError(f"uneven number of grib msgs associated with dimension: {dim}\n unique values for {dim}: {frame[dim].unique()} ")
1781
1782        if len(dims) >= 1:  # dims may be empty if no extra dims on top of x,y
1783            frame = frame.sort_values(dims)
1784            frame = frame.set_index(dims)
1785
1786        cubes.append(c)
1787
1788        # miloc is multi-index integer location of msg in nd DataArray
1789        miloc = list(zip(*[frame.index.unique(level=dim).get_indexer(frame.index.get_level_values(dim)) for dim in dims]))
1790
1791        # set frame multi index
1792        if len(miloc) >= 1:  # miloc will be empty when no extra dims, thus no multiindex
1793            dim_ix = tuple([n + "_ix" for n in dims])
1794            frame = frame.set_index(pd.MultiIndex.from_tuples(miloc, names=dim_ix))
1795
1796        ordered_frames.append(frame)
1797
1798    # no variables
1799    if not cubes:
1800        cubes = [dict()]
1801
1802    # check geography of data and assign to cube
1803    if len(index.ny.unique()) > 1 or len(index.nx.unique()) > 1:
1804        raise ValueError("multiple grids not accommodated")
1805    for cube in cubes:
1806        cube["y"] = range(int(index.ny.iloc[0]))
1807        cube["x"] = range(int(index.nx.iloc[0]))
1808
1809    extra_geo = None
1810    msg = index.msg.iloc[0]
1811
1812    # we want the lat lons; make them via accessing a record; we are assuming
1813    # all records are the same grid because they have the same shape;
1814    # may want a unique grid identifier from grib2io to avoid assuming this
1815    latitude, longitude = msg.latlons()
1816    latitude = xr.DataArray(latitude, dims=["y", "x"])
1817    latitude.attrs["standard_name"] = "latitude"
1818    latitude.attrs["units"] = "degrees_north"
1819    longitude = xr.DataArray(longitude, dims=["y", "x"])
1820    longitude.attrs["standard_name"] = "longitude"
1821    longitude.attrs["units"] = "degrees_east"
1822    extra_geo = dict(latitude=latitude, longitude=longitude)
1823
1824    return ordered_frames, cubes, extra_geo
1825
1826
1827def interp_nd(
1828    a: np.ndarray,
1829    *,
1830    method: typing.Union[str, int],
1831    grid_def_in: grib2io.Grib2GridDef,
1832    grid_def_out: grib2io.Grib2GridDef,
1833    method_options: typing.Optional[typing.List[int]] = None,
1834    num_threads: int = 1,
1835) -> np.ndarray:
1836    """
1837    Perform multi-dimensional interpolation on a horizontal grid.
1838
1839    This function reshapes the input array to (N, ny, nx) before performing
1840    interpolation and then reshapes it back to its original dimensions plus
1841    the new grid dimensions.
1842
1843    Parameters
1844    ----------
1845    a : np.ndarray
1846        Input array with horizontal dimensions (..., ny, nx).
1847    method : str or int
1848        Interpolation method.
1849    grid_def_in : grib2io.Grib2GridDef
1850        Input grid definition.
1851    grid_def_out : grib2io.Grib2GridDef
1852        Output grid definition.
1853    method_options : list of int, optional
1854        Interpolation options.
1855    num_threads : int, optional
1856        Number of threads for parallel interpolation.
1857
1858    Returns
1859    -------
1860    np.ndarray
1861        Interpolated array with horizontal dimensions of the output grid.
1862    """
1863    front_shape = a.shape[:-2]
1864    a = a.reshape(-1, a.shape[-2], a.shape[-1])
1865    a = grib2io.interpolate(
1866        a,
1867        method,
1868        grid_def_in,
1869        grid_def_out,
1870        method_options=method_options,
1871        num_threads=num_threads,
1872    )
1873    a = a.reshape(front_shape + (a.shape[-2], a.shape[-1]))
1874    return a
1875
1876
1877def interp_nd_stations(
1878    a: np.ndarray,
1879    *,
1880    method: typing.Union[str, int],
1881    grid_def_in: grib2io.Grib2GridDef,
1882    lats: typing.Sequence[float],
1883    lons: typing.Sequence[float],
1884    method_options: typing.Optional[typing.List[int]] = None,
1885    num_threads: int = 1,
1886) -> np.ndarray:
1887    """
1888    Perform multi-dimensional interpolation to station points.
1889
1890    This function reshapes the input array to (N, ny, nx) before performing
1891    interpolation and then reshapes it back to its original dimensions plus
1892    the station dimension.
1893
1894    Parameters
1895    ----------
1896    a : np.ndarray
1897        Input array with horizontal dimensions (..., ny, nx).
1898    method : str or int
1899        Interpolation method.
1900    grid_def_in : grib2io.Grib2GridDef
1901        Input grid definition.
1902    lats : sequence of float
1903        Station latitudes.
1904    lons : sequence of float
1905        Station longitudes.
1906    method_options : list of int, optional
1907        Interpolation options.
1908    num_threads : int, optional
1909        Number of threads for parallel interpolation.
1910
1911    Returns
1912    -------
1913    np.ndarray
1914        Interpolated array with the last dimension representing stations.
1915    """
1916    front_shape = a.shape[:-2]
1917    a = a.reshape(-1, a.shape[-2], a.shape[-1])
1918    a = grib2io.interpolate_to_stations(
1919        a,
1920        method,
1921        grid_def_in,
1922        lats,
1923        lons,
1924        method_options=method_options,
1925        num_threads=num_threads,
1926    )
1927    a = a.reshape(front_shape + (len(lats),))
1928    return a
1929
1930
1931@xr.register_dataset_accessor("grib2io")
1932class Grib2ioDataSet:
1933    def __init__(self, xarray_obj):
1934        self._obj = xarray_obj
1935
1936    def griddef(self):
1937        return Grib2GridDef.from_section3(self._obj[list(self._obj.data_vars)[0]].attrs["GRIB2IO_section3"])
1938
1939    def interp(
1940        self,
1941        method: typing.Union[str, int],
1942        grid_def_out: grib2io.Grib2GridDef,
1943        method_options: typing.Optional[typing.List[int]] = None,
1944        num_threads: int = 1,
1945    ) -> xr.Dataset:
1946        """
1947        Perform grid spatial interpolation on all variables in the Dataset.
1948
1949        Parameters
1950        ----------
1951        method : str or int
1952            Interpolation method.
1953        grid_def_out : grib2io.Grib2GridDef
1954            Output grid definition.
1955        method_options : list of int, optional
1956            Interpolation options.
1957        num_threads : int, optional
1958            Number of threads.
1959
1960        Returns
1961        -------
1962        xarray.Dataset
1963            Interpolated dataset.
1964        """
1965        da = self._obj.to_array()
1966        da.attrs["GRIB2IO_section3"] = self._obj[list(self._obj.data_vars)[0]].attrs["GRIB2IO_section3"]
1967        da = da.grib2io.interp(method, grid_def_out, method_options=method_options, num_threads=num_threads)
1968        ds = da.to_dataset(dim="variable")
1969
1970        # Update history for provenance
1971        history = ds.attrs.get("history", "")
1972        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
1973        ds.attrs["history"] = f"{now}: Interpolated via {method} to {grid_def_out}\n{history}"
1974
1975        return ds
1976
1977    def interp_to_stations(
1978        self,
1979        method: typing.Union[str, int],
1980        calls: typing.Sequence[str],
1981        lats: typing.Sequence[float],
1982        lons: typing.Sequence[float],
1983        method_options: typing.Optional[typing.List[int]] = None,
1984        num_threads: int = 1,
1985    ) -> xr.Dataset:
1986        """
1987        Perform spatial interpolation to station points on all variables.
1988
1989        Parameters
1990        ----------
1991        method : str or int
1992            Interpolation method.
1993        calls : sequence of str
1994            Station call signs.
1995        lats : sequence of float
1996            Station latitudes.
1997        lons : sequence of float
1998            Station longitudes.
1999        method_options : list of int, optional
2000            Interpolation options.
2001        num_threads : int, optional
2002            Number of threads.
2003
2004        Returns
2005        -------
2006        xarray.Dataset
2007            Dataset interpolated to stations.
2008        """
2009        da = self._obj.to_array()
2010        da.attrs["GRIB2IO_section3"] = self._obj[list(self._obj.data_vars)[0]].attrs["GRIB2IO_section3"]
2011        da = da.grib2io.interp_to_stations(
2012            method,
2013            calls,
2014            lats,
2015            lons,
2016            method_options=method_options,
2017            num_threads=num_threads,
2018        )
2019        ds = da.to_dataset(dim="variable")
2020
2021        # Update history for provenance
2022        history = ds.attrs.get("history", "")
2023        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
2024        ds.attrs["history"] = f"{now}: Interpolated to {len(calls)} stations via {method}\n{history}"
2025
2026        return ds
2027
2028    def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
2029        """
2030        Write a DataSet to a grib2 file.
2031
2032        Parameters
2033        ----------
2034        filename
2035            Name of the grib2 file to write to.
2036        mode: {"x", "w", "a"}, optional, default="x"
2037            Persistence mode
2038
2039            | mode | Description                       |
2040            | :---:| :---:                             |
2041            | 'x'  | create (fail if exists)           |
2042            | 'w'  | create (overwrite if exists)      |
2043            | 'a'  | append (create if does not exist) |
2044
2045        """
2046        ds = self._obj
2047
2048        for shortName in sorted(ds):
2049            # make a DataArray from the "Data Variables" in the DataSet
2050            da = ds[shortName]
2051
2052            da.grib2io.to_grib2(filename, mode=mode)
2053            mode = "a"
2054
2055    def update_attrs(self, **kwargs):
2056        """
2057        Raises an error because Datasets don't have a .attrs attribute.
2058
2059        Parameters
2060        ----------
2061        attrs
2062            Attributes to update.
2063        """
2064        raise ValueError(f"Datasets do not have a .attrs attribute; use .grib2io.update_attrs({kwargs}) on a DataArray instead.")
2065
2066    def subset(self, *, lats=None, lons=None) -> xr.Dataset:
2067        """
2068        Subset the Dataset to a box defined by latitudes and/or longitudes.
2069
2070        Parameters
2071        ----------
2072        lats
2073            Two item list or tuple of latitudes.  Default is None which will
2074            return a subset unbounded by latitude.  The first term defines the
2075            southern boundary and the second term defines the northern
2076            boundary.
2077        lons
2078            Two item list or tuple of longitudes.  Default is None which will
2079            return a subset unbounded by longitude.  The first term defines the
2080            western boundary and the second term defines the eastern
2081            boundary.  Can follow either: 0 to 360 postive eastward, or 0 to
2082            -180 westward / 0 to 180 eastward conventions.  The longitude
2083            boundaries cannot cross 0.
2084
2085        Returns
2086        -------
2087        subset
2088            Dataset subset to the bounding box created by input 'lats'/'lons'.
2089            All gridpoints with lat/lon matching contraints are included within
2090            subset.
2091        """
2092        ds = self._obj
2093
2094        newds = xr.Dataset()
2095        for shortName in ds:
2096            newds[shortName] = ds[shortName].grib2io.subset(lats=lats, lons=lons).copy()
2097
2098        # Update history for provenance
2099        history = newds.attrs.get("history", "")
2100        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
2101        newds.attrs["history"] = f"{now}: Subsetted to lats={lats}, lons={lons}\n{history}"
2102
2103        return newds
2104
2105    def compute(self, **kwargs):
2106        """
2107        Compute the Dask-backed Dataset with retries for transient errors.
2108
2109        Wraps :func:`grib2io.utils.compute_with_retries`.
2110
2111        Parameters
2112        ----------
2113        **kwargs
2114            Arguments passed to :func:`grib2io.utils.compute_with_retries`,
2115            e.g., `max_attempts` or `base_sleep`.
2116
2117        Returns
2118        -------
2119        xarray.Dataset
2120            The computed Dataset with NumPy-backed data.
2121        """
2122        from .utils import compute_with_retries
2123
2124        return compute_with_retries(self._obj, **kwargs)
2125
2126
2127@xr.register_dataarray_accessor("grib2io")
2128class Grib2ioDataArray:
2129    def __init__(self, xarray_obj):
2130        self._obj = xarray_obj
2131
2132    def griddef(self):
2133        return Grib2GridDef.from_section3(self._obj.attrs["GRIB2IO_section3"])
2134
2135    def interp(self, method, grid_def_out, method_options=None, num_threads=1) -> xr.DataArray:
2136        """
2137        Perform grid spatial interpolation.
2138
2139        Uses the [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip).
2140
2141        Parameters
2142        ----------
2143        method
2144            Interpolate method to use. This can either be an integer or string
2145            using the following mapping:
2146
2147            | Interpolate Scheme | Integer Value |
2148            | :---:              | :---:         |
2149            | 'bilinear'         | 0             |
2150            | 'bicubic'          | 1             |
2151            | 'neighbor'         | 2             |
2152            | 'budget'           | 3             |
2153            | 'spectral'         | 4             |
2154            | 'neighbor-budget'  | 6             |
2155        grid_def_out
2156            Grib2GridDef object of the output grid.
2157        method_options : list of ints, optional
2158            Interpolation options. See the NCEPLIBS-ip documentation for
2159            more information on how these are used.
2160        num_threads : int, optional
2161            Number of OpenMP threads to use for interpolation. The default
2162            value is 1. If grib2io_interp was not built with OpenMP, then
2163            this keyword argument and value will have no impact.
2164
2165        Returns
2166        -------
2167        interp
2168            DataSet interpolated to new grid definition.  The attribute
2169            GRIB2IO_section3 is replaced with the section3 array from the new
2170            grid definition.
2171        """
2172        da = self._obj
2173        # ensure that y, x are rightmost dims; they should be if opening with
2174        # grib2io engine
2175
2176        # gdtn and gdt is not the entirety of the new s3
2177        npoints = grid_def_out.npoints
2178        s3_new = np.array([0, npoints, 0, 0, grid_def_out.gdtn] + list(grid_def_out.gdt))
2179
2180        # make new lat lons
2181        lats, lons = Grib2Message(section3=s3_new, pdtn=0, drtn=0).grid()
2182        latitude = xr.DataArray(lats, dims=["y", "x"])
2183        longitude = xr.DataArray(lons, dims=["y", "x"])
2184
2185        # create new coords
2186        new_coords = dict(da.coords)
2187        del new_coords["latitude"]
2188        del new_coords["longitude"]
2189        new_coords["longitude"] = longitude
2190        new_coords["latitude"] = latitude
2191
2192        # make grid def in from section3 on da.attrs
2193        grid_def_in = self.griddef()
2194
2195        if da.chunks is None:
2196            data = interp_nd(
2197                da.data,
2198                method=method,
2199                grid_def_in=grid_def_in,
2200                grid_def_out=grid_def_out,
2201                method_options=method_options,
2202                num_threads=num_threads,
2203            )
2204        else:
2205            data = da.data.map_blocks(
2206                interp_nd,
2207                method=method,
2208                grid_def_in=grid_def_in,
2209                grid_def_out=grid_def_out,
2210                method_options=method_options,
2211                chunks=da.chunks[:-2] + latitude.shape,
2212                dtype=da.dtype,
2213            )
2214
2215        new_da = xr.DataArray(data, dims=da.dims, coords=new_coords, attrs=da.attrs)
2216
2217        new_da.attrs["GRIB2IO_section3"] = s3_new
2218        new_da.name = da.name
2219
2220        # Update history for provenance
2221        history = new_da.attrs.get("history", "")
2222        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
2223        new_da.attrs["history"] = f"{now}: Interpolated via {method} to {grid_def_out}\n{history}"
2224
2225        return new_da
2226
2227    def interp_to_stations(
2228        self,
2229        method: typing.Union[str, int],
2230        calls: typing.Sequence[str],
2231        lats: typing.Sequence[float],
2232        lons: typing.Sequence[float],
2233        method_options: typing.Optional[typing.List[int]] = None,
2234        num_threads: int = 1,
2235    ) -> xr.DataArray:
2236        """
2237        Perform spatial interpolation to station points.
2238
2239        Parameters
2240        ----------
2241        method : str or int
2242            Interpolate method to use. This can either be an integer or string
2243            using the following mapping:
2244
2245            | Interpolate Scheme | Integer Value |
2246            | :---:              | :---:         |
2247            | 'bilinear'         | 0             |
2248            | 'bicubic'          | 1             |
2249            | 'neighbor'         | 2             |
2250            | 'budget'           | 3             |
2251            | 'spectral'         | 4             |
2252            | 'neighbor-budget'  | 6             |
2253
2254        calls : sequence of str
2255            Station calls used for labeling new station index coordinate
2256        lats : sequence of float
2257            Latitudes of the station points.
2258        lons : sequence of float
2259            Longitudes of the station points.
2260        method_options : list of int, optional
2261            Interpolation options.
2262        num_threads : int, optional
2263            Number of threads.
2264
2265        Returns
2266        -------
2267        xarray.DataArray
2268            DataArray interpolated to lat and lon locations and labeled with
2269            dimension and coordinate 'station'. (..., y, x) -> (..., station)
2270        """
2271        da = self._obj
2272        # TODO ensure that y, x are rightmost dims; they should be if opening
2273        # with grib2io engine
2274
2275        calls = np.asarray(calls)
2276        lats = np.asarray(lats)
2277        lons = np.asarray(lons)
2278        latitude = xr.DataArray(lats, dims=["station"])
2279        longitude = xr.DataArray(lons, dims=["station"])
2280
2281        # create new coords
2282        new_coords = dict(da.coords)
2283        del new_coords["latitude"]
2284        del new_coords["longitude"]
2285        new_coords["longitude"] = longitude
2286        new_coords["latitude"] = latitude
2287        new_coords["station"] = calls
2288
2289        new_dims = da.dims[:-2] + ("station",)
2290
2291        # make grid def in from section3 on da attrs
2292        grid_def_in = self.griddef()
2293
2294        if da.chunks is None:
2295            data = interp_nd_stations(
2296                da.data,
2297                method=method,
2298                grid_def_in=grid_def_in,
2299                lats=lats,
2300                lons=lons,
2301                method_options=method_options,
2302                num_threads=num_threads,
2303            )
2304        else:
2305            data = da.data.map_blocks(
2306                interp_nd_stations,
2307                method=method,
2308                grid_def_in=grid_def_in,
2309                lats=lats,
2310                lons=lons,
2311                method_options=method_options,
2312                drop_axis=-1,
2313                chunks=da.chunks[:-2] + latitude.shape,
2314                dtype=da.dtype,
2315            )
2316
2317        new_da = xr.DataArray(data, dims=new_dims, coords=new_coords, attrs=da.attrs)
2318
2319        new_da.name = da.name
2320
2321        # Update history for provenance
2322        history = new_da.attrs.get("history", "")
2323        now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
2324        new_da.attrs["history"] = f"{now}: Interpolated to {len(calls)} stations via {method}\n{history}"
2325
2326        return new_da
2327
2328    def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
2329        """
2330        Write a DataArray to a grib2 file.
2331
2332        Parameters
2333        ----------
2334        filename
2335            Name of the grib2 file to write to.
2336        mode: {"x", "w", "a"}, optional, default="x"
2337            Persistence mode
2338
2339            +------+-----------------------------------+
2340            | mode | Description                       |
2341            +======+===================================+
2342            | x    | create (fail if exists)           |
2343            +------+-----------------------------------+
2344            | w    | create (overwrite if exists)      |
2345            +------+-----------------------------------+
2346            | a    | append (create if does not exist) |
2347            +------+-----------------------------------+
2348
2349        """
2350        da = self._obj.copy(deep=True)
2351
2352        coords_keys = sorted(da.coords.keys())
2353        coords_keys = [k for k in coords_keys if k in AVAILABLE_NON_GEO_COORDS]
2354
2355        # If there are dimension coordinates, the DataArray is a hypercube of
2356        # grib2 messages.
2357
2358        # Create `indexes` which is a list of lists of dictionaries for all
2359        # dimension coordinates. Each dictionary key is the dimension
2360        # coordinate name and the value is a list of the dimension coordinate
2361        # values.  This allows for easy iteration over all possible grib2
2362        # messages in the DataArray by using itertools.product.
2363        #
2364        # For example:
2365        # indexes = [
2366        #     [
2367        #         {"leadTime": 9},
2368        #         {"leadTime": 12},
2369        #     ],
2370        #     [
2371        #         {"valueOfFirstFixedSurface": 900},
2372        #         {"valueOfFirstFixedSurface": 925},
2373        #         {"valueOfFirstFixedSurface": 950},
2374        #     ],
2375        # ]
2376
2377        # assign loc indexes to dimensions without indexes for uniform selection by name
2378        loc_indexes = list()
2379        for dim in da.dims:
2380            if dim not in da.indexes:
2381                da = da.assign_coords({dim: range(da[dim].size)})
2382                loc_indexes.append(dim)
2383
2384        indexes = []
2385        for index in [i for i in AVAILABLE_NON_GEO_DIMS if i in da.dims]:
2386            values = da.coords[index].values
2387            if len(values) != len(set(values)):
2388                raise ValueError(
2389                    f"Dimension coordinate '{index}' has duplicate values, but to_grib2 requires unique values to find each GRIB2 message in the DataArray."
2390                )
2391            listeach = [{index: value} for value in sorted(values)]
2392            indexes.append(listeach)
2393
2394        # If `dim_coords` is [], then the DataArray is a single grib2 message and
2395        # itertools.product(*dim_coords) will run once with `selectors = ()`.
2396        for selectors in itertools.product(*indexes):
2397            # Need to find the correct data in the DataArray based on the
2398            # dimension coordinates.
2399            filters = {k: v for d in selectors for k, v in d.items()}
2400
2401            # If `filters` is {}, then the DataArray is a single grib2 message
2402            # and da.sel(indexers={}) returns the DataArray.
2403            selected = da.sel(indexers=filters)
2404
2405            newmsg = Grib2Message(
2406                selected.attrs["GRIB2IO_section0"],
2407                selected.attrs["GRIB2IO_section1"],
2408                selected.attrs["GRIB2IO_section2"],
2409                selected.attrs["GRIB2IO_section3"],
2410                selected.attrs["GRIB2IO_section4"],
2411                selected.attrs["GRIB2IO_section5"],
2412            )
2413            newmsg.data = np.array(selected.data)
2414
2415            # For dimension coordinates, set the grib2 message metadata to the
2416            # dimension coordinate value.
2417            for index, value in filters.items():
2418                if index not in loc_indexes:
2419                    setattr(newmsg, index, value)
2420
2421            # For non-dimension coordinates, set the grib2 message metadata to
2422            # the DataArray coordinate value.
2423            for index in [i for i in coords_keys if i not in da.dims]:
2424                setattr(newmsg, index, selected.coords[index].values)
2425
2426            # Set section 5 attributes to the da.encoding dictionary.
2427            for key, value in selected.encoding.items():
2428                if key in ["dtype", "chunks", "original_shape"]:
2429                    continue
2430                setattr(newmsg, key, value)
2431
2432            # write the message to file
2433            with grib2io.open(filename, mode=mode) as f:
2434                f.write(newmsg)
2435            mode = "a"
2436
2437    def update_attrs(self, **kwargs):
2438        """
2439        Update many of the attributes of the DataArray.
2440
2441        Parameters
2442        ----------
2443        **kwargs
2444            Attributes to update.  This can include many of the GRIB2IO message
2445            attributes that you can find when you print a GRIB2IO message. For
2446            conflicting updates, the last keyword will be used.
2447
2448            +-----------------------+------------------------------------------+
2449            | kwargs                | Description                              |
2450            +=======================+==========================================+
2451            | shortName="VTMP"      | Set shortName to "VTMP", along with      |
2452            |                       | appropriate discipline,                  |
2453            |                       | parameterCategory, parameterNumber,      |
2454            |                       | fullName and units.                      |
2455            +-----------------------+------------------------------------------+
2456            | discipline=0,         | Set shortName, discipline,               |
2457            | parameterCategory=0,  | parameterCategory, parameterNumber,      |
2458            | parameterNumber=1     | fullName and units appropriate for       |
2459            |                       | "Virtual Temperature".                   |
2460            +-----------------------+------------------------------------------+
2461            | discipline=0,         | Conflicting keywords but                 |
2462            | parameterCategory=0,  | 'shortName="TMP"' wins.  Set shortName,  |
2463            | parameterNumber=1,    | discipline, parameterCategory,           |
2464            | shortName="TMP"       | parameterNumber, fullName and units      |
2465            |                       | appropriate for "Temperature".           |
2466            +-----------------------+------------------------------------------+
2467
2468        Returns
2469        -------
2470        DataArray
2471            DataArray with updated attributes.
2472        """
2473        da = self._obj.copy(deep=True)
2474
2475        newmsg = Grib2Message(
2476            da.attrs["GRIB2IO_section0"],
2477            da.attrs["GRIB2IO_section1"],
2478            da.attrs["GRIB2IO_section2"],
2479            da.attrs["GRIB2IO_section3"],
2480            da.attrs["GRIB2IO_section4"],
2481            da.attrs["GRIB2IO_section5"],
2482        )
2483
2484        coords_keys = [k for k in da.coords.keys() if k in AVAILABLE_NON_GEO_COORDS]
2485
2486        for grib2_name, value in kwargs.items():
2487            if grib2_name == "gridDefinitionTemplateNumber":
2488                raise ValueError(
2489                    "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."
2490                )
2491            if grib2_name == "productDefinitionTemplateNumber":
2492                raise ValueError("The productDefinitionTemplateNumber attribute cannot be updated.")
2493            if grib2_name == "dataRepresentationTemplateNumber":
2494                raise ValueError("The dataRepresentationTemplateNumber attribute cannot be updated.")
2495            if grib2_name in coords_keys:
2496                warnings.warn(f"Skipping attribute '{grib2_name}' because it is a coordinate. Use da.assign_coords() to change coordinate values.")
2497                continue
2498            if hasattr(newmsg, grib2_name):
2499                setattr(newmsg, grib2_name, value)
2500            else:
2501                warnings.warn(f"Skipping attribute '{grib2_name}' because it is not a valid GRIB2 attribute for this message and cannot be updated.")
2502                continue
2503
2504        da.attrs["GRIB2IO_section0"] = newmsg.section0
2505        da.attrs["GRIB2IO_section1"] = newmsg.section1
2506        da.attrs["GRIB2IO_section2"] = newmsg.section2 or []
2507        da.attrs["GRIB2IO_section3"] = newmsg.section3
2508        da.attrs["GRIB2IO_section4"] = newmsg.section4
2509        da.attrs["GRIB2IO_section5"] = newmsg.section5
2510        da.attrs["fullName"] = newmsg.fullName
2511        da.attrs["shortName"] = newmsg.shortName
2512        da.attrs["units"] = newmsg.units
2513
2514        return da
2515
2516    def update_section3(self) -> xr.DataArray:
2517        """
2518        Update section3 attributes based on the latitude and longitude corners.
2519
2520        This makes the GRIB2IO_section3 attribute consistent with the grid's
2521        new corners after a change in the spatial extent.
2522        """
2523        da = self._obj
2524        if "GRIB2IO_section3" not in da.attrs:
2525            raise ValueError(
2526                "DataArray has no attr 'GRIB2IO_section3'.  This function only works with Datasets/DataArrrays opened with the 'grib2io' backend."
2527            )
2528        if "latitude" not in da.coords:
2529            raise ValueError("DataArray has no coord 'latitude'")
2530        if "longitude" not in da.coords:
2531            raise ValueError("DataArray has no coord 'longitude'")
2532
2533        grid = Grid(da.attrs["GRIB2IO_section3"])
2534
2535        if grid.gdtn not in [0, 1, 10, 20, 30, 31, 40, 110]:
2536            raise ValueError(
2537                textwrap.dedent("""\
2538                    update_section3 only works for:
2539
2540                    Latitude/Longitude, Equidistant Cylindrical, or Plate Carree (gdtn=0)
2541                    Rotated Latitude/Longitude (gdtn=1)
2542                    Mercator (gdtn=10)
2543                    Polar Stereographic (gdtn=20)
2544                    Lambert Conformal (gdtn=30)
2545                    Albers Equal-Area (gdtn=31)
2546                    Gaussian Latitude/Longitude (gdtn=40)
2547                    Equatorial Azimuthal Equidistant Projection (gdtn=110)
2548                    """)
2549            )
2550
2551        grid.latitudeFirstGridpoint = da.latitude.isel(y=0, x=0)
2552        grid.longitudeFirstGridpoint = da.longitude.isel(y=0, x=0)
2553        grid.nx = len(da.x)
2554        grid.ny = len(da.y)
2555
2556        # last gridpoint does not affect section3 for some gdt but set anyway
2557        grid.latitudeLastGridpoint = da.latitude.isel(y=-1, x=-1)
2558        grid.longitudeLastGridpoint = da.longitude.isel(y=-1, x=-1)
2559
2560        da.attrs["GRIB2IO_section3"] = grid.section3
2561
2562        return da
2563
2564    def subset(self, *, lats=None, lons=None) -> xr.DataArray:
2565        """
2566        Subset the DataArray to a box defined by latitudes and/or longitudes.
2567
2568        Parameters
2569        ----------
2570        lats
2571            Two item list or tuple of latitudes.  Default is None which will
2572            return a subset unbounded by latitude.  The first term defines the
2573            southern boundary and the second term defines the northern
2574            boundary.
2575        lons
2576            Two item list or tuple of longitudes.  Default is None which will
2577            return a subset unbounded by longitude.  The first term defines the
2578            western boundary and the second term defines the eastern
2579            boundary.  Can follow either: 0 to 360 postive eastward, or 0 to
2580            -180 westward / 0 to 180 eastward conventions.  The longitude
2581            boundaries cannot cross 0.
2582
2583        Returns
2584        -------
2585        subset
2586            DataArray subset to the bounding box created by input 'lats'/'lons'.
2587            All gridpoints with lat/lon matching contraints are included within
2588            subset.
2589        """
2590
2591        def slice_from_contiguous_mask(mask):
2592            indices = np.where(mask)[0]
2593
2594            if indices.size > 0:
2595                # slice(start, stop) - stop is exclusive, so we add 1
2596                my_slice = slice(indices[0], indices[-1] + 1)
2597            else:
2598                my_slice = slice(0, 0)
2599            return my_slice
2600
2601        da = self._obj.copy()
2602
2603        if lats is None:
2604            lats = (np.min(da.latitude), np.max(da.latitude))
2605        else:
2606            lats = (min(lats), max(lats))
2607
2608        if lons is None:
2609            lons = (np.min(da.longitude), np.max(da.longitude))
2610        else:
2611            lons = (min(lons), max(lons))
2612
2613        # Internally work in common lon data representation (0->360 positive eastward from 0)
2614        lons = np.mod(np.array(lons) + 360, 360)
2615        lon_da = np.mod(da.longitude + 360, 360)
2616
2617        snap_first_point = snap_to_nearest_cell_center(da.latitude, lon_da, lats[0], lons[0])
2618        snap_last_point = snap_to_nearest_cell_center(da.latitude, lon_da, lats[1], lons[1])
2619        lats = (snap_first_point[0], snap_last_point[0])
2620        lons = (snap_first_point[1], snap_last_point[1])
2621
2622        x = ((lon_da >= lons[0]) & (lon_da <= lons[1])).any("y")
2623        if x.chunks:
2624            x = x.compute()
2625
2626        y = ((da.latitude >= lats[0]) & (da.latitude <= lats[1])).any("x")
2627        if y.chunks:
2628            y = y.compute()
2629
2630        y_slice = slice_from_contiguous_mask(y)
2631        x_slice = slice_from_contiguous_mask(x)
2632
2633        da = da.isel(y=y_slice, x=x_slice)
2634        if da.size < 1:
2635            raise ValueError("None of grid data is within given lat/lon bounds.")
2636
2637        da = da.grib2io.update_section3()
2638
2639        return da
2640
2641    def compute(self, **kwargs):
2642        """
2643        Compute the Dask-backed DataArray with retries for transient errors.
2644
2645        Wraps :func:`grib2io.utils.compute_with_retries`.
2646
2647        Parameters
2648        ----------
2649        **kwargs
2650            Arguments passed to :func:`grib2io.utils.compute_with_retries`,
2651            e.g., `max_attempts` or `base_sleep`.
2652
2653        Returns
2654        -------
2655        xarray.DataArray
2656            The computed DataArray with NumPy-backed data.
2657        """
2658        from .utils import compute_with_retries
2659
2660        return compute_with_retries(self._obj, **kwargs)
2661
2662
2663def open_mfdataset(
2664    filenames: typing.Union[str, typing.Sequence[str]],
2665    *,
2666    drop_variables: typing.Optional[typing.List[str]] = None,
2667    save_index: bool = True,
2668    filters: typing.Mapping[str, typing.Any] = dict(),
2669    data_model: typing.Optional[str] = None,
2670    parallel: bool = False,
2671    preprocess: typing.Optional[typing.Callable] = None,
2672    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
2673    **kwargs,
2674) -> xr.Dataset:
2675    """
2676    Open multiple GRIB2 files as a single xarray Dataset.
2677
2678    This function is optimized for GRIB2 files by combining their indices
2679    and creating a single Dataset, which is often much faster than
2680    using ``xarray.open_mfdataset``. It supports parallel index reading
2681    and dataset opening when ``parallel=True`` and ``dask`` is installed.
2682
2683    Parameters
2684    ----------
2685    filenames : str or sequence of str
2686        GRIB2 files to be opened. Can be a glob pattern.
2687    drop_variables : list of str, optional
2688        List of variables to exclude from the dataset.
2689    save_index : bool, optional
2690        Whether to save the GRIB2 index to a file (default is True).
2691    filters : dict, optional
2692        Filter GRIB2 messages to a single hypercube. Dictionary keys can be
2693        any GRIB2 metadata attribute name.
2694    data_model : str, optional
2695        Parse GRIB metadata following a defined data model convention
2696        (e.g., "nws-viz").
2697    parallel : bool, optional
2698        If True, use ``dask`` to read indices and open datasets in parallel.
2699        Requires the ``dask`` package.
2700    preprocess : callable, optional
2701        A function to apply to each file's dataset before combining.
2702    chunks : int, dict or 'auto', optional
2703        If chunks is provided, it is used to load the dataset into a
2704        dask-backed dataset.
2705    **kwargs : optional
2706        Additional arguments passed to the combination logic.
2707        If ``combine='nested'``, passed to ``xarray.combine_nested``.
2708        If ``combine='by_coords'``, passed to ``xarray.combine_by_coords``.
2709        If ``combine='merge'``, passed to ``xarray.merge``.
2710        If no ``combine`` argument is provided, the function attempts
2711        ``xarray.combine_by_coords`` followed by ``xarray.merge``.
2712
2713    Returns
2714    -------
2715    xarray.Dataset
2716        Xarray dataset of grib2 messages.
2717
2718    Notes
2719    -----
2720    - This function uses a "fast path" when ``preprocess=None`` and no
2721      combination ``**kwargs`` are provided, which concatenates all indices
2722      into a single global index before building the Dataset.
2723    - All files must share the same horizontal grid (ny, nx).
2724    """
2725    if isinstance(filenames, str):
2726        import glob
2727
2728        filenames = sorted(glob.glob(filenames))
2729
2730    storage_options = kwargs.pop("storage_options", None)
2731
2732    def _get_index(fname_and_index: typing.Tuple[str, int]) -> pd.DataFrame:
2733        """
2734        Internal utility to read GRIB2 index from a file.
2735
2736        Parameters
2737        ----------
2738        fname_and_index : tuple of (str, int)
2739            Tuple containing the filename and its position in the file list.
2740
2741        Returns
2742        -------
2743        pandas.DataFrame
2744            The GRIB2 index for the specified file.
2745        """
2746        fname, i = fname_and_index
2747        with grib2io.open(
2748            fname,
2749            save_index=save_index,
2750            _xarray_backend=True,
2751            **(storage_options or {}),
2752        ) as f:
2753            idx = pd.DataFrame(f._index)
2754            idx = idx.assign(msg=list(f))
2755            idx["file_index"] = i
2756            return idx
2757
2758    if parallel:
2759        try:
2760            import dask
2761            from dask.bag import from_sequence
2762
2763            indices = from_sequence(zip(filenames, range(len(filenames)))).map(_get_index).compute()
2764        except ImportError:
2765            warnings.warn("dask not installed, falling back to sequential index reading.")
2766            parallel = False
2767            indices = [_get_index((fname, i)) for i, fname in enumerate(filenames)]
2768    else:
2769        indices = [_get_index((fname, i)) for i, fname in enumerate(filenames)]
2770
2771    if not indices:
2772        return xr.Dataset()
2773
2774    # Validate grid consistency across files using only the first message of each file
2775    grid_cols = ["ny", "nx"]
2776    first_msgs = pd.concat([idx.head(1) for idx in indices], ignore_index=True)
2777    unique_grids = first_msgs[grid_cols].drop_duplicates()
2778    if len(unique_grids) > 1:
2779        grid_list = unique_grids.to_dict("records")
2780        raise ValueError(f"Multiple grids detected in open_mfdataset. All files must have the same grid. Found grids: {grid_list}")
2781
2782    # Determine if we can use the fast path (single index concatenation)
2783    # The fast path is only available if no preprocess is provided and no combination kwargs are used
2784    # that would require individual datasets (like concat_dim for nested combination)
2785    use_fast_path = preprocess is None and not kwargs
2786
2787    if not use_fast_path:
2788        if parallel:
2789            import dask
2790
2791            @dask.delayed
2792            def _open_delayed(idx, fname):
2793                return _open_dataset_from_index(
2794                    idx,
2795                    fname,
2796                    filters,
2797                    data_model,
2798                    drop_variables=drop_variables,
2799                    chunks=chunks,
2800                )
2801
2802            datasets = dask.compute(*[_open_delayed(idx, fname) for idx, fname in zip(indices, filenames)])
2803        else:
2804            datasets = [
2805                _open_dataset_from_index(
2806                    idx,
2807                    fname,
2808                    filters,
2809                    data_model,
2810                    drop_variables=drop_variables,
2811                    chunks=chunks,
2812                )
2813                for idx, fname in zip(indices, filenames)
2814            ]
2815
2816        if preprocess is not None:
2817            datasets = [preprocess(ds) for ds in datasets]
2818
2819        combine_opt = kwargs.pop("combine", None)
2820        if combine_opt == "nested":
2821            ds = xr.combine_nested(datasets, **kwargs)
2822        elif combine_opt == "by_coords":
2823            ds = xr.combine_by_coords(datasets, **kwargs)
2824        elif combine_opt == "merge":
2825            ds = xr.merge(datasets, **kwargs)
2826        else:
2827            # Default behavior: try by_coords, then merge
2828            try:
2829                ds = xr.combine_by_coords(datasets, **kwargs)
2830            except Exception:
2831                ds = xr.merge(datasets, **kwargs)
2832    else:
2833        file_index = pd.concat(indices, ignore_index=True)
2834        ds = _open_dataset_from_index(
2835            file_index,
2836            list(filenames),
2837            filters,
2838            data_model,
2839            drop_variables=drop_variables,
2840            chunks=chunks,
2841        )
2842
2843    # Update history for provenance
2844    history = ds.attrs.get("history", "")
2845    now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
2846    ds.attrs["history"] = f"{now}: Initialized via grib2io.open_mfdataset from {len(filenames)} files\n{history}"
2847
2848    return ds
2849
2850
2851def _open_dataset_from_index(
2852    file_index: pd.DataFrame,
2853    filenames: typing.Union[str, typing.List[str]],
2854    filters: typing.Mapping[str, typing.Any] = dict(),
2855    data_model: typing.Optional[str] = None,
2856    drop_variables: typing.Optional[typing.List[str]] = None,
2857    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
2858) -> xr.Dataset:
2859    """
2860    Create an xarray Dataset from a GRIB2 index DataFrame.
2861
2862    This is an internal utility used by ``open_dataset`` and ``open_mfdataset``
2863    to build a Dataset structure from a pre-computed index of GRIB2 messages.
2864
2865    Parameters
2866    ----------
2867    file_index : pandas.DataFrame
2868        GRIB2 index DataFrame, expected to contain GRIB2 metadata and
2869        message pointers.
2870    filenames : str or list of str
2871        Path(s) to the GRIB2 file(s) referenced by the index.
2872    filters : dict, optional
2873        Filter GRIB2 messages to a single hypercube.
2874    data_model : str, optional
2875        Target data model for metadata normalization (e.g., "nws-viz").
2876    drop_variables : list of str, optional
2877        List of shortnames to exclude from the resulting Dataset.
2878    chunks : int, dict or 'auto', optional
2879        If chunks is provided, it is used to load the dataset into a
2880        dask-backed dataset.
2881
2882    Returns
2883    -------
2884    xarray.Dataset
2885        Dataset representing the GRIB2 messages.
2886    """
2887    # parse grib2io _index to dataframe and acquire non-geo possible dims
2888    # (scalar coord when not dim due to squeeze) parse_grib_index applies
2889    # filters to index and expands metadata based on product definition
2890    # template number
2891    file_index, dim_coords, attrs, coord_attrs = parse_grib_index(file_index, filters)
2892
2893    if drop_variables:
2894        file_index = file_index[~file_index["shortName"].isin(drop_variables)]
2895
2896    # Divide up records by variable
2897    frames, cubes, extra_geo = make_variables(file_index, filenames, dim_coords)  # have this return var_attrs
2898
2899    # return empty dataset if no data
2900    if frames is None:
2901        return xr.Dataset()
2902
2903    # create dataframe and add datarrays without any coords
2904    ds = xr.Dataset()
2905    for var_df, var_cube in zip(frames, cubes):
2906        da = build_da_without_coords(var_df, var_cube, filenames, attrs)
2907
2908        # Assign variable-specific coords from its cube
2909        coords = coords_from_cube(var_cube)
2910        da = da.assign_coords(coords)
2911
2912        # Assign extra index associated coords for this variable
2913        for dim_name, coord_names in dim_coords.items():
2914            retain_index_coord = False
2915            for name in coord_names:
2916                if name == dim_name:
2917                    retain_index_coord = True
2918                else:
2919                    if dim_name not in da.dims:
2920                        # for assigning scalar coords
2921                        coord_data = var_df[name].unique().item()
2922                        da = da.assign_coords({name: coord_data})
2923                    else:
2924                        # "ValueError: can only convert an array of size 1 to a Python scalar" indicates the coord is not compatible with the index
2925                        coord_data = [
2926                            var_df[var_df.index.get_level_values(f"{dim_name}_ix") == val][name].unique().item() for val in range(da[dim_name].size)
2927                        ]
2928                        coord = pd.Index(coord_data, name=dim_name)
2929                        da = da.assign_coords({name: (dim_name, coord)})
2930            if not retain_index_coord and dim_name in da.coords:
2931                da = da.drop_vars(dim_name)
2932
2933        ds[da.name] = da
2934
2935    # add coords and dataset meta
2936    # Pass first cube for common geo coords assignment
2937    ds = assign_xr_meta(ds, frames, cubes[0], dim_coords, extra_geo, coord_attrs)
2938
2939    if data_model is not None:
2940        ds = parse_data_model(ds, data_model)
2941
2942    # assign attributes
2943    ds.attrs["engine"] = "grib2io"
2944
2945    if chunks is not None:
2946        ds = ds.chunk(chunks)
2947
2948    return ds
2949
2950
2951def build_datatree_from_grib(
2952    filename: str,
2953    file_index: pd.DataFrame,
2954    filters: typing.Optional[typing.Mapping[str, typing.Any]] = None,
2955    stack_vertical: bool = False,
2956    drop_variables: typing.Optional[typing.List[str]] = None,
2957    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
2958) -> typing.Any:
2959    """
2960    Build a DataTree from GRIB2 messages.
2961
2962    This internal function organizes GRIB2 messages into a hierarchical
2963    tree structure based on level types, PDTNs, and other metadata.
2964
2965    Parameters
2966    ----------
2967    filename : str
2968        Path to the source GRIB2 file.
2969    file_index : pandas.DataFrame
2970        Index of GRIB2 messages.
2971    filters : dict, optional
2972        Filter criteria for selecting messages.
2973    stack_vertical : bool, optional
2974        If True, vertical levels will be stacked in a single dataset
2975        within each node, rather than creating separate nodes per level value.
2976    drop_variables : list of str, optional
2977        List of variable shortnames to exclude.
2978    chunks : int, dict or 'auto', optional
2979        If chunks is provided, it is used to load the dataset into a
2980        dask-backed dataset.
2981
2982    Returns
2983    -------
2984    xarray.DataTree
2985        A hierarchical tree representation of the GRIB2 data.
2986    """
2987    if filters is None:
2988        filters = {}
2989
2990    # Apply any filters from user
2991    for k, v in filters.items():
2992        if k not in file_index.columns:
2993            file_index = file_index.copy()
2994            file_index[k] = file_index.msg.apply(lambda msg: getattr(msg, k, None))
2995        file_index = filter_index(file_index, k, v)
2996
2997    # Make a copy to avoid the SettingWithCopyWarning
2998    file_index = file_index.copy()
2999
3000    # Extract metadata needed for tree organization
3001    # Use a safer approach to handle missing attributes
3002    def safe_getattr(obj, name):
3003        try:
3004            attr = getattr(obj, name)
3005            # Need to test if the attribute is Grib2Metadata. If so,
3006            # then get the value attribute.
3007            if isinstance(attr, grib2io.templates.Grib2Metadata):
3008                attr = attr.value
3009            return attr
3010        except (AttributeError, KeyError):
3011            return None
3012
3013    for attr in _TREE_HIERARCHY_LEVELS:
3014        if (attr not in file_index.columns) and (attr != "valueOfFirstFixedSurface"):
3015            file_index[attr] = file_index.msg.apply(lambda msg: safe_getattr(msg, attr))
3016
3017    # Also extract shortName for variable naming
3018    if "shortName" not in file_index.columns:
3019        file_index = file_index.assign(shortName=file_index.msg.apply(lambda msg: getattr(msg, "shortName", None)))
3020
3021    if drop_variables:
3022        file_index = file_index[~file_index["shortName"].isin(drop_variables)]
3023
3024    file_index = file_index.assign(nx=file_index.msg.apply(lambda msg: getattr(msg, "nx", None)))
3025    file_index = file_index.assign(ny=file_index.msg.apply(lambda msg: getattr(msg, "ny", None)))
3026
3027    # Create root DataTree
3028    root = xr.DataTree()
3029
3030    # Adjust hierarchy levels if we're stacking vertical levels
3031    hierarchy_levels = list(_TREE_HIERARCHY_LEVELS)  # This makes a copy
3032    if stack_vertical and "valueOfFirstFixedSurface" in hierarchy_levels:
3033        hierarchy_levels.remove("valueOfFirstFixedSurface")
3034
3035    # First group by level type
3036    level_groups = {}
3037
3038    # Create a dictionary to group data by level type
3039    for level_type in file_index["typeOfFirstFixedSurface"].unique():
3040        if pd.notna(level_type):  # Skip None/NaN values
3041            level_info = _LEVEL_NAME_MAPPING.get(level_type, f"level_{level_type}")
3042            level_name = level_info[0]
3043            # Get all rows for this level type
3044            level_data = file_index[file_index["typeOfFirstFixedSurface"] == level_type]
3045            level_groups[level_type] = {"name": level_name, "data": level_data}
3046
3047    # Process each level group
3048    for level_type, group_info in level_groups.items():
3049        level_name = group_info["name"]
3050        level_df = group_info["data"]
3051
3052        # Create a branch for this level type
3053        level_tree = xr.DataTree()
3054
3055        # Process this branch based on PDTN, perturbation number, etc.
3056        process_level_branch(level_tree, level_df, filename, chunks=chunks)
3057
3058        # Add this branch to the main tree
3059        root[level_name] = level_tree
3060
3061    return root
3062
3063
3064def process_level_branch(
3065    level_tree: typing.Any,
3066    df: pd.DataFrame,
3067    filename: str,
3068    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
3069):
3070    """
3071    Process a level type branch of the data tree.
3072
3073    Organizes the tree by PDTN and other attributes.
3074
3075    Parameters
3076    ----------
3077    level_tree : xarray.DataTree
3078        The DataTree node for this level type.
3079    df : pandas.DataFrame
3080        DataFrame of messages for this level type.
3081    filename : str
3082        Path to the GRIB2 file.
3083    chunks : int, dict or 'auto', optional
3084        If chunks is provided, it is used to load the dataset into a
3085        dask-backed dataset.
3086    """
3087    # Group by PDTN
3088    pdtn_groups = {}
3089
3090    # Group data by PDTN first
3091    for pdtn_value in df["productDefinitionTemplateNumber"].unique():
3092        if pd.notna(pdtn_value):
3093            pdtn_df = df[df["productDefinitionTemplateNumber"] == pdtn_value]
3094            pdtn_groups[pdtn_value] = pdtn_df
3095
3096    # If there's only one PDTN value, skip creating PDTN branch level
3097    if len(pdtn_groups) == 1:
3098        pdtn, pdtn_df = next(iter(pdtn_groups.items()))
3099
3100        pdtn_name = f"pdtn_{int(pdtn)}"
3101
3102        # Check if we need to further subdivide by perturbation number
3103        has_perturbations = "perturbationNumber" in pdtn_df.columns and len(pdtn_df["perturbationNumber"].dropna().unique()) > 1
3104
3105        # Check if we need to further subdivide by probabilities unique for each variable.
3106        has_probabilities = "typeOfProbability" in pdtn_df.columns and len(pdtn_df["typeOfProbability"].dropna().unique()) > 1
3107
3108        if has_perturbations:
3109            # Process perturbations directly on the level tree
3110            process_perturbation_groups(level_tree, pdtn_df, filename, chunks=chunks)
3111        elif has_probabilities:
3112            # Process probability groups
3113            process_probability_groups(level_tree, pdtn_df, filename, chunks=chunks)
3114        else:
3115            # Try to create dataset directly on level
3116            try:
3117                dss = create_datasets_from_df(pdtn_df, filename, chunks=chunks)
3118                if dss is not None:
3119                    dt = xr.DataTree()
3120                    if len(dss) == 1:
3121                        dt.ds = dss[0]
3122                    else:
3123                        for ds in dss:
3124                            varname = list(ds.data_vars)[0]
3125                            dt[f"var_{varname}"] = ds
3126                    level_tree[pdtn_name] = dt
3127                else:
3128                    # Try to separate by variable name as a fallback
3129                    try_process_by_variables(level_tree, pdtn_df, filename, chunks=chunks)
3130            except Exception as e:
3131                print(f"Error creating dataset for level with pdtn {int(pdtn)}: {e}")
3132
3133                # Try to separate by variable name as a fallback
3134                try_process_by_variables(level_tree, pdtn_df, filename, chunks=chunks)
3135    else:
3136        # Multiple PDTN values, process each group with PDTN branch nodes
3137        for pdtn, pdtn_df in pdtn_groups.items():
3138            # Use a simple node name that's easy to use in code
3139            pdtn_name = f"pdtn_{int(pdtn)}"
3140
3141            # Check if we need to further subdivide by perturbation number
3142            has_perturbations = "perturbationNumber" in pdtn_df.columns and len(pdtn_df["perturbationNumber"].dropna().unique()) > 1
3143
3144            # Check if we need to further subdivide by probabilities unique for each variable.
3145            has_probabilities = "typeOfProbability" in pdtn_df.columns and len(pdtn_df["typeOfProbability"].dropna().unique()) > 1
3146
3147            if has_perturbations:
3148                # Create a branch for this PDTN
3149                pdtn_tree = xr.DataTree()
3150
3151                # Process perturbation groups
3152                process_perturbation_groups(pdtn_tree, pdtn_df, filename, chunks=chunks)
3153
3154                # Only add the PDTN branch if it has children
3155                if len(pdtn_tree.children) > 0 or pdtn_tree.ds is not None:
3156                    level_tree[pdtn_name] = pdtn_tree
3157            elif has_probabilities:
3158                # Create a branch for this PDTN
3159                pdtn_tree = xr.DataTree()
3160
3161                # Process probability groups
3162                process_probability_groups(pdtn_tree, pdtn_df, filename, chunks=chunks)
3163
3164                # Only add the PDTN branch if it has children
3165                if len(pdtn_tree.children) > 0 or pdtn_tree.ds is not None:
3166                    level_tree[pdtn_name] = pdtn_tree
3167            else:
3168                # Create a subtree for this PDTN
3169                pdtn_tree = xr.DataTree()
3170
3171                # Try to create dataset directly on level
3172                try:
3173                    dss = create_datasets_from_df(pdtn_df, filename, chunks=chunks)
3174                    if dss is not None:
3175                        if len(dss) == 1:
3176                            pdtn_tree.ds = dss[0]
3177                        else:
3178                            for ds in dss:
3179                                varname = list(ds.data_vars)[0]
3180                                pdtn_tree[f"var_{varname}"] = ds
3181                        level_tree[pdtn_name] = pdtn_tree
3182                    else:
3183                        # Try to separate by variable name as a fallback
3184                        try_process_by_variables(pdtn_tree, pdtn_df, filename, chunks=chunks)
3185                        level_tree[pdtn_name] = pdtn_tree
3186                except Exception as e:
3187                    print(f"Error creating dataset for level with pdtn {int(pdtn)}: {e}")
3188
3189                    # Try to separate by variable name as a fallback
3190                    try_process_by_variables(pdtn_tree, pdtn_df, filename, chunks=chunks)
3191                    level_tree[pdtn_name] = pdtn_tree
3192
3193
3194def process_probability_groups(
3195    target_tree: typing.Any,
3196    pdtn_df: pd.DataFrame,
3197    filename: str,
3198    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
3199) -> bool:
3200    """
3201    Process probability groups and add them to the target tree.
3202
3203    Parameters
3204    ----------
3205    target_tree : xarray.DataTree
3206        The tree node to add probability groups to.
3207    pdtn_df : pandas.DataFrame
3208        DataFrame of messages for a specific PDTN.
3209    filename : str
3210        Path to the GRIB2 file.
3211    chunks : int, dict or 'auto', optional
3212        If chunks is provided, it is used to load the dataset into a
3213        dask-backed dataset.
3214
3215    Returns
3216    -------
3217    bool
3218        True if successful.
3219    """
3220    success = False
3221    # Group by type of probability
3222    prob_groups = {}
3223    for prob_value in pdtn_df["typeOfProbability"].unique():
3224        if pd.notna(prob_value):
3225            prob_df = pdtn_df[pdtn_df["typeOfProbability"] == prob_value]
3226            prob_groups[prob_value] = prob_df
3227
3228    # Process each probability group
3229    for prob_num, prob_df in prob_groups.items():
3230        prob_name = f"prob_{int(prob_num)}"
3231
3232        # Try to create dataset for this probability group
3233        try:
3234            dss = create_datasets_from_df(prob_df, filename, chunks=chunks)
3235            dt = xr.DataTree()
3236            if len(dss) == 1:
3237                dt.ds = dss[0]
3238                target_tree[prob_name] = dt
3239            elif len(dss) > 1:
3240                for ds in dss:
3241                    dt[f"var_{ds.data_vars[0]}"] = ds
3242            target_tree[prob_name] = dt
3243        except Exception as e:
3244            # Log error but continue processing other groups
3245            print(f"Error creating dataset for type of probability {prob_name}: {e}")
3246
3247    return success
3248
3249
3250def process_perturbation_groups(
3251    target_tree: typing.Any,
3252    pdtn_df: pd.DataFrame,
3253    filename: str,
3254    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
3255) -> bool:
3256    """
3257    Process perturbation groups and add them to the target tree.
3258
3259    Parameters
3260    ----------
3261    target_tree : xarray.DataTree
3262        The tree node to add perturbation groups to.
3263    pdtn_df : pandas.DataFrame
3264        DataFrame of messages for a specific PDTN.
3265    filename : str
3266        Path to the GRIB2 file.
3267    chunks : int, dict or 'auto', optional
3268        If chunks is provided, it is used to load the dataset into a
3269        dask-backed dataset.
3270
3271    Returns
3272    -------
3273    bool
3274        True if at least one perturbation was successfully processed.
3275    """
3276    success = False
3277    # Group by perturbation number
3278    pert_groups = {}
3279    for pert_value in pdtn_df["perturbationNumber"].unique():
3280        if pd.notna(pert_value):
3281            pert_df = pdtn_df[pdtn_df["perturbationNumber"] == pert_value]
3282            pert_groups[pert_value] = pert_df
3283
3284    # Process each perturbation group
3285    for pert_num, pert_df in pert_groups.items():
3286        pert_name = f"pert_{int(pert_num)}"
3287
3288        ## Try to create dataset for this perturbation group
3289        # try:
3290        #    dss = create_datasets_from_df(pert_df, filename)
3291        #    if dss is not None:
3292        #        if len(dss) == 1:
3293        #            target_tree.ds = dss[0]
3294        #        else:
3295        #            dss_dict = {f"ds_{i}": ds for i, ds in enumerate(dss)}
3296        #            atree = xr.DataTree(dss_dict)
3297        #            target_tree[prob_name] = atree
3298        #        success = True
3299        # except Exception as e:
3300        #    # Log error but continue processing other groups
3301        #    print(f"Error creating dataset for perturbation {pert_name}: {e}")
3302
3303        # Try to create dataset for this perturbation group
3304        try:
3305            dss = create_datasets_from_df(pert_df, filename, chunks=chunks)
3306            dt = xr.DataTree()
3307            if len(dss) == 1:
3308                dt.ds = dss[0]
3309                target_tree[pert_name] = dt
3310            elif len(dss) > 1:
3311                for ds in dss:
3312                    dt[f"pert{ds.data_vars[0]}"] = ds
3313            target_tree[pert_name] = dt
3314        except Exception as e:
3315            # Log error but continue processing other groups
3316            print(f"Error creating dataset for perturbation {pert_name}: {e}")
3317
3318    return success
3319
3320
3321def try_process_by_variables(
3322    target_tree: typing.Any,
3323    df: pd.DataFrame,
3324    filename: str,
3325    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
3326) -> bool:
3327    """
3328    Try to separate data by variable names and create datasets.
3329
3330    Parameters
3331    ----------
3332    target_tree : xarray.DataTree
3333        The tree node to add variable datasets to.
3334    df : pandas.DataFrame
3335        DataFrame of messages.
3336    filename : str
3337        Path to the GRIB2 file.
3338    chunks : int, dict or 'auto', optional
3339        If chunks is provided, it is used to load the dataset into a
3340        dask-backed dataset.
3341
3342    Returns
3343    -------
3344    bool
3345        True if at least one variable was successfully processed.
3346    """
3347    success = False
3348
3349    try:
3350        for var_name in df["shortName"].unique():
3351            if pd.notna(var_name):
3352                var_df = df[df["shortName"] == var_name]
3353                try:
3354                    var_ds = create_datasets_from_df(var_df, filename, chunks=chunks)
3355                    if var_ds is not None:
3356                        target_tree[f"var_{var_name}"] = var_ds[0]
3357                        success = True
3358                except Exception as var_e:
3359                    print(f"Error creating dataset for variable {var_name}: {var_e}")
3360    except Exception as nested_e:
3361        print(f"Failed to process variables: {nested_e}")
3362
3363    return success
3364
3365
3366def create_datasets_from_df(
3367    df: pd.DataFrame,
3368    filename: str,
3369    verbose: bool = False,
3370    chunks: typing.Optional[typing.Union[int, typing.Dict[typing.Any, typing.Any], typing.Literal["auto"]]] = None,
3371) -> typing.Optional[typing.List[xr.Dataset]]:
3372    """
3373    Create a list of xarray Datasets from a DataFrame of messages.
3374
3375    Parameters
3376    ----------
3377    df : pandas.DataFrame
3378        DataFrame of GRIB messages.
3379    filename : str
3380        Path to the GRIB2 file.
3381    verbose : bool, optional
3382        If True, prints detailed debugging information.
3383    chunks : int, dict or 'auto', optional
3384        If chunks is provided, it is used to load the dataset into a
3385        dask-backed dataset.
3386
3387    Returns
3388    -------
3389    list of xarray.Dataset, optional
3390        List of Datasets, or None if creation failed.
3391    """
3392    try:
3393        # Use parse_grib_index to get dimensions and attributes
3394        file_index, dim_coords, attrs, coord_attrs = parse_grib_index(df, {})
3395
3396        # Divide up records by variable
3397        frames, cubes, extra_geo = make_variables(file_index, filename, dim_coords, allow_uneven_dims=True)
3398
3399        if frames is None:
3400            return None
3401
3402        ds_list = []
3403        for var_df, var_cube in zip(frames, cubes):
3404            da = build_da_without_coords(var_df, var_cube, filename, attrs)
3405
3406            # Assign variable-specific coords from its cube
3407            coords = coords_from_cube(var_cube)
3408            da = da.assign_coords(coords)
3409
3410            # Assign extra index associated coords for this variable
3411            for dim_name, coord_names in dim_coords.items():
3412                retain_index_coord = False
3413                for name in coord_names:
3414                    if name == dim_name:
3415                        retain_index_coord = True
3416                    else:
3417                        if dim_name not in da.dims:
3418                            # for assigning scalar coords
3419                            coord_data = var_df[name].unique().item()
3420                            da = da.assign_coords({name: coord_data})
3421                        else:
3422                            # Handle non-scalar coords
3423                            coord_data = [
3424                                var_df[var_df.index.get_level_values(f"{dim_name}_ix") == val][name].unique().item()
3425                                for val in range(da[dim_name].size)
3426                            ]
3427                            coord = pd.Index(coord_data, name=dim_name)
3428                            da = da.assign_coords({name: (dim_name, coord)})
3429                if not retain_index_coord and dim_name in da.coords:
3430                    da = da.drop_vars(dim_name)
3431
3432            # Create a dataset for this variable
3433            var_ds = xr.Dataset({da.name: da})
3434
3435            # Assign metadata and common coords
3436            var_ds = assign_xr_meta(var_ds, [var_df], var_cube, dim_coords, extra_geo, coord_attrs)
3437
3438            if chunks is not None:
3439                var_ds = var_ds.chunk(chunks)
3440
3441            ds_list.append(var_ds)
3442
3443        return ds_list
3444    except Exception as e:
3445        if verbose:
3446            print(f"Error in create_datasets_from_df: {e}")
3447        return None
3448
3449
3450if _HAS_DATATREE:
3451
3452    @xr.register_datatree_accessor("grib2io")
3453    class Grib2ioDataTree:
3454        """
3455        DataTree accessor for GRIB2 files.
3456
3457        This accessor provides methods for working with GRIB2 data organized
3458        in a hierarchical tree structure.
3459        """
3460
3461        def __init__(self, datatree_obj):
3462            self._obj = datatree_obj
3463
3464        def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
3465            """
3466            Write all datasets in the DataTree to a GRIB2 file.
3467
3468            Parameters
3469            ----------
3470            filename : str
3471                Name of the GRIB2 file to write to.
3472            mode : {"x", "w", "a"}, optional
3473                Persistence mode, default is "x" (create, fail if exists)
3474            """
3475            # Start with the specified mode
3476            current_mode = mode
3477
3478            # Function to recursively process the tree
3479            def process_tree(node):
3480                nonlocal current_mode
3481
3482                # If this is a Dataset node with data variables
3483                if node.ds is not None and node.ds.data_vars:
3484                    # Write dataset to GRIB2 file
3485                    node.ds.grib2io.to_grib2(filename, mode=current_mode)
3486                    # Switch to append mode after first write
3487                    current_mode = "a"
3488
3489                # Process children
3490                for child_name, child_node in node.children.items():
3491                    process_tree(child_node)
3492
3493            # Start processing from the root
3494            process_tree(self._obj)
3495
3496        def griddef(self):
3497            """
3498            Get the grid definition from the first dataset in the tree that has one.
3499
3500            Returns
3501            -------
3502            grib2io.Grib2GridDef
3503                Grid definition object
3504            """
3505
3506            # Function to find first dataset with GRIB2IO_section3
3507            def find_griddef(node):
3508                if node.ds is not None and node.ds.data_vars:
3509                    for var_name in node.ds.data_vars:
3510                        if "GRIB2IO_section3" in node.ds[var_name].attrs:
3511                            return Grib2GridDef.from_section3(node.ds[var_name].attrs["GRIB2IO_section3"])
3512
3513                # Check children
3514                for child_name, child_node in node.children.items():
3515                    griddef = find_griddef(child_node)
3516                    if griddef is not None:
3517                        return griddef
3518
3519                return None
3520
3521            return find_griddef(self._obj)
3522
3523        def interp(
3524            self,
3525            method: typing.Union[str, int],
3526            grid_def_out: grib2io.Grib2GridDef,
3527            method_options: typing.Optional[typing.List[int]] = None,
3528            num_threads: int = 1,
3529        ) -> typing.Any:
3530            """
3531            Interpolate all datasets in the tree to a new grid.
3532
3533            Parameters
3534            ----------
3535            method : str or int
3536                Interpolation method to use.
3537            grid_def_out : grib2io.Grib2GridDef
3538                Target grid definition.
3539            method_options : list of int, optional
3540                Options for interpolation method.
3541            num_threads : int, optional
3542                Number of threads to use for interpolation.
3543
3544            Returns
3545            -------
3546            xarray.DataTree
3547                New DataTree with interpolated data.
3548            """
3549            new_tree = xr.DataTree()
3550
3551            # Function to recursively process the tree
3552            def process_tree(node, new_parent):
3553                # If this is a Dataset node with data variables
3554                if node.ds is not None and node.ds.data_vars:
3555                    # Interpolate dataset
3556                    interp_ds = node.ds.grib2io.interp(
3557                        method,
3558                        grid_def_out,
3559                        method_options=method_options,
3560                        num_threads=num_threads,
3561                    )
3562
3563                    # Add to new tree at the same path
3564                    if node == self._obj:  # Root node
3565                        new_parent.ds = interp_ds
3566                    else:
3567                        new_parent.ds = interp_ds
3568
3569                # Process children
3570                for child_name, child_node in node.children.items():
3571                    # Create same child in new tree
3572                    new_child = xr.DataTree()
3573                    new_parent[child_name] = new_child
3574                    process_tree(child_node, new_child)
3575
3576            # Start processing from the root
3577            process_tree(self._obj, new_tree)
3578
3579            return new_tree
3580
3581        def subset(self, lats: typing.Sequence[float], lons: typing.Sequence[float]) -> typing.Any:
3582            """
3583            Subset all datasets in the tree to a region.
3584
3585            Parameters
3586            ----------
3587            lats : sequence of float
3588                Latitude bounds [min_lat, max_lat].
3589            lons : sequence of float
3590                Longitude bounds [min_lon, max_lon].
3591
3592            Returns
3593            -------
3594            xarray.DataTree
3595                New DataTree with subset data.
3596            """
3597            new_tree = xr.DataTree()
3598
3599            # Function to recursively process the tree
3600            def process_tree(node, new_parent):
3601                # If this is a Dataset node with data variables
3602                if node.ds is not None and node.ds.data_vars:
3603                    # Subset dataset
3604                    subset_ds = node.ds.grib2io.subset(lats, lons)
3605
3606                    # Add to new tree at the same path
3607                    if node == self._obj:  # Root node
3608                        new_parent.ds = subset_ds
3609                    else:
3610                        new_parent.ds = subset_ds
3611
3612                # Process children
3613                for child_name, child_node in node.children.items():
3614                    # Create same child in new tree
3615                    new_child = xr.DataTree()
3616                    new_parent[child_name] = new_child
3617                    process_tree(child_node, new_child)
3618
3619            # Start processing from the root
3620            process_tree(self._obj, new_tree)
3621
3622            return new_tree