Skip to content

Commit

Permalink
Added new da.grib2io.update_attrs to update message attributes
Browse files Browse the repository at this point in the history 
New function updates the DataArray attrs section dictionaries
using the message attribute setters.
  • Loading branch information
@TimothyCera-NOAA
TimothyCera-NOAA committed Jul 26, 2024
1 parent 242c96c commit 6c81bd0
Show file tree
Hide file tree
Showing 3 changed files with 416 additions and 14 deletions.
24 changes: 21 additions & 3 deletions src/grib2io/templates.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,41 +169,47 @@ def __get__(self, obj, objtype=None):
return obj.section1[5]
def __set__(self, obj, value):
obj.section1[5] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class Month:
"""Month of reference time"""
def __get__(self, obj, objtype=None):
return obj.section1[6]
def __set__(self, obj, value):
obj.section1[6] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class Day:
"""Day of reference time"""
def __get__(self, obj, objtype=None):
return obj.section1[7]
def __set__(self, obj, value):
obj.section1[7] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class Hour:
"""Hour of reference time"""
def __get__(self, obj, objtype=None):
return obj.section1[8]
def __set__(self, obj, value):
obj.section1[8] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class Minute:
"""Minute of reference time"""
def __get__(self, obj, objtype=None):
return obj.section1[9]
def __set__(self, obj, value):
obj.section1[9] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class Second:
"""Second of reference time"""
def __get__(self, obj, objtype=None):
return obj.section1[10]
def __set__(self, obj, value):
obj.section1[10] = value
setattr(obj, "refDate", datetime.datetime(*obj.section1[5:11]))

class RefDate:
"""Reference Date. NOTE: This is a `datetime.datetime` object."""
Expand Down Expand Up @@ -927,21 +933,33 @@ class FullName:
def __get__(self, obj, objtype=None):
return tables.get_varinfo_from_table(obj.section0[2],*obj.section4[2:4],isNDFD=obj._isNDFD)[0]
def __set__(self, obj, value):
raise RuntimeError
raise RuntimeError(
"Cannot set the fullName of the message. Instead set shortName OR set the appropriate discipline, parameterCategory, and parameterNumber. The fullName will be set automatically from these other attributes."
)

class Units:
"""Units of the Variable."""
def __get__(self, obj, objtype=None):
return tables.get_varinfo_from_table(obj.section0[2],*obj.section4[2:4],isNDFD=obj._isNDFD)[1]
def __set__(self, obj, value):
raise RuntimeError
raise RuntimeError(
"Cannot set the units of the message. Instead set shortName OR set the appropriate discipline, parameterCategory, and parameterNumber. The units will be set automatically from these other attributes."
)

class ShortName:
""" Short name of the variable (i.e. the variable abbreviation)."""
def __get__(self, obj, objtype=None):
return tables.get_varinfo_from_table(obj.section0[2],*obj.section4[2:4],isNDFD=obj._isNDFD)[2]
def __set__(self, obj, value):
raise RuntimeError
metadata = tables.get_metadata_from_shortname(value)
if len(metadata) > 1:
raise ValueError(
f"shortName={value} is ambiguous within the GRIB2 standard and you have to set instead with discipline, parameterCategory, and parameterNumber.\n{metadata}"
)
for attr, val in metadata[0].items():
if attr in ["fullName", "units"]:
continue
setattr(obj, attr, val)

class TypeOfGeneratingProcess:
"""[Type of Generating Process](https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-3.shtml)"""
Expand Down
134 changes: 123 additions & 11 deletions src/grib2io/xarray_backend.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import itertools
import logging
import typing
from warnings import warn

import numpy as np
import pandas as pd
Expand All @@ -27,6 +28,17 @@

LOCK = SerializableLock()

AVAILABLE_NON_GEO_DIMS = [
"duration",
"leadTime",
"percentileValue",
"perturbationNumber",
"refDate",
"thresholdLowerLimit",
"thresholdUpperLimit",
"valueOfFirstFixedSurface",
]


class GribBackendEntrypoint(BackendEntrypoint):
"""
Expand Down Expand Up @@ -749,6 +761,19 @@ def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
da.grib2io.to_grib2(filename, mode=mode)
mode = "a"

def update_attrs(self, **kwargs):
"""
Raises an error because Datasets don't have a .attrs attribute.
Parameters
----------
attrs
Attributes to update.
"""
raise ValueError(
f"Datasets do not have a .attrs attribute; use .grib2io.update_attrs({kwargs}) on a DataArray instead."
)


@xr.register_dataarray_accessor("grib2io")
class Grib2ioDataArray:
Expand Down Expand Up @@ -934,19 +959,9 @@ def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
"""
da = self._obj.copy(deep=True)

available_non_geo_dims = [
"duration",
"leadTime",
"percentileValue",
"perturbationNumber",
"refDate",
"thresholdLowerLimit",
"thresholdUpperLimit",
"valueOfFirstFixedSurface",
]

coords_keys = sorted(da.coords.keys())
coords_keys = [k for k in coords_keys if k in available_non_geo_dims]
coords_keys = [k for k in coords_keys if k in AVAILABLE_NON_GEO_DIMS]

# If there are dimension coordinates, the DataArray is a hypercube of
# grib2 messages.
Expand Down Expand Up @@ -1010,7 +1025,104 @@ def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
for index in [i for i in coords_keys if i not in da.dims]:
setattr(newmsg, index, selected.coords[index].values)

# Set section 5 attributes to the da.encoding dictionary.
for key, value in selected.encoding.items():
if key in ["dtype", "chunks", "original_shape"]:
continue
setattr(newmsg, key, value)

# write the message to file
with grib2io.open(filename, mode=mode) as f:
f.write(newmsg)
mode = "a"

def update_attrs(self, **kwargs):
"""
Update many of the attributes of the DataArray.
Parameters
----------
**kwargs
Attributes to update. This can include many of the GRIB2IO message
attributes that you can find when you print a GRIB2IO message. For
conflicting updates, the last keyword will be used.
+-----------------------+------------------------------------------+
| kwargs | Description |
+=======================+==========================================+
| shortName="VTMP" | Set shortName to "VTMP", along with |
| | appropriate discipline, |
| | parameterCategory, parameterNumber, |
| | fullName and units. |
+-----------------------+------------------------------------------+
| discipline=0, | Set shortName, discipline, |
| parameterCategory=0, | parameterCategory, parameterNumber, |
| parameterNumber=1 | fullName and units appropriate for |
| | "Virtual Temperature". |
+-----------------------+------------------------------------------+
| discipline=0, | Conflicting keywords but |
| parameterCategory=0, | 'shortName="TMP"' wins. Set shortName, |
| parameterNumber=1, | discipline, parameterCategory, |
| shortName="TMP" | parameterNumber, fullName and units |
| | appropriate for "Temperature". |
+-----------------------+------------------------------------------+
Returns
-------
DataArray
DataArray with updated attributes.
"""
da = self._obj.copy(deep=True)

newmsg = Grib2Message(
da.attrs["GRIB2IO_section0"],
da.attrs["GRIB2IO_section1"],
da.attrs["GRIB2IO_section2"],
da.attrs["GRIB2IO_section3"],
da.attrs["GRIB2IO_section4"],
da.attrs["GRIB2IO_section5"],
)

coords_keys = [
k
for k in da.coords.keys()
if (k in AVAILABLE_NON_GEO_DIMS) and (k in da.dims)
]

for grib2_name, value in kwargs.items():
if grib2_name == "gridDefinitionTemplateNumber":
raise ValueError(
"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."
)
if grib2_name == "productDefinitionTemplateNumber":
raise ValueError(
"The productDefinitionTemplateNumber attribute cannot be updated."
)
if grib2_name == "dataRepresentationTemplateNumber":
raise ValueError(
"The dataRepresentationTemplateNumber attribute cannot be updated."
)
if grib2_name in coords_keys:
warn(
f"Skipping attribute '{grib2_name}' because it is a dimension coordinate and cannot be updated."
)
continue
if hasattr(newmsg, grib2_name):
setattr(newmsg, grib2_name, value)
else:
warn(
f"Skipping attribute '{grib2_name}' because it is not a valid GRIB2 attribute for this message and cannot be updated."
)
continue

da.attrs["GRIB2IO_section0"] = newmsg.section0
da.attrs["GRIB2IO_section1"] = newmsg.section1
da.attrs["GRIB2IO_section2"] = newmsg.section2 or []
da.attrs["GRIB2IO_section3"] = newmsg.section3
da.attrs["GRIB2IO_section4"] = newmsg.section4
da.attrs["GRIB2IO_section5"] = newmsg.section5
da.attrs["fullName"] = newmsg.fullName
da.attrs["shortName"] = newmsg.shortName
da.attrs["units"] = newmsg.units

return da
Loading

0 comments on commit 6c81bd0

Please sign in to comment.