composite_reflectivity

#!/usr/bin/env python
# Copyright (c) 2020 Jon Thielen.
# Distributed under the terms of the Apache 2.0 License.
# SPDX-License-Identifier: Apache-2.0
"""
Simple plot of composite reflectivity (from UPP grib files).
------------------------------------------------------------

"""

from lib.grib import (
    metpy_parse_variable,
    open_datasets,
    search_variable_unique
)
from lib.params import (
    add_optional_argument,
    add_required_argument,
    config,
    create_parser,
    get_args
)
from lib.plots import (
    add_boundaries,
    cmaps,
    set_extent_max_min_x_y,
    single_panel_with_colorbar
)
from lib.vendor.xarray import _infer_interval_breaks

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

from metpy.units import units

# Command line arguments
parser = create_parser("Simple plot of composite reflectivity")
add_required_argument(
    parser,
    "grid_subdir",
    "g",
    "Subdirectory for grid of interest",
    "25km"
)
add_required_argument(
    parser,
    "init_hr",
    "i",
    "Initialization time for cycle",
    "0"
)
add_required_argument(
    parser,
    "f_hr",
    "f",
    "Forecast hour",
    "21"
)
add_optional_argument(
    parser,
    "base_dir",
    "b",
    "Base directory of files"
)
add_optional_argument(
    parser,
    "output",
    "O",
    "Output filename template",
    "comp_refl_{f_hr:02d}.png"
)
args = get_args(parser)

# Get the data
#datasets = open_datasets(config["upp_filepath_template"], collection="dawp", **args)
#refl = metpy_parse_variable(search_variable_unique(
#    datasets,
#    attributes={"long_name": "Maximum/Composite radar reflectivity"}
#))
ds = xr.open_dataset(
    config["upp_filepath_template"].format(collection='dawp', **args),
    engine='cfgrib',
    backend_kwargs={
        'indexpath': '',
        'filter_by_keys': {'typeOfLevel': 'unknown', 'stepType': 'instant'}
    }
)
refl = metpy_parse_variable(ds['refc'])

# Mask out low values
refl.data[refl.data < 0] = np.nan

# Create plot
fig, ax, cax = single_panel_with_colorbar(
    config=config,
    projection=refl.metpy.cartopy_crs
)
pcm = ax.pcolormesh(
    _infer_interval_breaks(refl['x']),
    _infer_interval_breaks(refl['y']),
    refl.data,
    cmap=cmaps["radar"],
    vmin=0,
    vmax=70,
    zorder=2
)
set_extent_max_min_x_y(ax, refl)
add_boundaries(ax)
cb = plt.colorbar(pcm, cax=cax, orientation="vertical")
cb.set_label("Composite radar reflectivity\ndBZ")
ax.set_title(f"Composite Reflectivity\nForecast Hour: {args['f_hr']}")

# Output plot
plt.savefig("output/" + args["output"].format(**args), dpi=300, bbox_inches="tight")