Source code for radarx.vis.plots

#!/usr/bin/env python
# Copyright (c) 2024-2026, Radarx developers.
# Distributed under the MIT License. See LICENSE for more info.

"""
Radarx Plot Helpers
===================

.. autosummary::
   :nosignatures:
   :toctree: generated/

   {}
"""

from __future__ import annotations

__all__ = ["plot_cappi", "plot_ppi", "plot_rhi"]

__doc__ = __doc__.format("\n   ".join(__all__))

import os

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


def _get_dataarray(ds, data_var):
    if isinstance(ds, xr.DataArray):
        return ds
    if data_var is None:
        raise ValueError("data_var must be provided when plotting an xarray.Dataset.")
    return ds[data_var]


def _get_figure_and_axis(ax=None, figsize=(7, 6)):
    if ax is not None:
        return ax.figure, ax
    return plt.subplots(figsize=figsize)


def _get_coord(da, name):
    if name in da.coords:
        return np.asarray(da.coords[name].values, dtype=float)
    raise ValueError(f"Coordinate '{name}' is required for this plot.")


def _get_time_token(ds):
    if "time" not in ds:
        return "unknown_time"

    value = np.asarray(ds["time"].values).reshape(-1)[0]
    value_da = xr.DataArray(value)
    if np.issubdtype(value_da.dtype, np.datetime64):
        return value_da.dt.strftime("%Y%m%d%H%M%S").item()
    return str(value)


def _get_radar_name(ds):
    radar_name = (
        ds.attrs.get("instrument_name") or ds.attrs.get("radar_name") or "Radar"
    )
    if isinstance(radar_name, bytes):
        radar_name = radar_name.decode()
    return str(radar_name)


def _finalize_plot(
    fig, ax, ds, title, savedir=None, dpi=100, show_figure=True, add_slogan=False
):
    if add_slogan:
        fig.text(
            0.5,
            0.02,
            "Plot by Radarx | Powered by Xradar",
            fontsize=9,
            fontname="Courier New",
            ha="center",
        )

    if savedir:
        filename = f"{title}_{_get_radar_name(ds)}_{_get_time_token(ds)}.png"
        filepath = os.path.join(savedir, filename)
        fig.savefig(filepath, dpi=dpi, bbox_inches="tight")

    if show_figure:
        plt.show()
    else:
        plt.close(fig)

    return ax


def _plot_plan_view(
    ds,
    data_var,
    title,
    cmap=None,
    vmin=None,
    vmax=None,
    colorbar=True,
    ax=None,
    dpi=100,
    savedir=None,
    show_figure=True,
    add_slogan=False,
    **kwargs,
):
    da = _get_dataarray(ds, data_var)
    x = _get_coord(da, "x") / 1000.0
    y = _get_coord(da, "y") / 1000.0

    kwargs.setdefault("shading", "nearest")

    fig, ax = _get_figure_and_axis(ax=ax)
    mesh = ax.pcolormesh(
        x, y, np.asarray(da.values), cmap=cmap, vmin=vmin, vmax=vmax, **kwargs
    )
    ax.set_aspect("equal")
    ax.set_xlabel("x (km)")
    ax.set_ylabel("y (km)")
    ax.set_title(title)

    if colorbar:
        colorbar_obj = fig.colorbar(mesh, ax=ax)
        colorbar_obj.set_label(da.attrs.get("units", ""))

    return _finalize_plot(
        fig,
        ax,
        ds if isinstance(ds, xr.Dataset) else da.to_dataset(name=da.name or data_var),
        title,
        savedir=savedir,
        dpi=dpi,
        show_figure=show_figure,
        add_slogan=add_slogan,
    )


[docs] def plot_ppi( ds, data_var, cmap=None, vmin=None, vmax=None, title=None, colorbar=True, ax=None, dpi=100, savedir=None, show_figure=True, add_slogan=False, **kwargs, ): """ Plot a georeferenced plan-position indicator (PPI)-style view. Parameters ---------- ds : xarray.Dataset or xarray.DataArray Sweep dataset containing 2D data on georeferenced ``x`` and ``y`` coordinates. data_var : str Variable name to plot when ``ds`` is a dataset. """ title = title or f"PPI {data_var}" return _plot_plan_view( ds, data_var, title=title, cmap=cmap, vmin=vmin, vmax=vmax, colorbar=colorbar, ax=ax, dpi=dpi, savedir=savedir, show_figure=show_figure, add_slogan=add_slogan, **kwargs, )
[docs] def plot_cappi( ds, data_var, cmap=None, vmin=None, vmax=None, title=None, colorbar=True, ax=None, dpi=100, savedir=None, show_figure=True, add_slogan=False, **kwargs, ): """ Plot a CAPPI dataset on the horizontal ``x``/``y`` plane. Parameters ---------- ds : xarray.Dataset or xarray.DataArray CAPPI dataset containing horizontal ``x`` and ``y`` coordinates. data_var : str Variable name to plot when ``ds`` is a dataset. """ height = None if "z" in ds: height = float(np.asarray(ds["z"].values).reshape(-1)[0]) / 1000.0 title = title or ( f"CAPPI {height:.1f} km {data_var}" if height is not None else f"CAPPI {data_var}" ) return _plot_plan_view( ds, data_var, title=title, cmap=cmap, vmin=vmin, vmax=vmax, colorbar=colorbar, ax=ax, dpi=dpi, savedir=savedir, show_figure=show_figure, add_slogan=add_slogan, **kwargs, )
[docs] def plot_rhi( ds, data_var, cmap=None, vmin=None, vmax=None, title=None, colorbar=True, ax=None, dpi=100, savedir=None, show_figure=True, add_slogan=False, **kwargs, ): """ Plot a range-height indicator (RHI)-style vertical section. Parameters ---------- ds : xarray.Dataset or xarray.DataArray Sweep dataset containing georeferenced ``x``, ``y``, and ``z`` coordinates. data_var : str Variable name to plot when ``ds`` is a dataset. """ da = _get_dataarray(ds, data_var) x = _get_coord(da, "x") y = _get_coord(da, "y") z = _get_coord(da, "z") ground_range = np.hypot(x, y) / 1000.0 height = z / 1000.0 kwargs.setdefault("shading", "auto") fig, ax = _get_figure_and_axis(ax=ax) mesh = ax.pcolormesh( ground_range, height, np.asarray(da.values), cmap=cmap, vmin=vmin, vmax=vmax, **kwargs, ) ax.set_xlabel("Ground Range (km)") ax.set_ylabel("Height (km)") ax.set_title(title or f"RHI {data_var}") if colorbar: colorbar_obj = fig.colorbar(mesh, ax=ax) colorbar_obj.set_label(da.attrs.get("units", "")) return _finalize_plot( fig, ax, ds if isinstance(ds, xr.Dataset) else da.to_dataset(name=da.name or data_var), title or f"RHI {data_var}", savedir=savedir, dpi=dpi, show_figure=show_figure, add_slogan=add_slogan, )