diff --git a/docs/_toc.yml b/docs/_toc.yml index 293ef40..5523d4a 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -9,3 +9,4 @@ chapters: - file: it_dpc_precipitation - file: uk_metoffice_precipitation - file: be_rmi_precipitation +- file: plotting diff --git a/docs/plotting.ipynb b/docs/plotting.ipynb new file mode 100644 index 0000000..afb0603 --- /dev/null +++ b/docs/plotting.ipynb @@ -0,0 +1,213 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Plotting & Diagnostics\n", + "\n", + "The `mlcast_datasets.plotting` subpackage provides publication-quality diagnostic\n", + "figures for any dataset in the MLCast catalog. This notebook demonstrates every\n", + "plotting function using the Italian DPC SRI radar precipitation dataset as an\n", + "example.\n", + "\n", + "All sample sizes are kept deliberately small so that this notebook runs quickly\n", + "in CI. Increase `n_samples` / `n_coverage_samples` for higher-quality figures.\n", + "\n", + "## Installation\n", + "\n", + "The plotting functions require optional dependencies (cartopy, matplotlib, dask).\n", + "Install them with the `plotting` extra:\n", + "\n", + "```bash\n", + "pip install 'mlcast-datasets[plotting]'\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import mlcast_datasets\n", + "from mlcast_datasets import plotting" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Load the dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "cat = mlcast_datasets.open_catalog()\n", + "ds = cat.precipitation.it_dpc_sri_5min.to_dask()\n", + "ds" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "## Domain map\n", + "\n", + "Shows the spatial footprint of the radar coverage." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.plot_domain_map(ds, n_coverage_samples=10);" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Temporal coverage\n", + "\n", + "Year-by-month completeness heatmap and yearly timestep count." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.plot_temporal_coverage(ds);" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "## Spatial coverage\n", + "\n", + "Fraction of valid (non-NaN) observations per grid cell." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.plot_spatial_coverage(ds, n_samples=50);" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "## Precipitation statistics\n", + "\n", + "Three-panel map (mean, max, std) and a log-log histogram of non-zero values." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "fig_maps, fig_hist = plotting.plot_precipitation_stats(ds, n_samples=50)" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "## Sample precipitation event\n", + "\n", + "Multi-panel view of a specific event. Here we show the Emilia-Romagna flood\n", + "of May 2023." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.plot_sample_precipitation(\n", + " ds,\n", + " time_slice=slice(\"2023-05-16\", \"2023-05-17T06:00\"),\n", + ");" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "## Monthly climatology\n", + "\n", + "Box plot of domain-mean precipitation grouped by calendar month." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.plot_monthly_cycle(ds, n_samples=2000);" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## Summary table\n", + "\n", + "Compact overview of key dataset properties." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "plotting.generate_summary_table(ds)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 7b90610..1b070ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ dynamic = ["version"] description = "Intake catalog for datasets relevant for machine learning based nowcasting" authors = [ {name = "Leif Denby", email = "lcd@dmi.dk"}, + {name = "Gabriele Franch", email = "franch@fbk.eu"}, ] readme = "README.md" requires-python = ">=3.11" @@ -21,6 +22,15 @@ license-files = [ "LICENSE-BSD", ] +[project.optional-dependencies] +plotting = [ + "cartopy>=0.25.0", + "dask>=2024.0.0", + "matplotlib>=3.10.0", + "numpy>=1.24", + "pandas>=2.0", +] + [dependency-groups] dev = [ "ipdb>=0.13.13", @@ -62,3 +72,6 @@ mlcast_datasets = ["catalog/*.yaml", "catalog/**/*.yaml"] [tool.setuptools_scm] version_scheme = "post-release" + +[tool.isort] +profile = "black" diff --git a/src/mlcast_datasets/plotting/__init__.py b/src/mlcast_datasets/plotting/__init__.py new file mode 100644 index 0000000..e892d79 --- /dev/null +++ b/src/mlcast_datasets/plotting/__init__.py @@ -0,0 +1,40 @@ +"""Reusable plotting and diagnostic utilities for mlcast radar datasets.""" + +try: + import cartopy # noqa: F401 + import matplotlib # noqa: F401 +except ImportError as exc: + raise ImportError( + "mlcast-datasets plotting extras are not installed. " + "Run: pip install 'mlcast-datasets[plotting]'" + ) from exc + +from ._metadata import ( + get_data_crs, + get_domain_extent, + get_variable_label, + infer_time_step_minutes, + select_plot_variable, +) +from .domain_map import plot_domain_map +from .monthly_cycle import plot_monthly_cycle +from .precipitation_stats import plot_precipitation_stats +from .sample_precipitation import plot_sample_precipitation +from .spatial_coverage import plot_spatial_coverage +from .summary_table import generate_summary_table +from .temporal_coverage import plot_temporal_coverage + +__all__ = [ + "plot_domain_map", + "plot_monthly_cycle", + "plot_precipitation_stats", + "plot_sample_precipitation", + "plot_spatial_coverage", + "plot_temporal_coverage", + "generate_summary_table", + "select_plot_variable", + "get_data_crs", + "get_domain_extent", + "infer_time_step_minutes", + "get_variable_label", +] diff --git a/src/mlcast_datasets/plotting/_map_helpers.py b/src/mlcast_datasets/plotting/_map_helpers.py new file mode 100644 index 0000000..17d5fc3 --- /dev/null +++ b/src/mlcast_datasets/plotting/_map_helpers.py @@ -0,0 +1,127 @@ +"""Shared plotting utilities: map features, colormaps, savefig.""" + +from __future__ import annotations + +from pathlib import Path + +import cartopy.crs as ccrs +import cartopy.feature as cfeature +import matplotlib.pyplot as plt +from matplotlib.colors import BoundaryNorm, ListedColormap +from matplotlib.figure import Figure + +RCPARAMS = { + "font.size": 8, + "axes.labelsize": 8, + "axes.titlesize": 9, + "xtick.labelsize": 7, + "ytick.labelsize": 7, + "legend.fontsize": 7, + "figure.dpi": 150, + "savefig.dpi": 300, + "savefig.bbox": "tight", + "savefig.pad_inches": 0.05, +} + +PLATE_CARREE = ccrs.PlateCarree() + + +def setup_rcparams() -> None: + """Apply publication-quality matplotlib rcParams. + + Updates ``plt.rcParams`` with font sizes, DPI, and save settings + suitable for journal-quality figures. + """ + plt.rcParams.update(RCPARAMS) + + +def add_map_features(ax, resolution: str = "50m"): + """Add coastlines, borders, ocean/land fills, and gridlines to a cartopy axis. + + Parameters + ---------- + ax : cartopy.mpl.geoaxes.GeoAxes + A cartopy-aware matplotlib axes to decorate. + resolution : str, optional + Natural Earth feature resolution, one of ``'10m'``, ``'50m'``, or + ``'110m'``. Default is ``'50m'``. + + Returns + ------- + cartopy.mpl.gridliner.Gridliner + The gridliner artist added to the axes. + """ + ax.add_feature(cfeature.OCEAN.with_scale(resolution), facecolor="#ddeeff", zorder=0) + ax.add_feature(cfeature.LAND.with_scale(resolution), facecolor="#f5f5f0", zorder=0) + ax.coastlines(resolution=resolution, linewidth=0.6, color="0.2") + ax.add_feature( + cfeature.BORDERS.with_scale(resolution), + linewidth=0.4, + linestyle="--", + edgecolor="0.4", + ) + gl = ax.gridlines(draw_labels=True, linewidth=0.3, color="0.7", alpha=0.5) + gl.top_labels = False + gl.right_labels = False + return gl + + +def get_precipitation_cmap( + levels: list[float] | None = None, + colors: list[str] | None = None, +) -> tuple[ListedColormap, BoundaryNorm]: + """Return a discrete precipitation colormap with BoundaryNorm. + + Parameters + ---------- + levels : list of float or None, optional + Boundary values for the discrete colour bins. Default is + ``[0, 0.1, 0.5, 1, 2, 5, 10, 20, 50, 100]``. + colors : list of str or None, optional + Hex colour strings, one fewer than *levels*. Default is a 9-colour + blue-green-orange-red ramp. + + Returns + ------- + cmap : matplotlib.colors.ListedColormap + The discrete colormap. + norm : matplotlib.colors.BoundaryNorm + The norm mapping data values to colormap indices. + """ + if levels is None: + levels = [0, 0.1, 0.5, 1, 2, 5, 10, 20, 50, 100] + if colors is None: + colors = [ + "#f7f7f7", + "#c6dbef", + "#9ecae1", + "#6baed6", + "#3182bd", + "#31a354", + "#f7dc6f", + "#e67e22", + "#e74c3c", + ] + cmap = ListedColormap(colors) + norm = BoundaryNorm(levels, ncolors=len(colors)) + return cmap, norm + + +def savefig(fig: Figure, path: str | Path, close: bool = True) -> None: + """Save a matplotlib figure to disk, creating parent directories as needed. + + Parameters + ---------- + fig : matplotlib.figure.Figure + The figure to save. + path : str or pathlib.Path + Destination file path (e.g. ``'figures/map.pdf'``). + close : bool, optional + If ``True``, close the figure after saving to free memory. + Default is ``True``. + """ + Path(path).parent.mkdir(parents=True, exist_ok=True) + fig.savefig(path) + if close: + plt.close(fig) + print(f"Saved: {path}") diff --git a/src/mlcast_datasets/plotting/_metadata.py b/src/mlcast_datasets/plotting/_metadata.py new file mode 100644 index 0000000..b4c4e33 --- /dev/null +++ b/src/mlcast_datasets/plotting/_metadata.py @@ -0,0 +1,194 @@ +"""Dataset introspection utilities for automatic CRS, variable, and extent detection.""" + +from __future__ import annotations + +import cartopy.crs as ccrs +import numpy as np +import pandas as pd +import xarray as xr + + +def select_plot_variable(ds: xr.Dataset) -> str: + """Find the first 3-D (time, y, x) data variable in the dataset. + + Parameters + ---------- + ds : xr.Dataset + Input dataset to inspect. + + Returns + ------- + str + Name of the first data variable that has a ``time`` dimension and + at least three dimensions. + + Raises + ------ + ValueError + If no suitable variable is found. + """ + for var_name, da in ds.data_vars.items(): + if "time" in da.dims and da.ndim >= 3: + return var_name + raise ValueError("Could not find a 2D spatial variable with a time dimension.") + + +def get_data_crs(ds: xr.Dataset, var_name: str | None = None) -> ccrs.CRS: + """Extract a cartopy CRS from the ``grid_mapping`` attribute of a variable. + + Falls back to ``PlateCarree`` if no ``grid_mapping`` is found. + + Parameters + ---------- + ds : xr.Dataset + Input dataset containing the variable and its grid-mapping variable. + var_name : str or None, optional + Name of the data variable whose ``grid_mapping`` attribute is read. + Auto-detected via :func:`select_plot_variable` when ``None``. + + Returns + ------- + cartopy.crs.CRS + The coordinate reference system parsed from the ``crs_wkt`` + attribute, or ``PlateCarree`` as a fallback. + """ + if var_name is None: + var_name = select_plot_variable(ds) + grid_mapping_name = ds[var_name].attrs.get("grid_mapping") + if grid_mapping_name and grid_mapping_name in ds: + crs_wkt = ds[grid_mapping_name].attrs.get("crs_wkt") + if crs_wkt: + return ccrs.Projection(crs_wkt) + return ccrs.PlateCarree() + + +def get_domain_extent(ds: xr.Dataset, var_name: str | None = None) -> list[float]: + """Compute the geographic bounding box in PlateCarree degrees. + + Uses 2-D ``lat``/``lon`` coordinate variables if present, otherwise + reprojects the 1-D spatial coordinates from the native CRS. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with spatial coordinate variables. + var_name : str or None, optional + Name of the data variable used to determine spatial dimensions and + CRS. Auto-detected via :func:`select_plot_variable` when ``None``. + + Returns + ------- + list of float + Bounding box as ``[lon_min, lon_max, lat_min, lat_max]`` in + decimal degrees (PlateCarree). + + Raises + ------ + ValueError + If the variable does not have exactly two spatial dimensions. + """ + if "lat" in ds and "lon" in ds: + lat = ds["lat"].values + lon = ds["lon"].values + if lat.ndim == 2: + return [ + float(np.nanmin(lon)), + float(np.nanmax(lon)), + float(np.nanmin(lat)), + float(np.nanmax(lat)), + ] + elif lat.ndim == 1: + return [ + float(lon.min()), + float(lon.max()), + float(lat.min()), + float(lat.max()), + ] + + # Fallback: use 1D spatial coordinates with CRS reprojection + if var_name is None: + var_name = select_plot_variable(ds) + da = ds[var_name] + spatial_dims = [d for d in da.dims if d != "time"] + if len(spatial_dims) != 2: + raise ValueError(f"Expected two spatial dims, got {spatial_dims}") + y_dim, x_dim = spatial_dims + + x = ds[x_dim].values + y = ds[y_dim].values + data_crs = get_data_crs(ds, var_name) + plate_carree = ccrs.PlateCarree() + + corners = np.array( + [ + [x.min(), y.min()], + [x.max(), y.min()], + [x.min(), y.max()], + [x.max(), y.max()], + ] + ) + transformed = plate_carree.transform_points(data_crs, corners[:, 0], corners[:, 1]) + return [ + float(transformed[:, 0].min()), + float(transformed[:, 0].max()), + float(transformed[:, 1].min()), + float(transformed[:, 1].max()), + ] + + +def infer_time_step_minutes(ds: xr.Dataset, sample_size: int = 128) -> float: + """Infer the dominant time step in minutes from the first timestamps. + + Computes the median of consecutive time differences over the first + *sample_size* timesteps. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with a ``time`` coordinate. + sample_size : int, optional + Number of leading timesteps to consider. Default is 128. + + Returns + ------- + float + Median time step in minutes, or ``np.nan`` if it cannot be + determined (e.g. fewer than 2 timesteps). + """ + sample = pd.DatetimeIndex(ds["time"].isel(time=slice(0, sample_size)).values) + if len(sample) < 2: + return np.nan + deltas = sample.to_series().diff().dropna() + if deltas.empty: + return np.nan + return deltas.median().total_seconds() / 60.0 + + +def get_variable_label(ds: xr.Dataset, var_name: str | None = None) -> str: + """Build a human-readable axis label from variable attributes. + + Combines ``long_name`` and ``units`` into a string such as + ``'Precipitation rate (mm h⁻¹)'``. + + Parameters + ---------- + ds : xr.Dataset + Input dataset. + var_name : str or None, optional + Name of the variable to label. Auto-detected via + :func:`select_plot_variable` when ``None``. + + Returns + ------- + str + Label string in the form ``'{long_name} ({units})'``, or just + ``'{long_name}'`` if ``units`` is not set. + """ + if var_name is None: + var_name = select_plot_variable(ds) + attrs = ds[var_name].attrs + long_name = attrs.get("long_name", var_name) + units = attrs.get("units", "") + if units: + return f"{long_name} ({units})" + return long_name diff --git a/src/mlcast_datasets/plotting/_stats.py b/src/mlcast_datasets/plotting/_stats.py new file mode 100644 index 0000000..88d3a94 --- /dev/null +++ b/src/mlcast_datasets/plotting/_stats.py @@ -0,0 +1,188 @@ +"""Statistical computation helpers decoupled from plotting.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +from ._metadata import select_plot_variable + + +def _uniform_indices(n_total: int, n_samples: int) -> np.ndarray: + """Return evenly-spaced integer indices spanning an index range. + + Parameters + ---------- + n_total : int + Total number of elements in the range ``[0, n_total - 1]``. + n_samples : int + Desired number of indices. Clamped to *n_total* if larger. + + Returns + ------- + np.ndarray + 1-D integer array of length ``min(n_samples, n_total)``. + """ + return np.linspace(0, n_total - 1, min(n_samples, n_total), dtype=int) + + +def compute_welford_stats( + ds: xr.Dataset, var_name: str | None = None, n_samples: int = 2000 +) -> dict[str, np.ndarray]: + """Compute mean, std, max maps and a value histogram over sampled frames. + + All statistics are computed in a single dask scheduler pass so the + data is read from Zarr only once. + + Parameters + ---------- + ds : xr.Dataset + Input dataset backed by Zarr or in-memory arrays. + var_name : str or None, optional + Name of the data variable to analyse. Auto-detected via + :func:`select_plot_variable` when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 2000. + + Returns + ------- + dict of str to np.ndarray + Dictionary with the following keys: + + - ``'mean'`` : 2-D float array (ny, nx) -- temporal mean. + - ``'std'`` : 2-D float array (ny, nx) -- temporal standard deviation. + - ``'max_val'`` : 2-D float array (ny, nx) -- temporal maximum. + - ``'n_valid'`` : 2-D int64 array (ny, nx) -- count of valid samples. + - ``'hist_counts'`` : 1-D int64 array (199 bins) -- histogram counts. + - ``'hist_bins'`` : 1-D float64 array (200 edges) -- log-spaced bin + edges from 0.01 to 1000. + """ + if var_name is None: + var_name = select_plot_variable(ds) + + N = ds.sizes["time"] + indices = _uniform_indices(N, n_samples) + + sampled = ds[var_name].isel(time=indices) + hist_bins = np.logspace(-2, 3, 200) + + # Build lazy computations for all pixel-level stats + mean_da = sampled.mean(dim="time", skipna=True) + std_da = sampled.std(dim="time", skipna=True, ddof=1) + max_da = sampled.max(dim="time", skipna=True) + n_valid_da = sampled.count(dim="time") + + try: + import dask + import dask.array as da + from dask.diagnostics import ProgressBar as DaskProgressBar + + # NaN / zeros / negatives fall outside the log bins and are ignored. + hist_da, _ = da.histogram(sampled.data.ravel(), bins=hist_bins) + + with DaskProgressBar(dt=0.5): + c_mean, c_std, c_max, c_n, hist_counts = dask.compute( + mean_da, std_da, max_da, n_valid_da, hist_da + ) + except (ImportError, AttributeError): + # Fallback for non-dask-backed datasets + c_mean = mean_da.values + c_std = std_da.values + c_max = max_da.values + c_n = n_valid_da.values + hist_counts, _ = np.histogram(np.asarray(sampled).ravel(), bins=hist_bins) + + mean_arr = np.asarray(c_mean) + std_arr = np.asarray(c_std) + max_arr = np.asarray(c_max) + n_arr = np.asarray(c_n).astype(np.int64) + + return { + "mean": np.where(n_arr > 0, mean_arr, np.nan), + "std": np.where(n_arr > 1, std_arr, np.nan), + "max_val": np.where(n_arr > 0, max_arr, np.nan), + "n_valid": n_arr, + "hist_counts": np.asarray(hist_counts).astype(np.int64), + "hist_bins": hist_bins, + } + + +def compute_spatial_coverage( + ds: xr.Dataset, var_name: str | None = None, n_samples: int = 1000 +) -> xr.DataArray: + """Compute the fraction of valid (non-NaN) observations per pixel. + + Parameters + ---------- + ds : xr.Dataset + Input dataset. + var_name : str or None, optional + Name of the data variable. Auto-detected via + :func:`select_plot_variable` when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 1000. + + Returns + ------- + xr.DataArray + 2-D DataArray with values in ``[0, 1]`` representing the fraction + of sampled timesteps with valid data at each grid cell. + """ + if var_name is None: + var_name = select_plot_variable(ds) + N = ds[var_name].shape[0] + sample_indices = np.linspace(0, N - 1, n_samples, dtype=int) + ds_filtered = ds[var_name].isel(time=sample_indices) + return ds_filtered.map_blocks(lambda b: ~np.isnan(b)).sum(dim="time") / len( + sample_indices + ) + + +def compute_monthly_stats( + ds: xr.Dataset, var_name: str | None = None, n_samples: int = 3000 +) -> dict[int, list[float]]: + """Compute domain-mean values grouped by calendar month. + + Parameters + ---------- + ds : xr.Dataset + Input dataset. + var_name : str or None, optional + Name of the data variable. Auto-detected via + :func:`select_plot_variable` when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 3000. + + Returns + ------- + dict of int to list of float + Mapping from month number (1--12) to a list of spatial-mean values + for the sampled timesteps falling in that month. + """ + if var_name is None: + var_name = select_plot_variable(ds) + + N = ds.sizes["time"] + indices = _uniform_indices(N, n_samples) + + sampled = ds[var_name].isel(time=indices) + spatial_dims = [d for d in sampled.dims if d != "time"] + lazy_means = sampled.mean(dim=spatial_dims, skipna=True) + + try: + from dask.diagnostics import ProgressBar as DaskProgressBar + + with DaskProgressBar(dt=0.5): + spatial_means = lazy_means.values + except ImportError: + spatial_means = lazy_means.values + + months = pd.DatetimeIndex(ds.time.values[indices]).month + valid = np.isfinite(spatial_means) + + monthly_data: dict[int, list[float]] = {m: [] for m in range(1, 13)} + for m in range(1, 13): + monthly_data[m] = spatial_means[(months == m) & valid].tolist() + + return monthly_data diff --git a/src/mlcast_datasets/plotting/domain_map.py b/src/mlcast_datasets/plotting/domain_map.py new file mode 100644 index 0000000..d1c027c --- /dev/null +++ b/src/mlcast_datasets/plotting/domain_map.py @@ -0,0 +1,116 @@ +"""Domain overview map with radar coverage footprint.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import xarray as xr +from matplotlib.figure import Figure + +from ._map_helpers import PLATE_CARREE, add_map_features, savefig, setup_rcparams +from ._metadata import get_data_crs, get_domain_extent, select_plot_variable +from ._stats import compute_spatial_coverage + + +def plot_domain_map( + ds: xr.Dataset, + var_name: str | None = None, + n_coverage_samples: int = 100, + title: str | None = None, + figsize: tuple[float, float] = (6, 8), + cmap: str = "Blues", + alpha: float = 0.4, + map_resolution: str = "10m", + show_admin_boundaries: bool = False, + output_path: str | None = None, +) -> Figure: + """Plot a domain overview map showing the radar coverage footprint. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with 2-D ``lat``/``lon`` coordinate variables. + var_name : str or None, optional + Data variable used to determine coverage. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + n_coverage_samples : int, optional + Number of frames sampled to compute the domain mask. + Default is 100. + title : str or None, optional + Figure title. Auto-generated from dataset attributes when ``None``. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(6, 8)``. + cmap : str, optional + Matplotlib colormap name for the coverage fill. + Default is ``'Blues'``. + alpha : float, optional + Opacity of the coverage fill, in ``[0, 1]``. Default is 0.4. + map_resolution : str, optional + Natural Earth feature resolution (``'10m'``, ``'50m'``, or + ``'110m'``). Default is ``'10m'``. + show_admin_boundaries : bool, optional + If ``True``, overlay sub-national administrative boundaries. + Default is ``False``. + output_path : str or None, optional + If given, save the figure to this file path. + + Returns + ------- + matplotlib.figure.Figure + The generated domain map figure. + """ + setup_rcparams() + if var_name is None: + var_name = select_plot_variable(ds) + data_crs = get_data_crs(ds, var_name) + extent = get_domain_extent(ds, var_name) + + lat = ds["lat"].values + lon = ds["lon"].values + coverage = ( + compute_spatial_coverage(ds, var_name, n_samples=n_coverage_samples).values > 0 + ) + + if title is None: + dataset_title = ds.attrs.get( + "title", ds.attrs.get("mlcast_dataset_identifier", var_name) + ) + title = f"{dataset_title} Spatial Domain" + + fig, ax = plt.subplots(figsize=figsize, subplot_kw={"projection": data_crs}) + + ax.pcolormesh( + lon, + lat, + coverage, + transform=PLATE_CARREE, + cmap=cmap, + vmin=0, + vmax=2, + shading="auto", + zorder=1, + alpha=alpha, + rasterized=True, + ) + + add_map_features(ax, resolution=map_resolution) + + if show_admin_boundaries: + import cartopy.feature as cfeature + + provinces = cfeature.NaturalEarthFeature( + "cultural", + "admin_1_states_provinces_lines", + map_resolution, + facecolor="none", + edgecolor="0.6", + linewidth=0.3, + ) + ax.add_feature(provinces) + + ax.set_extent(extent, crs=PLATE_CARREE) + ax.set_title(title, fontweight="bold") + + if output_path: + savefig(fig, output_path, close=False) + return fig diff --git a/src/mlcast_datasets/plotting/monthly_cycle.py b/src/mlcast_datasets/plotting/monthly_cycle.py new file mode 100644 index 0000000..34eaf2f --- /dev/null +++ b/src/mlcast_datasets/plotting/monthly_cycle.py @@ -0,0 +1,125 @@ +"""Monthly precipitation climatology box plot.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import xarray as xr +from matplotlib.figure import Figure +from matplotlib.patches import Patch + +from ._map_helpers import savefig, setup_rcparams +from ._metadata import get_variable_label, select_plot_variable +from ._stats import compute_monthly_stats + +SEASON_COLORS = ( + ["#2196F3"] * 2 # Jan-Feb (winter) + + ["#4CAF50"] * 3 # Mar-May (spring) + + ["#FF9800"] * 3 # Jun-Aug (summer) + + ["#F44336"] * 3 # Sep-Nov (autumn) + + ["#2196F3"] # Dec (winter) +) + +MONTH_LABELS = [ + "Jan", + "Feb", + "Mar", + "Apr", + "May", + "Jun", + "Jul", + "Aug", + "Sep", + "Oct", + "Nov", + "Dec", +] + + +def plot_monthly_cycle( + ds: xr.Dataset, + var_name: str | None = None, + n_samples: int = 3000, + title: str = "Monthly Precipitation Climatology", + ylabel: str | None = None, + figsize: tuple[float, float] = (7, 4), + show_fliers: bool = False, + season_colors: list[str] | None = None, + output_path: str | None = None, +) -> Figure: + """Plot a boxplot of domain-mean values grouped by calendar month. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with a 3-D ``(time, y, x)`` data variable. + var_name : str or None, optional + Data variable to analyse. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 3000. + title : str, optional + Figure title. Default is ``'Monthly Precipitation Climatology'``. + ylabel : str or None, optional + Y-axis label. Auto-generated from the variable's ``long_name`` + and ``units`` attributes when ``None``. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(7, 4)``. + show_fliers : bool, optional + Whether to show boxplot outlier points. Default is ``False``. + season_colors : list of str or None, optional + List of 12 colour strings (one per month, Jan--Dec). Defaults to + a seasonal colour scheme (blue/green/orange/red). + output_path : str or None, optional + If given, save the figure to this file path. + + Returns + ------- + matplotlib.figure.Figure + The generated monthly-cycle boxplot figure. + """ + setup_rcparams() + if var_name is None: + var_name = select_plot_variable(ds) + if ylabel is None: + ylabel = f"Domain-mean {get_variable_label(ds, var_name)}" + if season_colors is None: + season_colors = SEASON_COLORS + + monthly_data = compute_monthly_stats(ds, var_name, n_samples=n_samples) + + fig, ax = plt.subplots(figsize=figsize) + + data_lists = [monthly_data[m] for m in range(1, 13)] + bp = ax.boxplot( + data_lists, + tick_labels=MONTH_LABELS, + patch_artist=True, + showfliers=show_fliers, + widths=0.6, + whiskerprops=dict(linestyle="--", linewidth=0.8), + medianprops=dict(color="black", linewidth=1.5), + capprops=dict(linewidth=0.8), + ) + + for patch, color in zip(bp["boxes"], season_colors): + patch.set_facecolor(color) + patch.set_alpha(0.6) + patch.set_edgecolor("0.3") + + ax.set_ylabel(ylabel) + ax.set_xlabel("Month") + ax.set_title(title, fontweight="bold") + + legend_elements = [ + Patch(facecolor="#2196F3", alpha=0.6, edgecolor="0.3", label="Winter"), + Patch(facecolor="#4CAF50", alpha=0.6, edgecolor="0.3", label="Spring"), + Patch(facecolor="#FF9800", alpha=0.6, edgecolor="0.3", label="Summer"), + Patch(facecolor="#F44336", alpha=0.6, edgecolor="0.3", label="Autumn"), + ] + ax.legend(handles=legend_elements, loc="upper right", fontsize=7) + ax.grid(axis="y", linewidth=0.3, alpha=0.5) + + if output_path: + savefig(fig, output_path, close=False) + return fig diff --git a/src/mlcast_datasets/plotting/precipitation_stats.py b/src/mlcast_datasets/plotting/precipitation_stats.py new file mode 100644 index 0000000..19d6c65 --- /dev/null +++ b/src/mlcast_datasets/plotting/precipitation_stats.py @@ -0,0 +1,205 @@ +"""Precipitation statistics: mean/max/std maps and value histogram.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from matplotlib.colors import LogNorm +from matplotlib.figure import Figure + +from ._map_helpers import PLATE_CARREE, add_map_features, savefig, setup_rcparams +from ._metadata import get_data_crs, get_domain_extent, select_plot_variable +from ._stats import compute_welford_stats + + +def plot_precipitation_stats( + ds: xr.Dataset, + var_name: str | None = None, + n_samples: int = 2000, + figsize_maps: tuple[float, float] = (15, 5.5), + figsize_hist: tuple[float, float] = (6, 4), + cmap_mean: str = "YlGnBu", + cmap_max: str = "YlOrRd", + cmap_std: str = "OrRd", + max_vmin: float = 1, + max_vmax: float = 300, + hist_color: str = "#3182bd", + hist_xlim: tuple[float, float] = (0.01, 1000), + map_resolution: str = "50m", + output_path_maps: str | None = None, + output_path_hist: str | None = None, +) -> tuple[Figure, Figure]: + """Plot 3-panel stat maps (mean, max, std) and a precipitation histogram. + + Produces two figures: a triptych of spatial maps and a log--log + histogram of non-zero precipitation values. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with 2-D ``lat``/``lon`` coordinate variables. + var_name : str or None, optional + Data variable to analyse. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 2000. + figsize_maps : tuple of float, optional + Figure size ``(width, height)`` for the 3-panel map. + Default is ``(15, 5.5)``. + figsize_hist : tuple of float, optional + Figure size ``(width, height)`` for the histogram. + Default is ``(6, 4)``. + cmap_mean : str, optional + Colormap for the mean panel. Default is ``'YlGnBu'``. + cmap_max : str, optional + Colormap for the maximum panel. Default is ``'YlOrRd'``. + cmap_std : str, optional + Colormap for the standard-deviation panel. Default is ``'OrRd'``. + max_vmin : float, optional + Lower bound of the log-scale range for the max panel. + Default is 1. + max_vmax : float, optional + Upper bound of the log-scale range for the max panel. + Default is 300. + hist_color : str, optional + Bar fill colour for the histogram. Default is ``'#3182bd'``. + hist_xlim : tuple of float, optional + ``(xmin, xmax)`` limits for the histogram x-axis. + Default is ``(0.01, 1000)``. + map_resolution : str, optional + Natural Earth feature resolution (``'10m'``, ``'50m'``, or + ``'110m'``). Default is ``'50m'``. + output_path_maps : str or None, optional + If given, save the map figure to this file path. + output_path_hist : str or None, optional + If given, save the histogram figure to this file path. + + Returns + ------- + fig_maps : matplotlib.figure.Figure + The 3-panel spatial statistics figure. + fig_hist : matplotlib.figure.Figure + The precipitation value histogram figure. + """ + setup_rcparams() + if var_name is None: + var_name = select_plot_variable(ds) + data_crs = get_data_crs(ds, var_name) + extent = get_domain_extent(ds, var_name) + + stats = compute_welford_stats(ds, var_name, n_samples=n_samples) + + lat = ds["lat"].values + lon = ds["lon"].values + + # ===================== 3-panel map ===================== + fig_maps, axes = plt.subplots( + 1, + 3, + figsize=figsize_maps, + subplot_kw={"projection": data_crs}, + constrained_layout=True, + ) + + # (a) Mean + im_a = axes[0].pcolormesh( + lon, + lat, + stats["mean"], + transform=PLATE_CARREE, + cmap=cmap_mean, + vmin=0, + vmax=np.nanpercentile(stats["mean"], 99), + shading="auto", + rasterized=True, + ) + add_map_features(axes[0], resolution=map_resolution) + axes[0].set_extent(extent, crs=PLATE_CARREE) + axes[0].set_title("(a) Mean", fontweight="bold") + fig_maps.colorbar( + im_a, + ax=axes[0], + orientation="horizontal", + shrink=0.8, + pad=0.04, + label="mm h$^{-1}$", + ) + + # (b) Max + im_b = axes[1].pcolormesh( + lon, + lat, + stats["max_val"], + transform=PLATE_CARREE, + cmap=cmap_max, + norm=LogNorm(vmin=max_vmin, vmax=max_vmax), + shading="auto", + rasterized=True, + ) + add_map_features(axes[1], resolution=map_resolution) + axes[1].set_extent(extent, crs=PLATE_CARREE) + axes[1].set_title("(b) Maximum", fontweight="bold") + fig_maps.colorbar( + im_b, + ax=axes[1], + orientation="horizontal", + shrink=0.8, + pad=0.04, + label="mm h$^{-1}$", + ) + + # (c) Std + im_c = axes[2].pcolormesh( + lon, + lat, + stats["std"], + transform=PLATE_CARREE, + cmap=cmap_std, + vmin=0, + vmax=np.nanpercentile(stats["std"], 99), + shading="auto", + rasterized=True, + ) + add_map_features(axes[2], resolution=map_resolution) + axes[2].set_extent(extent, crs=PLATE_CARREE) + axes[2].set_title("(c) Std. dev.", fontweight="bold") + fig_maps.colorbar( + im_c, + ax=axes[2], + orientation="horizontal", + shrink=0.8, + pad=0.04, + label="mm h$^{-1}$", + ) + + if output_path_maps: + savefig(fig_maps, output_path_maps, close=False) + + # ===================== Histogram ===================== + fig_hist, ax = plt.subplots(figsize=figsize_hist) + + bins = stats["hist_bins"] + bin_centers = np.sqrt(bins[:-1] * bins[1:]) + widths = np.diff(bins) + ax.bar( + bin_centers, + stats["hist_counts"], + width=widths, + align="center", + color=hist_color, + edgecolor="none", + alpha=0.8, + ) + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlabel("Precipitation rate (mm h$^{-1}$)") + ax.set_ylabel("Count (from sampled timesteps)") + ax.set_title("Distribution of Non-Zero Precipitation Values", fontweight="bold") + ax.set_xlim(*hist_xlim) + + if output_path_hist: + savefig(fig_hist, output_path_hist, close=False) + + return fig_maps, fig_hist diff --git a/src/mlcast_datasets/plotting/sample_precipitation.py b/src/mlcast_datasets/plotting/sample_precipitation.py new file mode 100644 index 0000000..bf04f5a --- /dev/null +++ b/src/mlcast_datasets/plotting/sample_precipitation.py @@ -0,0 +1,158 @@ +"""Sample precipitation maps from a specific event window.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from matplotlib.figure import Figure + +from ._map_helpers import ( + PLATE_CARREE, + add_map_features, + get_precipitation_cmap, + savefig, + setup_rcparams, +) +from ._metadata import get_data_crs, get_domain_extent, select_plot_variable + + +def plot_sample_precipitation( + ds: xr.Dataset, + time_slice: slice, + var_name: str | None = None, + time_spacing_hours: int = 3, + max_frames: int = 6, + title: str | None = None, + figsize: tuple[float, float] = (12, 9), + levels: list[float] | None = None, + colors: list[str] | None = None, + map_resolution: str = "50m", + output_path: str | None = None, +) -> Figure: + """Plot a grid of precipitation maps from a specific event window. + + Selects frames from the time window at a minimum spacing of + *time_spacing_hours* and renders each on a separate map panel with + a shared discrete precipitation colorbar. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with 2-D ``lat``/``lon`` coordinate variables. + time_slice : slice + Time range to select, e.g. + ``slice('2023-05-16', '2023-05-17T06:00')``. **Required.** + var_name : str or None, optional + Data variable to plot. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + time_spacing_hours : int, optional + Minimum spacing in hours between consecutive displayed frames. + Default is 3. + max_frames : int, optional + Maximum number of map panels to show. Default is 6. + title : str or None, optional + Figure super-title. Auto-generated from the first timestamp + when ``None``. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(12, 9)``. + levels : list of float or None, optional + Discrete colourbar boundaries. Default is + ``[0, 0.1, 0.5, 1, 2, 5, 10, 20, 50, 100]``. + colors : list of str or None, optional + Hex colour strings for the discrete bins (one fewer than + *levels*). Default is a 9-colour ramp. + map_resolution : str, optional + Natural Earth feature resolution (``'10m'``, ``'50m'``, or + ``'110m'``). Default is ``'50m'``. + output_path : str or None, optional + If given, save the figure to this file path. + + Returns + ------- + matplotlib.figure.Figure + The generated multi-panel precipitation figure. + """ + setup_rcparams() + if var_name is None: + var_name = select_plot_variable(ds) + data_crs = get_data_crs(ds, var_name) + extent = get_domain_extent(ds, var_name) + + event_window = ds.sel(time=time_slice) + all_times = event_window.time.values + + # Pick timestamps at least time_spacing_hours apart + picked = [all_times[0]] + for t in all_times[1:]: + if (t - picked[-1]) >= np.timedelta64(time_spacing_hours, "h"): + picked.append(t) + if len(picked) == max_frames: + break + event = event_window.sel(time=picked) + + data = event[var_name].values + timestamps = event.time.values + + lat = ds["lat"].values + lon = ds["lon"].values + + cmap, norm = get_precipitation_cmap(levels=levels, colors=colors) + + if title is None: + t0 = np.datetime_as_string(timestamps[0], unit="D") + title = f"Precipitation Event, {t0}" + + n_frames = len(timestamps) + ncols = min(3, n_frames) + nrows = (n_frames + ncols - 1) // ncols + + fig, axes = plt.subplots( + nrows, + ncols, + figsize=figsize, + subplot_kw={"projection": data_crs}, + constrained_layout=True, + ) + if nrows == 1 and ncols == 1: + axes = np.array([[axes]]) + elif nrows == 1 or ncols == 1: + axes = axes.reshape(nrows, ncols) + + im = None + for idx, ax in enumerate(axes.flat): + if idx < n_frames: + im = ax.pcolormesh( + lon, + lat, + data[idx], + transform=PLATE_CARREE, + cmap=cmap, + norm=norm, + shading="auto", + rasterized=True, + ) + add_map_features(ax, resolution=map_resolution) + ax.set_extent(extent, crs=PLATE_CARREE) + ts = np.datetime_as_string(timestamps[idx], unit="m") + ax.set_title(ts.replace("T", " ") + " UTC", fontsize=9) + else: + ax.set_visible(False) + + if im is not None: + cbar = fig.colorbar( + im, + ax=axes.ravel().tolist(), + orientation="horizontal", + shrink=0.6, + pad=0.03, + aspect=40, + ) + cbar.set_label("Precipitation rate (mm h$^{-1}$)") + + fig.suptitle(title, fontsize=13, fontweight="bold") + + if output_path: + savefig(fig, output_path, close=False) + return fig diff --git a/src/mlcast_datasets/plotting/spatial_coverage.py b/src/mlcast_datasets/plotting/spatial_coverage.py new file mode 100644 index 0000000..19a4871 --- /dev/null +++ b/src/mlcast_datasets/plotting/spatial_coverage.py @@ -0,0 +1,113 @@ +"""Spatial coverage map: fraction of valid data per pixel.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import xarray as xr +from matplotlib.figure import Figure + +from ._map_helpers import PLATE_CARREE, add_map_features, savefig, setup_rcparams +from ._metadata import get_data_crs, get_domain_extent, select_plot_variable +from ._stats import compute_spatial_coverage + + +def plot_spatial_coverage( + ds: xr.Dataset, + var_name: str | None = None, + n_samples: int = 1000, + title: str = "Spatial Data Coverage", + figsize: tuple[float, float] = (6, 8), + cmap: str = "YlOrRd_r", + vmin: float = 0, + vmax: float = 1, + contour_levels: tuple[float, ...] = (0.5, 0.9), + contour_colors: tuple[str, ...] = ("#c0392b", "#2c3e50"), + map_resolution: str = "10m", + output_path: str | None = None, +) -> Figure: + """Plot the fraction of valid observations per pixel as a map. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with a 3-D ``(time, y, x)`` data variable and + 2-D ``lat``/``lon`` coordinate variables. + var_name : str or None, optional + Data variable to analyse. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + n_samples : int, optional + Number of uniformly spaced timesteps to sample. Default is 1000. + title : str, optional + Figure title. Default is ``'Spatial Data Coverage'``. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(6, 8)``. + cmap : str, optional + Matplotlib colormap name. Default is ``'YlOrRd_r'``. + vmin : float, optional + Lower bound of the colourbar. Default is 0. + vmax : float, optional + Upper bound of the colourbar. Default is 1. + contour_levels : tuple of float, optional + Iso-contour levels to overlay on the map. + Default is ``(0.5, 0.9)``. + contour_colors : tuple of str, optional + Colours for the contour lines, matched 1-to-1 with + *contour_levels*. Default is ``('#c0392b', '#2c3e50')``. + map_resolution : str, optional + Natural Earth feature resolution (``'10m'``, ``'50m'``, or + ``'110m'``). Default is ``'10m'``. + output_path : str or None, optional + If given, save the figure to this file path. + + Returns + ------- + matplotlib.figure.Figure + The generated spatial coverage figure. + """ + setup_rcparams() + if var_name is None: + var_name = select_plot_variable(ds) + data_crs = get_data_crs(ds, var_name) + extent = get_domain_extent(ds, var_name) + + frac = compute_spatial_coverage(ds, var_name, n_samples=n_samples).values + + lat = ds["lat"].values + lon = ds["lon"].values + + fig, ax = plt.subplots(figsize=figsize, subplot_kw={"projection": data_crs}) + + im = ax.pcolormesh( + lon, + lat, + frac, + transform=PLATE_CARREE, + cmap=cmap, + vmin=vmin, + vmax=vmax, + shading="auto", + rasterized=True, + ) + + if contour_levels: + ax.contour( + lon, + lat, + frac, + levels=list(contour_levels), + colors=list(contour_colors), + linewidths=[0.8 + 0.2 * i for i in range(len(contour_levels))], + transform=PLATE_CARREE, + ) + + add_map_features(ax, resolution=map_resolution) + ax.set_extent(extent, crs=PLATE_CARREE) + + cbar = fig.colorbar(im, ax=ax, orientation="horizontal", shrink=0.7, pad=0.08) + cbar.set_label("Fraction of valid observations") + ax.set_title(title, fontsize=12, fontweight="bold") + + if output_path: + savefig(fig, output_path, close=False) + return fig diff --git a/src/mlcast_datasets/plotting/summary_table.py b/src/mlcast_datasets/plotting/summary_table.py new file mode 100644 index 0000000..571fe10 --- /dev/null +++ b/src/mlcast_datasets/plotting/summary_table.py @@ -0,0 +1,136 @@ +"""Build a summary DataFrame of dataset metadata.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd +import xarray as xr + +from ._metadata import get_data_crs, infer_time_step_minutes, select_plot_variable + + +def generate_summary_table( + ds: xr.Dataset, + var_name: str | None = None, + compressed_size_bytes: int | None = None, + output_path: str | None = None, +) -> pd.DataFrame: + """Build a summary table of dataset metadata. + + Extracts time range, grid dimensions, CRS, compression statistics, + and other properties into a two-column DataFrame. + + Parameters + ---------- + ds : xr.Dataset + Input dataset to summarise. + var_name : str or None, optional + Data variable to describe. Auto-detected via + :func:`~mlcast_datasets.plotting._metadata.select_plot_variable` + when ``None``. + compressed_size_bytes : int or None, optional + On-disk compressed size in bytes. When provided, compression ratio + and compressed volume are included in the table. + output_path : str or None, optional + If given, save the DataFrame as a CSV file to this path. + + Returns + ------- + pd.DataFrame + Two-column DataFrame with columns ``['Property', 'Value']`` + containing one row per metadata property. + """ + if var_name is None: + var_name = select_plot_variable(ds) + + n_time = ds.sizes["time"] + n_y = ds.sizes["y"] + n_x = ds.sizes["x"] + + t0 = pd.Timestamp(ds.time.values[0]) + t1 = pd.Timestamp(ds.time.values[-1]) + time_range_str = f"{t0.strftime('%Y-%m-%d')} to {t1.strftime('%Y-%m-%d')}" + + if "missing_times" in ds: + n_missing = ds.sizes["missing_times"] + missing_pct = n_missing / (n_missing + n_time) * 100 + missing_str = f"{n_missing:,} ({missing_pct:.1f}%)" + else: + missing_str = "N/A" + + attrs = ds[var_name].attrs + long_name = attrs.get("long_name", var_name) + units = attrs.get("units", "unknown") + + data_crs = get_data_crs(ds, var_name) + crs_name = type(data_crs).__name__ + grid_mapping_name = ds[var_name].attrs.get("grid_mapping", "") + if grid_mapping_name and grid_mapping_name in ds: + gm_attrs = ds[grid_mapping_name].attrs + gm_name = gm_attrs.get("grid_mapping_name", crs_name) + crs_display = gm_name.replace("_", " ").title() + else: + crs_display = crs_name + + uncompressed_bytes = n_time * n_y * n_x * 4 + uncompressed_tb = uncompressed_bytes / 1e12 + + if compressed_size_bytes is not None: + compressed_gb = compressed_size_bytes / 1e9 + compression_ratio = uncompressed_bytes / compressed_size_bytes + compressed_str = f"{compressed_gb:.0f} GB" + ratio_str = f"{compression_ratio:.0f}:1" + else: + compressed_str = "N/A" + ratio_str = "N/A" + + freq_str = ds.attrs.get("base_frequencies", "") + if freq_str: + freq_bands = [] + for part in freq_str.split(";"): + if not part.strip(): + continue + freq_part, range_part = part.strip().split(":", 1) + start, end = range_part.split("/") + end_display = end.strip() if end.strip().lower() != "none" else "present" + freq_bands.append( + f"{freq_part.strip()} ({start.strip()[:10]}-{end_display[:10]})" + ) + freq_display = ", ".join(freq_bands) + else: + step = infer_time_step_minutes(ds) + freq_display = f"{int(step)} min" if np.isfinite(step) else "variable" + + if "x" in ds and len(ds["x"]) > 1: + dx = abs(float(ds["x"].values[1] - ds["x"].values[0])) + res_str = f"{dx / 1000:.0f} km" if dx > 100 else f"{dx:.0f} m" + else: + res_str = "N/A" + + rows = [ + ("Time range", time_range_str), + ("Total timesteps", f"{n_time:,}"), + ("Missing timesteps", missing_str), + ("Grid dimensions", f"{n_y} × {n_x} pixels"), + ("Spatial resolution", res_str), + ("Data variable", f"{var_name} ({long_name})"), + ("Units", units), + ("Data type", "float32"), + ("Uncompressed volume", f"{uncompressed_tb:.1f} TB"), + ("Compressed volume", compressed_str), + ("Compression ratio", ratio_str), + ("Temporal frequency", freq_display), + ("CRS", crs_display), + ("License", ds.attrs.get("license", "unknown")), + ] + + df = pd.DataFrame(rows, columns=["Property", "Value"]) + + if output_path: + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + df.to_csv(output_path, index=False) + print(f"Saved: {output_path}") + + return df diff --git a/src/mlcast_datasets/plotting/temporal_coverage.py b/src/mlcast_datasets/plotting/temporal_coverage.py new file mode 100644 index 0000000..8c3d507 --- /dev/null +++ b/src/mlcast_datasets/plotting/temporal_coverage.py @@ -0,0 +1,229 @@ +"""Temporal data completeness heatmap and yearly bar chart.""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from matplotlib.colors import BoundaryNorm +from matplotlib.figure import Figure +from matplotlib.patches import Patch + +from ._map_helpers import savefig, setup_rcparams +from ._metadata import infer_time_step_minutes + + +def _parse_base_frequencies( + freq_str: str, +) -> list[tuple[int, pd.Timestamp, pd.Timestamp]]: + """Parse the ``base_frequencies`` dataset attribute into structured tuples. + + The attribute string uses the format + ``'15min:2010-01-01/2014-06-01; 10min:2014-06-01/2020-07-01; ...'``. + + Parameters + ---------- + freq_str : str + Semi-colon-separated frequency bands, each as + ``'{freq}min:{start}/{end}'``. An empty string returns an empty + list. + + Returns + ------- + list of tuple(int, pd.Timestamp, pd.Timestamp) + Each tuple contains ``(freq_min, start, end)`` where *freq_min* is + the temporal frequency in minutes. An ``end`` value of ``'None'`` + in the input is mapped to ``pd.Timestamp('2100-01-01')``. + """ + bands = [] + if not freq_str: + return bands + for part in freq_str.split(";"): + if not part.strip(): + continue + fp, rp = part.strip().split(":", 1) + freq_min = int(fp.replace("min", "")) + start_s, end_s = rp.split("/") + start = pd.Timestamp(start_s.strip()) + end = ( + pd.Timestamp(end_s.strip()) + if end_s.strip().lower() != "none" + else pd.Timestamp("2100-01-01") + ) + bands.append((freq_min, start, end)) + return bands + + +def plot_temporal_coverage( + ds: xr.Dataset, + title: str = "Monthly Data Completeness", + figsize: tuple[float, float] = (8, 7.5), + cmap=None, + completeness_levels: list[float] | None = None, + output_path: str | None = None, +) -> Figure: + """Plot a heatmap of monthly data completeness and a yearly timestep bar chart. + + Gracefully handles datasets with or without ``missing_times`` and + ``base_frequencies`` attributes. + + Parameters + ---------- + ds : xr.Dataset + Input dataset with a ``time`` coordinate. Optionally contains a + ``missing_times`` coordinate and a ``base_frequencies`` global + attribute. + title : str, optional + Title for the heatmap panel. + Default is ``'Monthly Data Completeness'``. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(8, 7.5)``. + cmap : matplotlib.colors.Colormap or None, optional + Colormap for the completeness heatmap. Default is + ``plt.cm.RdYlGn``. + completeness_levels : list of float or None, optional + Boundary values for the ``BoundaryNorm`` applied to the heatmap. + Default is ``[0, 50, 70, 80, 90, 95, 98, 100]``. + output_path : str or None, optional + If given, save the figure to this file path. + + Returns + ------- + matplotlib.figure.Figure + The generated two-panel figure (heatmap + bar chart). + """ + setup_rcparams() + + if cmap is None: + cmap = plt.cm.RdYlGn + if completeness_levels is None: + completeness_levels = [0, 50, 70, 80, 90, 95, 98, 100] + + times = pd.DatetimeIndex(ds.time.values) + + if "missing_times" in ds: + missing = pd.DatetimeIndex(ds.missing_times.values) + else: + step_min = infer_time_step_minutes(ds) + if np.isfinite(step_min): + expected = pd.date_range( + times.min(), times.max(), freq=f"{int(step_min)}min" + ) + missing = expected.difference(times) + else: + missing = pd.DatetimeIndex([]) + + bands = _parse_base_frequencies(ds.attrs.get("base_frequencies", "")) + + years = list(range(times.year.min(), times.year.max() + 1)) + months = list(range(1, 13)) + + actual_ym = times.to_series().groupby([times.year, times.month]).count() + missing_ym = ( + missing.to_series().groupby([missing.year, missing.month]).count() + if len(missing) > 0 + else pd.Series(dtype=int) + ) + + completeness = np.full((len(years), 12), np.nan) + actual_counts = np.zeros((len(years), 12), dtype=int) + + for i, yr in enumerate(years): + for j, mo in enumerate(months): + act = actual_ym.get((yr, mo), 0) + mis = missing_ym.get((yr, mo), 0) if len(missing_ym) > 0 else 0 + actual_counts[i, j] = act + total = act + mis + if total > 0: + completeness[i, j] = act / total * 100 + + # --- Figure --- + fig, (ax1, ax2) = plt.subplots( + 2, + 1, + figsize=figsize, + gridspec_kw={"height_ratios": [3, 1.2]}, + constrained_layout=True, + ) + + # Top: heatmap + norm = BoundaryNorm(completeness_levels, ncolors=cmap.N) + im = ax1.imshow( + completeness, cmap=cmap, norm=norm, aspect="auto", interpolation="nearest" + ) + + month_labels = ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"] + ax1.set_xticks(range(12)) + ax1.set_xticklabels(month_labels) + ax1.set_yticks(range(len(years))) + ax1.set_yticklabels(years) + + for i in range(len(years)): + for j in range(12): + val = completeness[i, j] + if np.isfinite(val): + color = "white" if val < 60 else "black" + ax1.text( + j, + i, + f"{val:.0f}", + ha="center", + va="center", + fontsize=6, + color=color, + ) + + cb = fig.colorbar(im, ax=ax1, orientation="vertical", shrink=0.8, pad=0.02) + cb.set_label("Data completeness (%)") + ax1.set_title(title, fontweight="bold") + + # Bottom: bar chart by year + yearly_counts = actual_counts.sum(axis=1) + + freq_color_map = {5: "#FF9800", 10: "#4CAF50", 15: "#2196F3"} + default_color = "#2196F3" + + bar_colors = [] + legend_entries = {} + for yr in years: + ts = pd.Timestamp(f"{yr}-07-01") + color = default_color + freq_label = None + for freq_min, start, end in bands: + if start <= ts < end: + color = freq_color_map.get(freq_min, default_color) + freq_label = f"{freq_min} min" + break + bar_colors.append(color) + if freq_label: + legend_entries[freq_label] = color + + ax2.bar( + range(len(years)), + yearly_counts, + color=bar_colors, + edgecolor="0.3", + linewidth=0.3, + ) + ax2.set_xticks(range(len(years))) + ax2.set_xticklabels(years, rotation=45, ha="right") + ax2.set_ylabel("Timesteps / year") + ax2.set_xlim(-0.5, len(years) - 0.5) + + if legend_entries: + legend_elements = [ + Patch(facecolor=c, edgecolor="0.3", label=lbl) + for lbl, c in legend_entries.items() + ] + ax2.legend( + handles=legend_elements, + loc="upper left", + title="Frequency", + fontsize=7, + title_fontsize=7, + ) + + if output_path: + savefig(fig, output_path, close=False) + return fig diff --git a/src/mlcast_datasets/tests/test_plotting.py b/src/mlcast_datasets/tests/test_plotting.py new file mode 100644 index 0000000..3ab9cf1 --- /dev/null +++ b/src/mlcast_datasets/tests/test_plotting.py @@ -0,0 +1,224 @@ +"""Unit tests for mlcast_datasets.plotting using a synthetic dataset.""" + +import matplotlib +import numpy as np +import pytest +import xarray as xr + +matplotlib.use("Agg") + + +@pytest.fixture +def synthetic_ds(): + """Create a minimal synthetic dataset mimicking an mlcast zarr store.""" + rng = np.random.default_rng(42) + ny, nx, nt = 20, 15, 100 + + time = np.arange( + "2020-01-01", "2020-01-01T08:20", np.timedelta64(5, "m"), dtype="datetime64[ns]" + )[:nt] + y = np.arange(ny, dtype=np.float64) * 1000 + x = np.arange(nx, dtype=np.float64) * 1000 + + # Simple lat/lon grids (approx equirectangular) + lon_2d, lat_2d = np.meshgrid( + np.linspace(10, 12, nx), + np.linspace(50, 52, ny), + ) + + data = rng.exponential(scale=2.0, size=(nt, ny, nx)).astype(np.float32) + # Inject some NaN (missing data) + data[:, :3, :] = np.nan + + # Use UTM 32N via epsg() — its WKT includes BBOX, so ccrs.Projection() round-trip works + import cartopy.crs as ccrs + + crs_wkt = ccrs.epsg(32632).to_wkt() + + ds = xr.Dataset( + { + "RR": xr.DataArray( + data, + dims=["time", "y", "x"], + attrs={ + "grid_mapping": "crs", + "long_name": "Precipitation rate", + "standard_name": "rainfall_flux", + "units": "mm h-1", + }, + ), + "crs": xr.DataArray( + np.nan, + attrs={"crs_wkt": crs_wkt}, + ), + "lat": xr.DataArray(lat_2d, dims=["y", "x"]), + "lon": xr.DataArray(lon_2d, dims=["y", "x"]), + }, + coords={ + "time": time, + "y": y, + "x": x, + }, + attrs={ + "title": "Test Dataset", + "license": "CC-BY-4.0", + "mlcast_dataset_identifier": "TEST-DATASET", + }, + ) + return ds + + +# ---- _metadata tests ---- + + +def test_select_plot_variable(synthetic_ds): + from mlcast_datasets.plotting._metadata import select_plot_variable + + assert select_plot_variable(synthetic_ds) == "RR" + + +def test_get_data_crs(synthetic_ds): + from mlcast_datasets.plotting._metadata import get_data_crs + + crs = get_data_crs(synthetic_ds, "RR") + assert crs is not None + + +def test_get_domain_extent(synthetic_ds): + from mlcast_datasets.plotting._metadata import get_domain_extent + + extent = get_domain_extent(synthetic_ds, "RR") + assert len(extent) == 4 + lon_min, lon_max, lat_min, lat_max = extent + assert lon_min < lon_max + assert lat_min < lat_max + + +def test_infer_time_step_minutes(synthetic_ds): + from mlcast_datasets.plotting._metadata import infer_time_step_minutes + + step = infer_time_step_minutes(synthetic_ds) + assert step == 5.0 + + +def test_get_variable_label(synthetic_ds): + from mlcast_datasets.plotting._metadata import get_variable_label + + label = get_variable_label(synthetic_ds, "RR") + assert "Precipitation" in label + assert "mm" in label + + +# ---- _stats tests ---- + + +def test_compute_spatial_coverage(synthetic_ds): + from mlcast_datasets.plotting._stats import compute_spatial_coverage + + frac = compute_spatial_coverage(synthetic_ds, n_samples=20) + assert frac.shape == (20, 15) + assert frac.max() <= 1.0 + assert frac.min() >= 0.0 + # First 3 rows are always NaN -> should have 0 coverage + assert np.all(frac[:3, :] == 0.0) + + +def test_compute_welford_stats(synthetic_ds): + from mlcast_datasets.plotting._stats import compute_welford_stats + + stats = compute_welford_stats(synthetic_ds, n_samples=20) + assert "mean" in stats + assert "std" in stats + assert "max_val" in stats + assert stats["mean"].shape == (20, 15) + # Mean should be positive where there is data + valid_mean = stats["mean"][~np.isnan(stats["mean"])] + assert np.all(valid_mean >= 0) + + +def test_compute_monthly_stats(synthetic_ds): + from mlcast_datasets.plotting._stats import compute_monthly_stats + + monthly = compute_monthly_stats(synthetic_ds, n_samples=20) + assert 1 in monthly + assert len(monthly[1]) > 0 # January should have data + + +def test_compute_spatial_coverage_n_samples(synthetic_ds): + from mlcast_datasets.plotting._stats import compute_spatial_coverage + + frac_full = compute_spatial_coverage(synthetic_ds, n_samples=20) + frac_fewer = compute_spatial_coverage(synthetic_ds, n_samples=10) + assert frac_full.shape == frac_fewer.shape + + +# ---- Figure function tests (just check they produce a Figure) ---- + + +def test_plot_domain_map(synthetic_ds): + from mlcast_datasets.plotting import plot_domain_map + + fig = plot_domain_map(synthetic_ds, n_coverage_samples=3) + assert fig is not None + matplotlib.pyplot.close(fig) + + +def test_plot_spatial_coverage(synthetic_ds): + from mlcast_datasets.plotting import plot_spatial_coverage + + fig = plot_spatial_coverage(synthetic_ds, n_samples=20) + assert fig is not None + matplotlib.pyplot.close(fig) + + +def test_plot_monthly_cycle(synthetic_ds): + from mlcast_datasets.plotting import plot_monthly_cycle + + fig = plot_monthly_cycle(synthetic_ds, n_samples=20) + assert fig is not None + matplotlib.pyplot.close(fig) + + +def test_plot_precipitation_stats(synthetic_ds): + from mlcast_datasets.plotting import plot_precipitation_stats + + fig_maps, fig_hist = plot_precipitation_stats(synthetic_ds, n_samples=20) + assert fig_maps is not None + assert fig_hist is not None + matplotlib.pyplot.close(fig_maps) + matplotlib.pyplot.close(fig_hist) + + +def test_plot_sample_precipitation(synthetic_ds): + from mlcast_datasets.plotting import plot_sample_precipitation + + fig = plot_sample_precipitation( + synthetic_ds, + time_slice=slice("2020-01-01T00:00", "2020-01-01T08:00"), + time_spacing_hours=1, + max_frames=4, + ) + assert fig is not None + matplotlib.pyplot.close(fig) + + +def test_plot_temporal_coverage(synthetic_ds): + from mlcast_datasets.plotting import plot_temporal_coverage + + fig = plot_temporal_coverage(synthetic_ds) + assert fig is not None + matplotlib.pyplot.close(fig) + + +def test_generate_summary_table(synthetic_ds): + import pandas as pd + + from mlcast_datasets.plotting import generate_summary_table + + df = generate_summary_table(synthetic_ds) + assert isinstance(df, pd.DataFrame) + assert list(df.columns) == ["Property", "Value"] + props = dict(zip(df["Property"], df["Value"])) + assert "RR" in props["Data variable"] + assert props["Data type"] == "float32" + assert "2020" in props["Time range"]