#!/usr/bin/env python
# Copyright (c) 2011-2026, wradlib developers.
# Distributed under the MIT License. See LICENSE.txt for more info.
"""
Composition
^^^^^^^^^^^
Combine data from different radar locations on one common set of locations
.. autosummary::
:nosignatures:
:toctree: generated/
{}
"""
__all__ = [
"extract_circle",
"togrid",
"compose_ko",
"compose_weighted",
"CompMethods",
"transform_binned",
"sweep_to_raster",
]
__doc__ = __doc__.format("\n ".join(__all__))
from functools import singledispatch
import numpy as np
import pyproj
import xarray as xr
from wradlib.georef import (
add_raster_grid_mapping,
get_earth_projection,
get_raster_coordinates,
reproject,
)
from wradlib.ipol import Nearest, RectBin
from wradlib.util import XarrayMethods, docstring, get_apply_ufunc_variables
[docs]
@singledispatch
def togrid(src, trg, radius, center, data, interpol, *args, **kwargs):
"""Interpolate data from a radar location to the composite grid or set of \
locations
Parameters
----------
src : :class:`numpy:numpy.ndarray`
array of float of shape (numpoints, ndim),
cartesian x / y coordinates of the radar bins
trg : :class:`numpy:numpy.ndarray`
array of float of shape (numpoints, ndim),
cartesian x / y coordinates of the composite
radius : float
the radius of the radar circle (same units as src and trg)
center : :class:`numpy:numpy.ndarray`
array of float, the location coordinates of the radar
data : :class:`numpy:numpy.ndarray`
array of float, the data that should be transferred to composite
interpol : :class:`~wradlib.ipol.IpolBase`
an interpolation class name from :mod:`wradlib.ipol`
e.g. :class:`~wradlib.ipol.Nearest` or :class:`~wradlib.ipol.Idw`
Other Parameters
----------------
*args : dict
arguments of Interpolator (see class documentation)
Keyword Arguments
-----------------
**kwargs : dict
keyword arguments of Interpolator (see class documentation)
Returns
-------
output : :class:`numpy:numpy.ndarray`
array of float, data of the radar circle which is interpolated on
the composite grid
Note
----
Keyword arguments to be used while calling the interpolator can be issued as
`call_kwargs`, e.g. togrid(..., call_kwargs=dict(maxdist=10))
Examples
--------
See :ref:`togrid:togrid` and :ref:`notebooks:notebooks/basics/wradlib_workflow:gridding`.
"""
# get indices to select the subgrid from the composite grid
ix = extract_circle(center, radius, trg)
call_kwargs = kwargs.pop("call_kwargs", {})
# interpolate on subgrid
ip = interpol(src, trg[ix], *args, **kwargs)
data_on_subgrid = ip(data, **call_kwargs).reshape(len(ix))
# create container for entire grid
composegridshape = [len(trg)]
composegridshape.extend(data.shape[1:])
compose_grid = np.repeat(np.nan, len(trg) * np.prod(data.shape[1:])).reshape(
composegridshape
)
# push subgrid results into the large grid
compose_grid[ix] = data_on_subgrid
return compose_grid
@togrid.register(xr.DataArray)
@togrid.register(xr.Dataset)
def _togrid_xarray(obj, trg, *args, **kwargs):
dim0 = obj.wrl.util.dim0()
grid_xy = (
xr.concat(xr.broadcast(trg.y, trg.x), "xy")
.stack(npoints_cart=("y", "x"))
.transpose(..., "xy")
)
xy = (
xr.concat([obj.y, obj.x], "xy")
.stack(npoints_pol=(dim0, "range"))
.transpose(..., "xy")
.reset_coords(drop=True)
)
kwargs.setdefault("radius", obj.range.max().values)
kwargs.setdefault("center", (obj.y.mean().values, obj.x.mean().values))
kwargs.setdefault("interpol", Nearest)
# separate variables that will be applied with ufunc
if isinstance(obj, xr.Dataset):
obj_to_grid, keep = get_apply_ufunc_variables(obj, (dim0, "range"))
output_dtypes = [obj_to_grid[list(obj_to_grid.data_vars.keys())[0]].dtype]
else:
obj_to_grid = obj
output_dtypes = [obj_to_grid.dtype]
obj_stacked = obj_to_grid.stack(npoints_pol=(dim0, "range")).reset_coords(drop=True)
def wrapper(obj, xy, grid_xy, **kwargs):
radius = kwargs.get("radius")
center = kwargs.get("center")
ipol = kwargs.get("interpol")
out = togrid(xy, grid_xy, radius, center, obj, ipol)
return out
out = xr.apply_ufunc(
wrapper,
obj_stacked,
xy,
grid_xy,
input_core_dims=[
["npoints_pol"],
["npoints_pol", "xy"],
["npoints_cart", "xy"],
],
output_core_dims=[["npoints_cart"]],
output_dtypes=output_dtypes,
dask="parallelized",
kwargs=kwargs,
dask_gufunc_kwargs=dict(allow_rechunk=True),
on_missing_core_dim="drop",
)
out = out.unstack("npoints_cart")
if isinstance(obj, xr.Dataset):
out = xr.merge([out, keep], compat="no_conflicts")
else:
out.attrs = obj.attrs
out.name = f"{obj.name}.togrid"
return out
[docs]
@singledispatch
def compose_ko(radargrids, qualitygrids):
"""Composes grids according to quality information using quality \
information as a knockout criterion.
The value of the composed pixel is taken from the radargrid whose
quality grid has the highest value.
Parameters
----------
radargrids : list
radar data to be composited. Each item in the list corresponds to the
data of one radar location. All items must have the same shape.
qualitygrids : list
quality data to decide upon which radar site will contribute its pixel
to the composite. Then length of this list must be the same as that
of `radargrids`. All items must have the same shape and be aligned with
the items in `radargrids`.
Returns
-------
composite : :class:`numpy:numpy.ndarray`
"""
# first add a fallback array for all pixels having missing values in all
# radargrids
radarfallback = np.repeat(np.nan, np.prod(radargrids[0].shape)).reshape(
radargrids[0].shape
)
radargrids.append(radarfallback)
radarinfo = np.array(radargrids)
# then do the same for the quality grids
qualityfallback = np.repeat(-np.inf, np.prod(radargrids[0].shape)).reshape(
radargrids[0].shape
)
qualitygrids.append(qualityfallback)
qualityinfo = np.array(qualitygrids)
select = np.nanargmax(qualityinfo, axis=0)
composite = radarinfo.reshape((radarinfo.shape[0], -1))[
select.ravel(), np.arange(np.prod(radarinfo.shape[1:]))
].reshape(radarinfo.shape[1:])
radargrids.pop()
qualitygrids.pop()
return composite
@compose_ko.register(xr.DataArray)
def _compose_ko_xarray(radargrids, qualitygrids):
"""
Xarray knockout composition:
select radar pixel with the highest quality per grid point.
"""
# mask quality where radar data is missing
qualityinfo = qualitygrids.where(radargrids.notnull(), -np.inf)
# index of best radar per pixel
best = qualityinfo.argmax("radar")
# select radar value at that index
composite = radargrids.isel(radar=best)
composite.name = radargrids.name
composite.attrs = radargrids.attrs
return composite
[docs]
@singledispatch
def compose_weighted(radargrids, qualitygrids):
"""Composes grids according to quality information using a weighted \
averaging approach.
The value of the composed pixel is the weighted average of all radar
pixels with the quality values being the weights.
Parameters
----------
radargrids : list
list of arrays
qualitygrids : list
list of arrays
Returns
-------
composite : :class:`numpy:numpy.ndarray`
Examples
--------
See :ref:`togrid:togrid` and :doc:`notebooks:notebooks/workflow/recipe1`.
See Also
--------
:func:`~wradlib.comp.compose_ko`
"""
radarinfo = np.array(radargrids)
qualityinfo = np.array(qualitygrids)
# overall nanmask
nanmask = np.all(np.isnan(radarinfo), axis=0)
# quality grids must contain values only where radarinfo does
qualityinfo[np.isnan(radarinfo)] = np.nan
qualityinfo /= np.nansum(qualityinfo, axis=0)
composite = np.nansum(radarinfo * qualityinfo, axis=0)
composite[nanmask] = np.nan
return composite
@compose_weighted.register(xr.DataArray)
def _compose_weighted_xarray(radargrids, qualitygrids):
qualityinfo = qualitygrids.where(radargrids)
qualityinfo /= qualityinfo.sum("radar", skipna=True)
composite = (radargrids * qualityinfo).sum("radar", skipna=True)
composite.name = radargrids.name
composite.attrs = radargrids.attrs
return composite.where(radargrids.sum("radar"))
[docs]
def sweep_to_raster(sweep, raster, **kwargs):
"""Transform a radar sweep into a raster image.
Parameters
----------
sweep : :class:`xarray:xarray.Dataset` | :class:`xarray:xarray.DataArray`
radar sweep dataset/dataarray following WMO conventions
raster : :class:`xarray:xarray.Dataset`
raster image dataset following CF conventions
Keyword Arguments
-----------------
transform : :class:`wradlib.ipol`
a transformation object (if None, `transform_binned` will be called)
**kwargs : dict
keyword arguments of Interpolator (see class documentation)
Returns
-------
out : :class:`xarray:xarray.Dataset` | :class:`xarray:xarray.DataArray`
raster image with transformed sweep values
Examples
--------
See :ref:`max_reflectivity:sweep-to-raster` and :doc:`notebooks:notebooks/workflow/recipe1`.
"""
dim0 = sweep.wrl.util.dim0()
if kwargs.get("transform", None) is None:
kwargs["transform"] = transform_binned(sweep, raster)
def wrapper(obj, **kwargs):
transform = kwargs.pop("transform")
return transform(obj, **kwargs)
with xr.set_options(keep_attrs=True):
out = xr.apply_ufunc(
wrapper,
sweep,
input_core_dims=[[dim0, "range"]],
output_core_dims=[["y", "x"]],
dask="parallelized",
kwargs=kwargs,
on_missing_core_dim="drop",
)
out = out.assign_coords(x=raster.x, y=raster.y, spatial_ref=raster.spatial_ref)
out = add_raster_grid_mapping(out)
return out
[docs]
class CompMethods(XarrayMethods):
"""wradlib xarray SubAccessor methods for Comp Methods."""
[docs]
@docstring(togrid)
def togrid(self, *args, **kwargs):
if not isinstance(self, CompMethods):
return togrid(self, *args, **kwargs)
else:
return togrid(self._obj, *args, **kwargs)
[docs]
@docstring(sweep_to_raster)
def sweep_to_raster(self, *args, **kwargs):
if not isinstance(self, CompMethods):
return sweep_to_raster(self, *args, **kwargs)
else:
return sweep_to_raster(self._obj, *args, **kwargs)
[docs]
@docstring(compose_weighted)
def compose_weighted(self, *args, **kwargs):
if not isinstance(self, CompMethods):
return compose_weighted(self, *args, **kwargs)
else:
return compose_weighted(self._obj, *args, **kwargs)
[docs]
@docstring(compose_ko)
def compose_ko(self, *args, **kwargs):
if not isinstance(self, CompMethods):
return compose_ko(self, *args, **kwargs)
else:
return compose_ko(self._obj, *args, **kwargs)
if __name__ == "__main__":
print("wradlib: Calling module <comp> as main...")
