## Brian Blaylock
## April 23, 2021
"""Herbie's xarray accessors.
Other useful packages
- salem: mask data by a geographic region: https://github.com/fmaussion/salem
- xoak: xarray nearest neighbor https://github.com/xarray-contrib/xoak
"""
import functools
import pickle
import re
import warnings
from pathlib import Path
from typing import Literal, Optional, Union
import numpy as np
import pandas as pd
import xarray as xr
from pyproj import CRS
import herbie
_level_units = dict(
adiabaticCondensation="adiabatic condensation",
atmosphere="atmosphere",
atmosphereSingleLayer="atmosphere single layer",
boundaryLayerCloudLayer="boundary layer cloud layer",
cloudBase="cloud base",
cloudCeiling="cloud ceiling",
cloudTop="cloud top",
depthBelowLand="m",
equilibrium="equilibrium",
heightAboveGround="m",
heightAboveGroundLayer="m",
highCloudLayer="high cloud layer",
highestTroposphericFreezing="highest tropospheric freezing",
isobaricInhPa="hPa",
isobaricLayer="hPa",
isothermZero="0 C",
isothermal="K",
level="m",
lowCloudLayer="low cloud layer",
meanSea="MSLP",
middleCloudLayer="middle cloud layer",
nominalTop="nominal top",
pressureFromGroundLayer="Pa",
sigma="sigma",
sigmaLayer="sigmaLayer",
surface="surface",
)
def add_proj_info(ds: xr.Dataset):
"""Add projection info to a Dataset."""
raise NotImplementedError("This function `add_proj_info` is not yet implemented.")
# TODO: remove pyproj dependency
match = re.search(r'"source": "(.*?)"', ds.history)
FILE = Path(match.group(1))
# Get CF grid projection information with pygrib and pyproj because
# this is something cfgrib doesn't do (https://github.com/ecmwf/cfgrib/issues/251)
# NOTE: Assumes the projection is the same for all variables
with pygrib.open(str(FILE)) as grb:
msg = grb.message(1)
cf_params = CRS(msg.projparams).to_cf()
# Funny stuff with polar stereographic (https://github.com/pyproj4/pyproj/issues/856)
# TODO: Is there a better way to handle this? What about south pole?
if cf_params["grid_mapping_name"] == "polar_stereographic":
cf_params["latitude_of_projection_origin"] = cf_params.get(
"latitude_of_projection_origin", 90
)
# ----------------------
# Attach CF grid mapping
# ----------------------
# http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#appendix-grid-mappings
ds["gribfile_projection"] = None
ds["gribfile_projection"].attrs = cf_params
ds["gribfile_projection"].attrs["long_name"] = "model grid projection"
# Assign this grid_mapping for all variables
for var in list(ds):
if var == "gribfile_projection":
continue
ds[var].attrs["grid_mapping"] = "gribfile_projection"
[docs]
@xr.register_dataset_accessor("herbie")
class HerbieAccessor:
"""Accessor for xarray Datasets opened with Herbie."""
[docs]
def __init__(self, xarray_obj):
self._obj = xarray_obj
self._center = None
@property
def center(self) -> tuple[float, float]:
"""Return the geographic center point of this dataset."""
if self._center is None:
# we can use a cache on our accessor objects, because accessors
# themselves are cached on instances that access them.
lon = self._obj.latitude
lat = self._obj.longitude
self._center = (float(lon.mean()), float(lat.mean()))
return self._center
[docs]
def to_180(self) -> xr.Dataset:
"""Wrap longitude coordinates as range [-180,180]."""
ds = self._obj
ds["longitude"] = (ds["longitude"] + 180) % 360 - 180
return ds
[docs]
def to_360(self) -> xr.Dataset:
"""Wrap longitude coordinates as range [0,360]."""
ds = self._obj
ds["longitude"] = (ds["longitude"] - 360) % 360
return ds
@functools.cached_property
def crs(self):
"""
Cartopy coordinate reference system (crs) from a cfgrib Dataset.
Projection information is from the grib2 message for each variable.
Parameters
----------
ds : xarray.Dataset
An xarray.Dataset from a GRIB2 file opened by the cfgrib engine.
"""
try:
import metpy
except ModuleNotFoundError:
raise ModuleNotFoundError(
"metpy is an 'extra' requirement, please use "
"`pip install 'herbie-data[extras]'` for the full functionality."
)
ds = self._obj
# Get variables that have dimensions
# (this filters out the gribfile_projection variable)
variables = [i for i in list(ds) if len(ds[i].dims) > 0]
ds = ds.metpy.parse_cf(varname=variables)
crs = ds.metpy_crs.item().to_cartopy()
return crs
@functools.cached_property
def polygon(self):
"""Get a polygon of the domain boundary."""
try:
import cartopy.crs as ccrs
from shapely.geometry import Polygon
except ModuleNotFoundError:
raise ModuleNotFoundError(
"cartopy is an 'extra' requirements, please use "
"`pip install 'herbie-data[extras]'` for the full functionality."
)
ds = self._obj
LON = ds.longitude.data
LAT = ds.latitude.data
# Path of array outside border starting from the lower left corner
# and going around the array counter-clockwise.
outside = (
list(zip(LON[0, :], LAT[0, :]))
+ list(zip(LON[:, -1], LAT[:, -1]))
+ list(zip(LON[-1, ::-1], LAT[-1, ::-1]))
+ list(zip(LON[::-1, 0], LAT[::-1, 0]))
)
outside = np.array(outside)
###############################
# Polygon in Lat/Lon coordinates
x = outside[:, 0]
y = outside[:, 1]
domain_polygon_latlon = Polygon(zip(x, y))
###################################
# Polygon in projection coordinates
transform = self.crs.transform_points(ccrs.PlateCarree(), x, y)
# Remove any points that run off the projection map (i.e., point's value is `inf`).
transform = transform[~np.isinf(transform).any(axis=1)]
x = transform[:, 0]
y = transform[:, 1]
domain_polygon = Polygon(zip(x, y))
return domain_polygon, domain_polygon_latlon
[docs]
def with_wind(
self, which: Literal["both", "speed", "direction"] = "both"
) -> xr.Dataset:
"""Return Dataset with calculated wind speed and/or direction.
Consistent with the eccodes GRIB parameter database, variables
names are assigned as follows:
- "si10" : 10 metre wind speed (note this is not ws10 as you might expect)
- "wdir10" : 10 metre wind direction
- "ws" : wind speed
- "wdir" : wind direction
Refer to the eccodes database <https://codes.ecmwf.int/grib/param-db/>.
Parameters
----------
which : {'both', 'speed', 'direction'}
Specify which wind quantity to compute.
"""
ds = self._obj
n_computed = 0
if which in ("speed", "both"):
if {"u10", "v10"}.issubset(ds):
ds["si10"] = np.sqrt(ds.u10**2 + ds.v10**2)
ds["si10"].attrs["GRIB_paramId"] = 207
ds["si10"].attrs["long_name"] = "10 metre wind speed"
ds["si10"].attrs["units"] = "m s**-1"
ds["si10"].attrs["standard_name"] = "wind_speed"
ds["si10"].attrs["grid_mapping"] = ds.u10.attrs.get("grid_mapping")
n_computed += 1
if {"u100", "v100"}.issubset(ds):
ds["si100"] = np.sqrt(ds.u100**2 + ds.v100**2)
ds["si100"].attrs["GRIB_paramId"] = 228249
ds["si100"].attrs["long_name"] = "100 metre wind speed"
ds["si100"].attrs["units"] = "m s**-1"
ds["si100"].attrs["standard_name"] = "wind_speed"
ds["si100"].attrs["grid_mapping"] = ds.u100.attrs.get("grid_mapping")
n_computed += 1
if {"u80", "v80"}.issubset(ds):
ds["si80"] = np.sqrt(ds.u80**2 + ds.v80**2)
ds["si80"].attrs["long_name"] = "80 metre wind speed"
ds["si80"].attrs["units"] = "m s**-1"
ds["si80"].attrs["standard_name"] = "wind_speed"
ds["si80"].attrs["grid_mapping"] = ds.u80.attrs.get("grid_mapping")
n_computed += 1
if {"u", "v"}.issubset(ds):
ds["ws"] = np.sqrt(ds.u**2 + ds.v**2)
ds["ws"].attrs["GRIB_paramId"] = 10
ds["ws"].attrs["long_name"] = "wind speed"
ds["ws"].attrs["units"] = "m s**-1"
ds["ws"].attrs["standard_name"] = "wind_speed"
ds["ws"].attrs["grid_mapping"] = ds.u.attrs.get("grid_mapping")
n_computed += 1
if which in ("direction", "both"):
if {"u10", "v10"}.issubset(ds):
ds["wdir10"] = (
(270 - np.rad2deg(np.arctan2(ds.v10, ds.u10))) % 360
).where((ds.u10 != 0) & (ds.v10 != 0))
ds["wdir10"].attrs["GRIB_paramId"] = 260260
ds["wdir10"].attrs["long_name"] = "10 metre wind direction"
ds["wdir10"].attrs["units"] = "degree"
ds["wdir10"].attrs["standard_name"] = "wind_from_direction"
ds["wdir10"].attrs["grid_mapping"] = ds.u10.attrs.get("grid_mapping")
n_computed += 1
if {"u100", "v100"}.issubset(ds):
ds["wdir100"] = (
(270 - np.rad2deg(np.arctan2(ds.v100, ds.u100))) % 360
).where((ds.u100 != 0) & (ds.v100 != 0))
ds["wdir100"].attrs["long_name"] = "100 metre wind direction"
ds["wdir100"].attrs["units"] = "degree"
ds["wdir100"].attrs["standard_name"] = "wind_from_direction"
ds["wdir100"].attrs["grid_mapping"] = ds.u100.attrs.get("grid_mapping")
n_computed += 1
if {"u80", "v80"}.issubset(ds):
ds["wdir80"] = (
(270 - np.rad2deg(np.arctan2(ds.v80, ds.u80))) % 360
).where((ds.u80 != 0) & (ds.v80 != 0))
ds["wdir80"].attrs["long_name"] = "80 metre wind direction"
ds["wdir80"].attrs["units"] = "degree"
ds["wdir80"].attrs["standard_name"] = "wind_from_direction"
ds["wdir80"].attrs["grid_mapping"] = ds.u80.attrs.get("grid_mapping")
n_computed += 1
if {"u", "v"}.issubset(ds):
ds["wdir"] = ((270 - np.rad2deg(np.arctan2(ds.v, ds.u))) % 360).where(
(ds.u != 0) & (ds.v != 0)
)
ds["wdir"].attrs["GRIB_paramId"] = 3031
ds["wdir"].attrs["long_name"] = "wind direction"
ds["wdir"].attrs["units"] = "degree"
ds["wdir"].attrs["standard_name"] = "wind_from_direction"
ds["wdir"].attrs["grid_mapping"] = ds.u.attrs.get("grid_mapping")
n_computed += 1
if n_computed == 0:
warnings.warn("`with_wind()` did not do anything.")
return ds
[docs]
def pick_points(
self,
points: pd.DataFrame,
method: Literal["nearest", "weighted"] = "nearest",
*,
k: Optional[int] = None,
max_distance: Union[int, float] = 500,
use_cached_tree: Union[bool, Literal["replant"]] = True,
tree_name: Optional[str] = None,
verbose: bool = False,
) -> xr.Dataset:
"""Pick nearest neighbor grid values at selected points.
Parameters
----------
points : Pandas DataFrame
A DataFrame with columns 'latitude' and 'longitude'
representing the points to match to the model grid.
method : {'nearest', 'weighted'}
Method used to pick points.
- `nearest` : Gets grid value nearest the requested point.
- `weighted`: Gets four grid value nearest the requested
point and compute the inverse-distance-weighted mean.
k : None or int
If None and method is nearest, `k=1`.
If None and method is weighted, `k=4`.
Else, specify the number of neighbors to find.
max_distance : int or float
Maximum distance in kilometers allowed for nearest neighbor
search. Default is 500 km, which is very generous for any
model grid. This can help the case when a requested point
is off the grid.
use_cached_tree : {True, False, "replant"}
Controls if the BallTree object is caches for later use.
By "plant", I mean, "create a new BallTree object."
- `True` : Plant+save BallTree if it doesn't exist; load
saved BallTree if one exists.
- `False`: Plant the BallTree, even if one exists.
- `"replant"` : Plant a new BallTree and save a new pickle.
tree_name : str
If None, use the ds.model and domain size as the tree's name.
If ds.model does not exists, then the BallTree will not be
cached, unless you provide the tree_name.
Examples
--------
>>> H = Herbie("2024-03-28 00:00", model="hrrr")
>>> ds = H.xarray("TMP:[5,6,7,8,9][0,5]0 mb", remove_grib=False)
>>> points = pd.DataFrame(
... {
... "longitude": [-100, -105, -98.4],
... "latitude": [40, 29, 42.3],
... "stid": ["aa", "bb", "cc"],
... }
... )
Pick value at the nearest neighbor point
>>> dsp = ds.herbie.pick_points(points, method="nearest")
Get the weighted mean of the four nearest neighbor points
>>> dsp = ds.herbie.pick_points(points, method="weighted")
A Dataset is returned of the original grid reduced to the
requested points, with the values from the `points` dataset
added as new coordinates.
A user can easily convert the result to a Pandas DataFrame
>>> dsp.to_dataframe()
If you want to select points by a station name, swap the
dimension.
>>> dsp = dsp.swap_dims({"point": "point_stid"})
"""
try:
from sklearn.neighbors import BallTree
except ModuleNotFoundError:
raise ModuleNotFoundError(
"scikit-learn is an 'extra' requirement, please use "
"`pip install 'herbie-data[extras]'` for the full functionality."
)
def plant_tree(save_pickle: Optional[Union[Path, str]] = None):
"""Grow a new BallTree object from seedling."""
timer = pd.Timestamp("now")
print("INFO: 🌱 Growing new BallTree...", end="")
tree = BallTree(np.deg2rad(df_grid), metric="haversine")
print(
f"🌳 BallTree grew in {(pd.Timestamp('now')-timer).total_seconds():.2} seconds."
)
if save_pickle:
try:
Path(save_pickle).parent.mkdir(parents=True, exist_ok=True)
with open(save_pickle, "wb") as f:
pickle.dump(tree, f)
print(f"INFO: Saved BallTree to {save_pickle}")
except OSError:
print(f"ERROR: Could not save BallTree to {save_pickle}.")
return tree
ds = self._obj
# ---------------------
# Validate points input
if ("latitude" not in points) and ("longitude" not in points):
raise ValueError(
"`points` DataFrame must have columns 'latitude' and 'longitude'"
)
if not all(points.latitude.between(-90, 90, inclusive="both")):
raise ValueError("All latitude points must be [-90,90]")
if not all(points.longitude.between(0, 360, inclusive="both")):
if not all(points.longitude.between(-180, 180, inclusive="both")):
raise ValueError("All longitude points must be [-180,180] or [0,360]")
# ---------------------
# Validate method input
_method = set(["nearest", "weighted"])
if method == "nearest" and k is None:
# Get the value at the nearest grid point using BallTree
k = 1
elif method == "weighted" and k is None:
# Compute the value of each variable from the inverse-
# weighted distance of the values of the four nearest
# neighbors.
k = 4
elif method in _method and isinstance(k, int):
# Get the k nearest neighbors and return the values (nearest)
# or compute the distance-weighted mean (weighted).
pass
else:
raise ValueError(
f"`method` must be one of {_method} and `k` must be an int or None."
)
# Only consider variables that have dimensions.
ds = ds[[i for i in ds if ds[i].dims != ()]]
if "latitude" in ds.dims and "longitude" in ds.dims:
# Rename dims to x and y
# This is needed for regular latitude-longitude grids like
# GFS and IFS model data.
ds = ds.rename_dims({"latitude": "y", "longitude": "x"})
# Get Dataset's lat/lon grid and coordinate indices as a DataFrame.
df_grid = (
ds[["latitude", "longitude"]]
.drop_vars([i for i, j in ds.coords.items() if not j.ndim])
.to_dataframe()
)
# ---------------
# BallTree object
# Plant, plant+Save, or load
if tree_name is None:
tree_name = getattr(ds, "model", "UNKNOWN")
if use_cached_tree and tree_name == "UNKNOWN":
use_cached_tree = False
print(
"WARNING: Herbie won't cache the BallTree because it\n"
" doesn't know what to name it. Please specify\n"
" `tree_name` to cache the tree for use later."
)
pkl_BallTree_file = (
herbie.config["default"]["save_dir"]
/ "BallTree"
/ f"{tree_name}_{ds.x.size}-{ds.y.size}.pkl"
)
if not use_cached_tree:
# Create a new BallTree. Do not save pickle.
tree = plant_tree(save_pickle=False)
elif use_cached_tree == "replant" or not pkl_BallTree_file.exists():
# Create a new BallTree and save pickle.
tree = plant_tree(save_pickle=pkl_BallTree_file)
elif use_cached_tree:
# Load BallTree from pickle.
with open(pkl_BallTree_file, "rb") as f:
tree = pickle.load(f)
# -------------------------------------
# Query points to find nearest neighbor
# Note: Order matters, and lat/long must be in radians.
# TODO: Maybe add option to use MultiProcessing here, to split
# TODO: the Dataset into chunks; or maybe not needed because
# TODO: the method is fast enough without the added complexity.
dist, ind = tree.query(np.deg2rad(points[["latitude", "longitude"]]), k=k)
# Convert distance to km by multiplying by the radius of the Earth
dist *= 6371
# Pick grid values for each value of k
k_points = []
df_grid = df_grid.reset_index()
for i in range(k):
a = points.copy()
a["point_grid_distance"] = dist[:, i]
a["grid_index"] = ind[:, i]
a = pd.concat(
[
a,
df_grid.iloc[a.grid_index]
.add_suffix("_grid")
.reset_index(drop=True),
],
axis=1,
)
a.index.name = "point"
if max_distance:
flagged = a.loc[a.point_grid_distance > max_distance]
a = a.loc[a.point_grid_distance <= max_distance]
if len(flagged):
print(
f"WARNING: {len(flagged)} points removed for exceeding {max_distance=} km threshold."
)
print(f"{flagged}")
print("")
# Get corresponding values from xarray
# https://docs.xarray.dev/en/stable/user-guide/indexing.html#more-advanced-indexing
ds_points = ds.sel(
x=a.x_grid.to_xarray(),
y=a.y_grid.to_xarray(),
)
ds_points.coords["point_grid_distance"] = a.point_grid_distance.to_xarray()
ds_points["point_grid_distance"].attrs["long_name"] = (
"Distance between requested point and nearest grid point."
)
ds_points["point_grid_distance"].attrs["units"] = "km"
for i in points.columns:
ds_points.coords[f"point_{i}"] = a[i].to_xarray()
ds_points[f"point_{i}"].attrs["long_name"] = f"Requested grid point {i}"
k_points.append(ds_points.drop_vars("point"))
if method == "nearest" and k == 1:
return k_points[0]
elif method == "nearest" and k > 1:
# New dimension k is the index of the n-th nearest neighbor
return xr.concat(k_points, dim="k")
elif method == "weighted":
# Compute the inverse-distance weighted mean for each
# variable from the four nearest points.
b = xr.concat(k_points, dim="k")
# Note: clipping accounts for the "divide by zero" case when
# the requested point is exactly the nearest grid point.
weights = (1 / b.point_grid_distance).clip(max=1e6)
# Compute weighted mean of variables
sum_of_weights = weights.sum(dim="k")
weighted_sum = (b * weights).sum(dim="k")
c = weighted_sum / sum_of_weights
# Include some coordinates that were dropped as a result of
# the line `weights.sum(dim='k')`.
c.coords["latitude"] = b.coords["latitude"]
c.coords["longitude"] = b.coords["longitude"]
c.coords["point_grid_distance"] = b.coords["point_grid_distance"]
return c
else:
raise ValueError("I didn't expect to be here.")
[docs]
def nearest_points(self, points, names=None, verbose=True):
"""
Get the nearest latitude/longitude points from a xarray Dataset.
- Stack Overflow: https://stackoverflow.com/questions/58758480/xarray-select-nearest-lat-lon-with-multi-dimension-coordinates
- MetPy Details: https://unidata.github.io/MetPy/latest/tutorials/xarray_tutorial.html?highlight=assign_y_x
Parameters
----------
ds : xr.Dataset
A Herbie-friendly xarray Dataset
points : tuple, list of tuples, pd.DataFrame
Points to be plucked from the gridded Dataset.
There are multiple objects accepted.
1. Tuple of longitude and latitude (lon, lat) coordinate pair.
1. List of multiple (lon, lat) coordinate pair tuples.
1. Pandas DataFrame with ``longitude`` and ``latitude`` columns. Index will be used as point names, unless ``names`` is specified.
1. Shapeley Point or Points
names : list
A list of names for each point location (i.e., station name).
None will not append any names. names should be the same
length as points.
Notes
-----
This is **much** faster than my old "pluck_points" method.
For matching 1,948 points:
- `nearest_points` completed in 7.5 seconds.
- `pluck_points` completed in 2 minutes.
TODO: Explore alternatives
- Could Shapely nearest_points be used
https://shapely.readthedocs.io/en/latest/manual.html#nearest-points
- Or possibly scipy BallTree method.
"""
try:
import cartopy.crs as ccrs
import shapely
from shapely.geometry import MultiPoint, Point
except ModuleNotFoundError:
raise ModuleNotFoundError(
"cartopy is an 'extra' requirements, please use "
"`pip install 'herbie-data[extras]'` for the full functionality."
)
warnings.warn(
"The accessor `ds.herbie.nearest_points` is deprecated in "
"favor of the `ds.herbie.pick_points` which uses the "
"BallTree algorithm instead.",
DeprecationWarning,
stacklevel=2,
)
ds = self._obj
# Longitude and Latitude point DataFrame
if isinstance(points, pd.DataFrame):
point_df = points[["longitude", "latitude"]]
if names is not None:
point_df.index = names
elif np.shape(points) == (2,):
# points is a tuple (lon, lat) or list [lon, lat]
# and name is given as None or str
point_df = pd.DataFrame(
[points],
columns=["longitude", "latitude"],
index=[names],
)
elif isinstance(points, list):
# points given as a list of coordinate-pair tuples
# and name is given as a list of str
point_df = pd.DataFrame(
points,
columns=["longitude", "latitude"],
index=names,
)
elif isinstance(points, (MultiPoint, Point)):
# points is given as a Shapely object
point_df = pd.DataFrame(
shapely.get_coordinates(points),
columns=["longitude", "latitude"],
index=names,
)
else:
raise ValueError("The points supplied was not understood.")
# Check if MetPy has already parsed the CF metadata grid projection.
# Do that if it hasn't been done yet.
if "metpy_crs" not in ds:
ds = ds.metpy.parse_cf()
# Apply the MetPy method `assign_y_x` to the dataset
# https://unidata.github.io/MetPy/latest/api/generated/metpy.xarray.html?highlight=assign_y_x#metpy.xarray.MetPyDataArrayAccessor.assign_y_x
ds = ds.metpy.assign_y_x()
# Convert the requested [(lon,lat), (lon,lat)] points to map projection.
# Accept a list of point tuples, or Shapely Points object.
# We want to index the dataset at a single point.
# We can do this by transforming a lat/lon point to the grid location
crs = ds.metpy_crs.item().to_cartopy()
transformed_data = crs.transform_points(
ccrs.PlateCarree(), point_df.longitude, point_df.latitude
)
a = pd.DataFrame({"x": transformed_data[:, 0], "y": transformed_data[:, 1]})
a.index.name = "point"
# Select the nearest points from the projection coordinates.
# Get corresponding values from xarray
# https://docs.xarray.dev/en/stable/user-guide/indexing.html#more-advanced-indexing
#
new_ds = ds.sel(
x=a["x"].to_xarray(),
y=a["y"].to_xarray(),
method="nearest",
)
new_ds.coords["point"] = ("point", point_df.index.to_list())
new_ds.coords["point_latitude"] = ("point", point_df.latitude)
new_ds.coords["point_longitude"] = ("point", point_df.longitude)
return new_ds
[docs]
def plot(self, ax=None, common_features_kw={}, vars=None, **kwargs):
"""Plot data on a map.
TODO: Work in progress!
Parameters
----------
vars : list
List of variables to plot. Default None will plot all
variables in the DataSet.
"""
raise NotImplementedError("Plotting functionality is not working right now.")
try:
import matplotlib.pyplot as plt
from herbie import paint
from herbie.toolbox import EasyMap, pc
except ModuleNotFoundError:
raise ModuleNotFoundError(
"cartopy is an 'extra' requirement. Please use "
"`pip install 'herbie-data[extras]'` for the full functionality."
)
ds = self._obj
if isinstance(vars, str):
vars = [vars]
if vars is None:
vars = ds.data_vars
for i, var in enumerate(vars):
if "longitude" not in ds[var].coords:
# This is the case for the gribfile_projection variable
continue
print("cfgrib variable:", var)
print("GRIB_cfName", ds[var].attrs.get("GRIB_cfName"))
print("GRIB_cfVarName", ds[var].attrs.get("GRIB_cfVarName"))
print("GRIB_name", ds[var].attrs.get("GRIB_name"))
print("GRIB_units", ds[var].attrs.get("GRIB_units"))
print("GRIB_typeOfLevel", ds[var].attrs.get("GRIB_typeOfLevel"))
print()
ds[var].attrs["units"] = (
ds[var]
.attrs["units"]
.replace("**-1", "$^{-1}$")
.replace("**-2", "$^{-2}$")
)
defaults = dict(
scale="50m",
dpi=150,
figsize=(10, 5),
crs=ds.herbie.crs,
ax=ax,
)
common_features_kw = {**defaults, **common_features_kw}
ax = EasyMap(fignum=i, **common_features_kw).STATES().ax
title = ""
kwargs.setdefault("shading", "auto")
cbar_kwargs = dict(pad=0.01)
if ds[var].GRIB_cfVarName in ["d2m", "dpt"]:
ds[var].attrs["GRIB_cfName"] = "dew_point_temperature"
## Wind
wind_pair = {"u10": "v10", "u80": "v80", "u": "v"}
if ds[var].GRIB_cfName == "air_temperature":
kwargs = {**paint.NWSTemperature.kwargs2, **kwargs}
cbar_kwargs = {**paint.NWSTemperature.cbar_kwargs2, **cbar_kwargs}
if ds[var].GRIB_units == "K":
ds[var] -= 273.15
ds[var].attrs["GRIB_units"] = "C"
ds[var].attrs["units"] = "C"
elif ds[var].GRIB_cfName == "dew_point_temperature":
kwargs = {**paint.NWSDewPointTemperature.kwargs2, **kwargs}
cbar_kwargs = {
**paint.NWSDewPointTemperature.cbar_kwargs2,
**cbar_kwargs,
}
if ds[var].GRIB_units == "K":
ds[var] -= 273.15
ds[var].attrs["GRIB_units"] = "C"
ds[var].attrs["units"] = "C"
elif ds[var].GRIB_name == "Total Precipitation":
title = "-".join(
[f"F{int(i):02d}" for i in ds[var].GRIB_stepRange.split("-")]
)
ds[var] = ds[var].where(ds[var] != 0)
kwargs = {**paint.NWSPrecipitation.kwargs2, **kwargs}
cbar_kwargs = {**paint.NWSPrecipitation.cbar_kwargs2, **cbar_kwargs}
elif ds[var].GRIB_name == "Maximum/Composite radar reflectivity":
ds[var] = ds[var].where(ds[var] >= 0)
kwargs = {**paint.RadarReflectivity.kwargs2, **kwargs}
cbar_kwargs = {**paint.RadarReflectivity.cbar_kwargs2, **cbar_kwargs}
elif ds[var].GRIB_cfName == "relative_humidity":
kwargs = {**paint.NWSRelativeHumidity.kwargs2, **kwargs}
cbar_kwargs = {**paint.NWSRelativeHumidity.cbar_kwargs2, **cbar_kwargs}
elif ds[var].GRIB_name == "Orography":
if "lsm" in ds:
ds["orog"] = ds.orog.where(ds.lsm == 1, -100)
kwargs = {**paint.LandGreen.kwargs, **kwargs}
# cbar_kwargs = {**cm_terrain().cbar_kwargs, **cbar_kwargs}
elif "wind" in ds[var].GRIB_cfName or "wind" in ds[var].GRIB_name:
cbar_kwargs = {**cm_wind().cbar_kwargs, **cbar_kwargs}
kwargs = {**cm_wind().cmap_kwargs, **kwargs}
if ds[var].GRIB_cfName == "eastward_wind":
cbar_kwargs["label"] = "U " + cbar_kwargs["label"]
elif ds[var].GRIB_cfName == "northward_wind":
cbar_kwargs["label"] = "V " + cbar_kwargs["label"]
else:
cbar_kwargs = {
**dict(
label=f"{ds[var].GRIB_parameterName.strip().title()} ({ds[var].units})"
),
**cbar_kwargs,
}
p = ax.pcolormesh(
ds.longitude, ds.latitude, ds[var], transform=pc, **kwargs
)
plt.colorbar(p, ax=ax, **cbar_kwargs)
VALID = pd.to_datetime(ds.valid_time.data).strftime("%H:%M UTC %d %b %Y")
RUN = pd.to_datetime(ds.time.data).strftime("%H:%M UTC %d %b %Y")
FXX = f"F{pd.to_timedelta(ds.step.data).total_seconds()/3600:02.0f}"
level_type = ds[var].GRIB_typeOfLevel
if level_type in _level_units:
level_units = _level_units[level_type]
else:
level_units = "unknown"
if level_units.lower() in ["surface", "atmosphere"]:
level = f"{title} {level_units}"
else:
level = f"{ds[var][level_type].data:g} {level_units}"
ax.set_title(
f"Run: {RUN} {FXX}",
loc="left",
fontfamily="monospace",
fontsize="x-small",
)
ax.set_title(
f"{ds.model.upper()} {level}\n", loc="center", fontweight="semibold"
)
ax.set_title(
f"Valid: {VALID}",
loc="right",
fontfamily="monospace",
fontsize="x-small",
)
# Set extent so no whitespace shows around pcolormesh area
# TODO: Any better way to do this? With metpy.assign_y_x
# !!!!: The `metpy.assign_y_x` method could be used for pluck_point :)
try:
if "x" in ds.dims:
ds = ds.metpy.parse_cf()
ds = ds.metpy.assign_y_x()
ax.set_extent(
[
ds.x.min().item(),
ds.x.max().item(),
ds.y.min().item(),
ds.y.max().item(),
],
crs=ds.herbie.crs,
)
except Exception:
pass
return ax
