"""Binned Wind Climate Module
Extends windkit Binned Wind Climate module to add a function for creating a weibull
distribution out of the binned wind climate using the WAsP method. It contains a much
faster fortran based binning routine that can be used for example in combination with
mesoscale model output, where speed is crucial.
It also contains functions for calculating stability information efficiently using
a binning algorithm which computes the conditional mean stability based on the wind
histogram in a single pass without storing the whole time series of wind speed.
"""
__all__ = [
"bwc_from_tswc",
"calc_temp_scale",
"stability_histogram",
"create_histogram_z0",
"mean_z0",
"baroclinicity_histogram",
"bwc_resample_like",
"bwc_resample_sectors",
"bwc_resample_wsbins_like",
]
import logging
import numpy as np
import pyproj
import xarray as xr
from pywasp._batch import batch_layout
from pywasp._errors import PywaspError
from pywasp.wasp.config import Config
from pywasp._windkit_internal import (
_ALL_VARS_META,
_ANY_3D_SPATIAL,
_BWC_ATTRS,
_BWC_SCHEMA,
_HIS_ATTRS,
_HIS_BARO_ATTRS,
_HIS_STAB_ATTRS,
_MEAN_BARO_ATTRS,
_MEAN_STAB_ATTRS,
_TSWC_SCHEMA,
_TS_STAB_ATTRS,
_bwc_from_counts,
_create_time_attributes,
_update_history,
_update_var_attrs,
validate_input,
)
from windkit.wind_climate.time_series_wind_climate import WS, WD, DIM_TIME
from windkit import create_sector_coords
from windkit import create_wsbin_coords
from windkit import spatial as wksp
from windkit.spatial import (
to_point,
get_crs,
set_crs,
interp_unstructured_like,
are_spatially_equal,
)
from pywasp.core import Rvea0288, Rvea0326 # pylint:disable=no-name-in-module
from pywasp.wasp.time_series_wind_climate import tswc_rotate
logger = logging.getLogger(__name__)
# fortran calls
_CREATE_HIST = Rvea0288.histograms.add_histogram_points
_CALC_TSCALE = Rvea0326.calc_points.tstar_points
_COND_HIST = Rvea0288.histograms.conditional_hist
_MEAN_HIST = Rvea0288.histograms.mean_from_hist
_STD_HIST = Rvea0288.histograms.std_from_hist
_RESAMPLE_HIST = Rvea0288.histograms.histogram_resample_sectors
_RESAMPLE_HIST_WSBINS = Rvea0288.histograms.histogram_resample_wsbins
_RESAMPLE_HIST_WSBINS_AK = Rvea0288.histograms.histogram_resample_wsbins_ak
_DGDZ = Rvea0288.baroclin.compute_dgdz
_DGDZ_TIME = Rvea0288.baroclin.compute_dgdz_t
# constants
RTD = 180 / np.pi
# Amortize Dask graph overhead while capping coalesced input blocks at 64 MiB.
_HISTOGRAM_TARGET_SAMPLES = 24 * 31
_HISTOGRAM_MAX_INPUT_BYTES = 64 * 1024**2
def _block_layout(array, core_dims):
"""Wrap one eager gufunc block in a batch layout."""
batch_dims = tuple(f"batch_{index}" for index in range(array.ndim - len(core_dims)))
data = xr.DataArray(array, dims=(*batch_dims, *core_dims))
return batch_dims, data, batch_layout(data, core_dims=core_dims, required_coords=())
def _histogram_layout(ds, hist, revert_to_original):
"""Return inputs on either their native or point layout."""
if revert_to_original:
return ds, hist
return to_point(ds), None if hist is None else to_point(hist)
def _broadcast_block(array, array_dims, target_dims, target_shape):
"""Broadcast a NumPy gufunc block without expanding the xarray input."""
for axis, dim in enumerate(target_dims):
if dim not in array_dims:
array = np.expand_dims(array, axis)
return np.broadcast_to(array, target_shape)
def _run_histograms(wspd, wdir, values, bin_ws, bin_wd, n_points):
"""Run the compiled histogram once for each value array."""
outputs = []
count_bins_exceeded = None
for value in values:
# hist and ierr are intent(inout), so never pass caller-owned views.
histogram = np.zeros(
(bin_ws.size, bin_wd.size, n_points), dtype=np.float32, order="F"
)
errors = np.zeros(n_points, dtype=np.int32, order="F")
_CREATE_HIST(wspd, wdir, 0, bin_ws, bin_wd, value, histogram, errors)
outputs.append(histogram)
if count_bins_exceeded is None:
count_bins_exceeded = errors
return (*outputs, count_bins_exceeded)
def _histogram_kernel(wspd, wdir, *values, bin_ws, bin_wd, input_batch_dims):
"""Bin one NumPy block supplied by the lazy gufunc path."""
arrays = tuple(
_broadcast_block(array, dims, input_batch_dims[0], wspd.shape)
for array, dims in zip((wspd, wdir, *values), input_batch_dims, strict=True)
)
wspd, wdir, *values = arrays
batch_dims, speed, ctx = _block_layout(wspd, (DIM_TIME,))
dims = speed.dims
wspd = ctx.flatten(speed, dtype=np.float32)
wdir = ctx.flatten(xr.DataArray(wdir, dims=dims), dtype=np.float32)
values = (
tuple(
ctx.flatten(xr.DataArray(value, dims=dims), dtype=np.float32)
for value in values
)
if values
else (np.ones_like(wspd),)
)
output = _run_histograms(wspd, wdir, values, bin_ws, bin_wd, ctx.n_points)
restored = [
ctx.restore(array, core_dims=("wsbin", "sector"))
.transpose(*batch_dims, "wsbin", "sector")
.data
for array in output[:-1]
]
return (*restored, ctx.restore(output[-1], core_dims=()).data)
def _apply_histogram_kernel(inputs, names, bin_ws, bin_wd):
"""Apply the histogram kernel to one complete eager or Dask time chunk."""
return xr.apply_ufunc(
_histogram_kernel,
*inputs,
input_core_dims=[[DIM_TIME]] * len(inputs),
output_core_dims=[["wsbin", "sector"]] * len(names) + [[]],
kwargs={
"bin_ws": bin_ws,
"bin_wd": bin_wd,
"input_batch_dims": [
tuple(dim for dim in array.dims if dim != DIM_TIME) for array in inputs
],
},
dask="parallelized",
vectorize=False,
output_dtypes=[np.float32] * len(names) + [np.int32],
dask_gufunc_kwargs={
"output_sizes": {"wsbin": bin_ws.size, "sector": bin_wd.size}
},
)
def _histogram_time_slices(inputs, time_chunks):
"""Coalesce small time chunks without creating oversized input blocks."""
speed = inputs[0]
batch_size = np.prod(
[
max(speed.chunksizes.get(dim, (speed.sizes[dim],)))
for dim in speed.dims
if dim != DIM_TIME
],
dtype=np.int64,
)
bytes_per_sample = batch_size * sum(
np.dtype(array.dtype).itemsize for array in inputs
)
target = min(
_HISTOGRAM_TARGET_SAMPLES,
max(1, _HISTOGRAM_MAX_INPUT_BYTES // max(1, bytes_per_sample)),
)
slices = []
start = stop = 0
for chunk_size in time_chunks:
stop += chunk_size
if stop - start >= target:
slices.append(slice(start, stop))
start = stop
if start < stop:
slices.append(slice(start, stop))
return slices
def _slice_time_block(inputs, indexer):
"""Return one block with a single Dask chunk along its time core dimension."""
block = tuple(array.isel({DIM_TIME: indexer}) for array in inputs)
return tuple(
array.chunk({DIM_TIME: -1})
if array.chunks is not None and len(array.chunksizes[DIM_TIME]) > 1
else array
for array in block
)
def _histogram_sums(
ds, values=None, hist=None, wsbin_width=1, n_wsbins=30, n_sectors=12
):
"""Create eager or lazy weighted histograms without a zero template.
Parameters
----------
ds : xr.Dataset
Dataset containing wind speed and wind direction.
values : dict of str to xr.DataArray, optional
Values to sum by wind bin. If omitted, wv_count is produced.
hist : xr.Dataset, optional
Existing sums to add to the newly created histograms.
wsbin_width : float
Width of the wind-speed bins.
n_wsbins : int
Number of wind-speed bins.
n_sectors : int
Number of direction sectors.
Returns
-------
xr.Dataset
Histogram sums and count_bins_exceeded. NumPy input produces NumPy
output; Dask input remains lazy and spatially chunked.
"""
values = {} if values is None else values
names = tuple(values) if values else ("wv_count",)
# Compatibility for xarray 2024.9.0; remove when that version is unsupported.
crs_attrs = ds["crs"].attrs.copy() if "crs" in ds.coords else {}
if hist is not None:
# are_spatially_equal only reads the west_east/south_north coordinates,
# never a data variable, and Dataset.chunk() never chunks dimension
# coordinates -- so this stays cheap and adds no Dask tasks even when
# ds/hist's own data variables are lazy. Run it unconditionally.
if not are_spatially_equal(ds, hist):
raise PywaspError(
"Existing histogram and time series must have the same spatial structure"
)
if hist is None:
ws_bins = create_wsbin_coords(n_wsbins, width=wsbin_width)
wd_bins = create_sector_coords(n_sectors)
else:
ws_bins = hist["wsbin"]
wd_bins = hist["sector"]
bin_ws = np.asarray(ws_bins, dtype=np.float32, order="F")
bin_wd = np.asarray(wd_bins, dtype=np.float32, order="F")
# Let apply_ufunc align named dimensions; the kernel expands missing axes
# inside each NumPy block to avoid expensive dataset-wide broadcasting.
inputs = (ds[WS], ds[WD], *(values[name] for name in names if values))
if any(array.chunks is not None for array in inputs):
inputs = xr.unify_chunks(*inputs)
time_chunks = next(
(array.chunksizes[DIM_TIME] for array in inputs if array.chunks is not None), ()
)
if len(time_chunks) <= 1:
output = _apply_histogram_kernel(inputs, names, bin_ws, bin_wd)
else:
time_slices = _histogram_time_slices(inputs, time_chunks)
partials = [
_apply_histogram_kernel(
_slice_time_block(inputs, indexer),
names,
bin_ws,
bin_wd,
)
for indexer in time_slices
]
if len(partials) == 1:
output = partials[0]
else:
output = tuple(
xr.concat(arrays, dim="_histogram_time_chunk").sum(
"_histogram_time_chunk", dtype=arrays[0].dtype
)
for arrays in zip(*partials, strict=True)
)
result = xr.Dataset(attrs=hist.attrs if hist is not None else ds.attrs)
for name, array in zip(names, output[:-1], strict=True):
result[name] = array.transpose("wsbin", "sector", ...)
result["count_bins_exceeded"] = output[-1]
result = result.assign_coords({**ws_bins.coords, **wd_bins.coords})
# Compatibility for xarray 2024.9.0; remove when that version is unsupported.
if crs_attrs and "crs" in result.coords:
result["crs"].attrs = crs_attrs
if hist is not None:
required = {*names, "count_bins_exceeded"}
missing = required - set(hist.data_vars)
if missing:
raise PywaspError(
"Existing histogram is missing variables: " + ", ".join(sorted(missing))
)
for name in (*names, "count_bins_exceeded"):
try:
current, previous = xr.align(
result[name], hist[name], join="exact", copy=False
)
except ValueError as error:
raise PywaspError(
"Existing histogram coordinates do not match the time series"
) from error
result[name] = current + previous
return result
def _update_time_attributes_without_computing(result, ds, hist=None):
"""Update time metadata without executing a Dask-backed time series.
WindKit's metadata helper counts complete records with a dropna over the
data, which is exact for eager data but executes every Dask block.
count, count_missing and recovery_percentage need those real data values
to know which samples are missing; start_time, end_time and
count_expected only need the time index. Once any chunk in an
accumulation has been processed lazily, the true count for that chunk is
permanently unknown, so count/count_missing/recovery_percentage are
omitted from that point on -- for this call and every later one
accumulating on top of it -- rather than silently reporting a number
that is wrong, such as 100% recovery for data that in fact recovered 0%.
"""
all_eager = all(array.chunks is None for array in ds.data_vars.values())
hist_omitted_count = hist is not None and "count" not in hist.attrs
if all_eager and not hist_omitted_count:
attrs = _create_time_attributes(ds, hist=hist)
else:
dummy = np.broadcast_to(np.float32(0.0), (ds.sizes[DIM_TIME],))
proxy = xr.Dataset(
{WS: (DIM_TIME, dummy), WD: (DIM_TIME, dummy)},
coords={DIM_TIME: ds[DIM_TIME]},
)
if hist_omitted_count:
hist = xr.Dataset(attrs={**hist.attrs, "count": 0})
attrs = _create_time_attributes(proxy, hist=hist)
for key in ("count", "count_missing", "recovery_percentage"):
attrs.pop(key, None)
for key in ("count", "count_missing", "recovery_percentage"):
if key not in attrs:
result.attrs.pop(key, None)
result.attrs.update(attrs)
return result
def _conditional_histogram_kernel(hist, freq, percentile):
"""Apply the conditional-histogram core routine to one spatial block."""
batch_dims, histogram_da, ctx = _block_layout(hist, ("wsbin", "sector"))
dims = histogram_da.dims
histogram = ctx.flatten(histogram_da, dtype=np.float32)
frequency = ctx.flatten(
xr.DataArray(freq, dims=dims),
dtype=np.float32,
)
conditional, sector_frequency = _COND_HIST(histogram, frequency, percentile)
conditional = ctx.restore(conditional, core_dims=("sector",)).transpose(
*batch_dims, "sector"
)
sector_frequency = ctx.restore(sector_frequency, core_dims=("sector",)).transpose(
*batch_dims, "sector"
)
return conditional.data, sector_frequency.data
def _conditional_histogram(hist, freq, percentile):
"""Return conditional sums and frequencies eagerly or lazily."""
return xr.apply_ufunc(
_conditional_histogram_kernel,
hist,
freq,
input_core_dims=[["wsbin", "sector"], ["wsbin", "sector"]],
output_core_dims=[["sector"], ["sector"]],
kwargs={"percentile": percentile},
dask="parallelized",
vectorize=False,
output_dtypes=[np.float32, np.float32],
dask_gufunc_kwargs={"allow_rechunk": True},
)
def _sector_statistic_kernel(*arrays, statistic_name):
"""Apply a sector-wise WAsP statistic to one spatial block."""
batch_dims, first, ctx = _block_layout(arrays[0], ("sector",))
dims = first.dims
input_arrays = [xr.DataArray(array, dims=dims) for array in arrays]
fortran_arrays = [ctx.flatten(array, dtype=np.float32) for array in input_arrays]
if statistic_name == "mean":
result = _MEAN_HIST(*fortran_arrays)
elif statistic_name == "std":
result = _STD_HIST(*fortran_arrays)
else:
raise PywaspError(f"Unknown sector statistic: {statistic_name}")
return (
ctx.restore(result, core_dims=("sector",)).transpose(*batch_dims, "sector").data
)
def _sector_statistic(statistic_name, *arrays):
"""Return a sector statistic eagerly or lazily."""
return xr.apply_ufunc(
_sector_statistic_kernel,
*arrays,
input_core_dims=[["sector"]] * len(arrays),
output_core_dims=[["sector"]],
kwargs={"statistic_name": statistic_name},
dask="parallelized",
vectorize=False,
output_dtypes=[np.float32],
dask_gufunc_kwargs={"allow_rechunk": True},
)
[docs]
@validate_input(_TSWC_SCHEMA, run_extra_checks=False)
def bwc_from_tswc(
ds,
hist=None,
normalize=True,
revert_to_original=True,
wsbin_width=1,
n_wsbins=30,
n_sectors=12,
wind_dir_crs=None,
):
"""Add timeseries to histogram
Converts a time series with wind speed & direction to derive
summed histogram of the wind vector.
.. warning::
This function is experimental and its signature may change.
Parameters
----------
ds: xr.Dataset
PyWAsP Timeseries dataset. Can have any of the pywasp
spatial data structures
hist: xr.Dataset, optional
Histogram with dimensions point, wsbin, sector containing counts to
add to this time series. The input is not modified. Default is None,
creating a histogram using wsbin_width, n_wsbins, and n_sectors.
normalize: boolean
Normalize histogram for each sector, storing the frequency per sector in "wdfreq"
variable, Default True
revert_to_original: boolean
If True (default), operate on and return ds's own native spatial
structure directly -- there is no promotion/reversion round trip to
pay for any more, unlike before. If False, promote to point
structure and return that; still useful when composing with other
point-structured accumulators, e.g. across chunked calls.
wsbin_width: float
width of wind speed bins
n_wsbins: int
Number of wind speed bins
n_sectors: int
Number of sectors (wind direction bins)
wind_dir_crs : CRS-like, optional
CRS that defines the output wind-direction frame. If ``ds`` has a
``wind_dir_crs`` attribute, input directions are treated as already
expressed in that CRS frame. Otherwise, input directions are assumed
true-north-relative. Directions are rotated into a projected target
CRS's grid-north frame before binning; a geographic target CRS retains
them in the true-north frame.
Returns
-------
hist: xr.Dataset
Histogram with dimensions point, wsbin, sector containing the values
from the time series. The per-location count_bins_exceeded diagnostic
is retained for both normalized and unnormalized output.
Notes
-----
NumPy-backed inputs are evaluated eagerly. Dask-backed inputs remain lazy
and histogram each time chunk independently before a tree reduction,
preserving the input spatial chunks. Dask is therefore optional.
For lazy output, inspect count_bins_exceeded after computation instead of
relying on the eager bin-overflow warning.
The result's ``count``, ``count_missing`` and ``recovery_percentage``
attributes need the actual data values (a missing-sample count) and are
therefore only set for eager input; ``start_time``, ``end_time`` and
``count_expected`` are always set. For lazy input, or once any chunk in
an accumulation has been lazy, those three attributes are omitted rather
than reported as a number that could be wrong.
"""
existing_wind_dir_crs = ds.attrs.get("wind_dir_crs")
if existing_wind_dir_crs is not None:
existing_wind_dir_crs = pyproj.CRS.from_user_input(existing_wind_dir_crs)
target_wind_dir_crs = existing_wind_dir_crs
if wind_dir_crs is not None:
target_wind_dir_crs = pyproj.CRS.from_user_input(wind_dir_crs)
ds = wksp.reproject(ds, target_wind_dir_crs)
rotation = None
if wind_dir_crs is not None and existing_wind_dir_crs is None:
rotation = wksp.grid_convergence(ds)
elif (
wind_dir_crs is not None
and existing_wind_dir_crs is not None
and not existing_wind_dir_crs.equals(target_wind_dir_crs)
):
rotation = wksp.grid_convergence(ds, source=existing_wind_dir_crs)
if rotation is not None:
ds = tswc_rotate(ds, -rotation)
logger.debug(
f"Convert time series ({ds[WS].sizes[DIM_TIME]} samples) into histogram"
)
previous_hist = hist
anchor_ds, hist = _histogram_layout(ds, hist, revert_to_original)
hist = _histogram_sums(
anchor_ds,
hist=hist,
wsbin_width=wsbin_width,
n_wsbins=n_wsbins,
n_sectors=n_sectors,
)
hist.attrs.update(ds.attrs)
hist = _update_time_attributes_without_computing(hist, ds, previous_hist)
errors = hist["count_bins_exceeded"]
if errors.chunks is None:
n_failed_points = np.count_nonzero(errors.data > 0)
if n_failed_points > 0:
logger.warning(
"The wind speeds exceeded the binning range "
f"({float(hist.wsbin.max())}) in {n_failed_points} points"
)
if normalize:
hist = _normalize_hist(hist)
attrs = (
{
**_BWC_ATTRS,
"count_bins_exceeded": _HIS_ATTRS["count_bins_exceeded"],
}
if normalize
else _HIS_ATTRS
)
hist = _update_var_attrs(hist, attrs)
if target_wind_dir_crs is not None:
hist.attrs["wind_dir_crs"] = target_wind_dir_crs.to_wkt()
return _update_history(hist)
def _accumulate_wv_count(
ds,
hist,
wv_count,
hist_was_supplied,
histogram_name,
wsbin_width,
validate_bins=False,
):
"""Add the current time series to counts used for finalization."""
if wv_count is not None:
return bwc_from_tswc(
ds, hist=wv_count, normalize=False, revert_to_original=True
)
if hist_was_supplied:
raise PywaspError(
f"Finalizing a supplied {histogram_name} histogram requires the "
"matching accumulated wind-vector counts. Build them alongside "
"the histogram with bwc_from_tswc(..., normalize=False) for every "
"chunk already added, on the same spatial structure as hist, and "
"pass them as wv_count. Counts from only the final chunk would "
"silently inflate the means."
)
wsbin = hist["wsbin"].values
if validate_bins:
if wsbin.size < 2:
raise PywaspError(
f"{histogram_name.capitalize()} histograms require at least "
"two wind-speed bins"
)
wsbin_width = float(wsbin[1] - wsbin[0])
if wsbin_width <= 0.0 or not np.allclose(np.diff(wsbin), wsbin_width):
raise PywaspError(
f"{histogram_name.capitalize()} histogram wind-speed bins "
"must be uniformly spaced"
)
return bwc_from_tswc(
ds,
normalize=False,
revert_to_original=True,
wsbin_width=wsbin_width,
n_wsbins=hist.sizes["wsbin"],
n_sectors=hist.sizes["sector"],
)
[docs]
def calc_temp_scale(ds, revert_to_original=True, enhmix=1.0, use_humidity=False):
"""Calculate temperature scale
This calculates a surface layer temperature scale that can
be used to model the wind profile. Using the enhmix can be used
to control mixing over land, because ERA5 has a large bias
in $u_*$ over land. If you supply enhmix you will have to
add a landmask because it is only applied over land. If use_humidity
is true the humidity is used for calculating the air density.
.. warning::
This function is experimental and its signature may change.
Parameters
----------
ds : xarray.Dataset
Dataset containing variables
["PSFC", "T2", "UST", "HFX", "LH", "PBLH"]
Can have any pywasp spatial structure
revert_to_original: boolean
Return the histogram input in original format from the
input :class:`xarray.Dataset`? If false, keeps the data in stacked point
format which is more efficient when doing large
calculations for numerical wind atlases.
enhmix : float
Decrease mixing over land with this factor. If different from
1.0, requires a land mask variable "LANDMASK" since it is only
applied over land.
use_humidity: bool
If True, use the humdity for calculating the air density. If
true, requires variable "elevation" (for the terrain elevation)
and "Q2" to be present on the dataset. By default False
Returns
-------
hist : xr.DataArray
Height-independent ``temp_scale`` DataArray with the input's remaining
spatial and time dimensions.
"""
source = ds if revert_to_original else to_point(ds)
if "UST" not in source.data_vars:
raise KeyError("Required variable 'UST' not found in dataset!")
ctx = batch_layout(source, anchor="UST", core_dims=(DIM_TIME,), required_coords=())
def _over_time(name):
return ctx.flatten(source[name], dtype=np.float32)
def _per_location(name):
return ctx.flatten(source[name], core_dims=(), dtype=np.float32)
def _zeros(*shape):
return np.zeros(shape, dtype=np.float32, order="F")
try:
logger.debug("Calculating temperature scale")
ust = _over_time("UST")
# The land mask only matters when mixing is reduced over land, and the
# humidity fields only when use_humidity is set. Both stay zero-filled
# otherwise, even when the dataset happens to carry them.
land = _per_location("LANDMASK") if enhmix != 1.0 else _zeros(ctx.n_points)
if use_humidity:
elevation = _per_location("elevation")
q2 = _over_time("Q2")
else:
logger.debug("Using dummy variables Q2 and elevation")
elevation = _zeros(ctx.n_points)
q2 = _zeros(*ust.shape)
tstar = _CALC_TSCALE(
_over_time("PSFC"),
_over_time("T2"),
ust,
_over_time("HFX"),
_over_time("LH"),
_over_time("PBLH"),
elevation,
q2,
land,
enhmix,
int(use_humidity),
)
except KeyError as e:
raise KeyError(f"Required variable {e} not found in dataset!")
# The surface fields are height-independent by construction, so a height
# batch dimension is reduced away rather than carried into the result. The
# reduction is checked: restore proves the values really are constant.
drop_height = ("height",) if "height" in ctx.batch_dims else ()
res = ctx.restore(
tstar,
core_dims=(DIM_TIME,),
drop_batch_dims=drop_height,
name="temp_scale",
).to_dataset()
res = _update_var_attrs(res, _TS_STAB_ATTRS)
return _update_history(res["temp_scale"])
[docs]
def stability_histogram(
ds,
hist=None,
finalize=True,
revert_to_original=True,
wsbin_width=1,
n_wsbins=40,
n_sectors=12,
percentile=0.5,
wv_count=None,
landmask=None,
):
"""Add timeseries to existing histogram
Uses the stability parameters "temp_scale" and "ustar_over_pblh".
See the paper "Using observed and modelled heat fluxes for
improved extrapolation of wind distributions" for the definitions
of these variables. The temp_scale variable is defined as $T_*$
in that paper, whereas the "ustar_over_pblh" should be defined as
$u_*/pblh$ as input for this function. This is done to give priority
to the smallest pblh that are most important for the wind profile.
The output is reversed ($**-1$) to give the input format as required
by WAsP ($pblh/u_*$).
The time dimension is copied so that one can see what was the
period that was used to generate the histogram.
This method is optimized for large datasets, where storing
all values in memory is not possible. It uses a binning so that
the stability parameters are calculated as a conditional mean (on the
wind speed and direction) in a single pass through the data.
The histogram is created if not provided, but it can also be
provided to add more data to an existing histogram. The idea
is than that you sum to the existing histogram and then
finalize at the end using finalize=True. When using finalize=True,
the histogram is converted into a mean and standard deviation
of the stability parameters per wind speed and direction bin. To
calculate both sea and land stability, you will have to provide
a landmask so that the finalize step can separate between sea
and land points.
.. warning::
This function is experimental and its signature may change.
Parameters
----------
ds : xarray.Dataset
Dataset containing variables ['wind_speed', 'wind_direction',
'temp_scale', 'ustar_over_pblh']
Can have any pywasp spatial structure
hist : xarray.Dataset
Histogram with dimensions point, wsbin, sector
finalize : boolean
Convert the stabilility histogram into a mean and
standard deviation?
revert_to_original : boolean
If True (default), operate on and return ds's own native spatial
structure directly; there is no promotion/reversion round trip. If
False, promote to point structure and return that; still useful when
composing with other point-structured accumulators, e.g. across
chunked calls.
wsbin_width : float
width of wind speed bins
n_wsbins : int
Number of wind speed bins
n_sectors : int
Number of sectors (wind direction bins)
percentile : float
Percentile of the wind speed distribution to use for
calculating the stability parameters. Default is 0.5,
which is the median.
wv_count : xarray.Dataset, optional
Histogram with dimensions point, wsbin, sector
containg a count. If provided, this histogram is used
for the finalization step. If None, a new histogram
is created from the input ``ds``. Default is None.
landmask: xr.DataArray
Landmask with values 0 (sea) and 1 (land)
Returns
-------
hist : xarray.Dataset
Histogram with dimensions point, wsbin, sector
with the values from ds-timeseries added.
Raises
------
ValueError
If landmask is provided, but does not have the same
spatial structure as ds.
PywaspError
If hist is provided, but does not have the same
spatial structure as ds; or if a supplied hist is finalized
without the matching accumulated wv_count.
ValueError
If landmask is provided, finalize=True and
there are only land or sea points in ds.
"""
if landmask is not None:
if not are_spatially_equal(ds, landmask):
raise ValueError(
"Input ds and landmask must have the same spatial structure"
)
# A supplied histogram may already hold accumulated sums from earlier chunks,
# so counts covering only the current chunk must never be used to finalize it.
hist_is_accumulated = hist is not None
ds_sub = ds[[WS, WD, "temp_scale", "ustar_over_pblh"]]
dss, previous_hist = _histogram_layout(ds_sub, hist, revert_to_original)
hist = _histogram_sums(
dss,
values={
"sum_temp_scale": dss["temp_scale"],
"sum_squared_temp_scale": dss["temp_scale"] ** 2,
"sum_pblh": dss["ustar_over_pblh"],
},
hist=previous_hist,
wsbin_width=wsbin_width,
n_wsbins=n_wsbins,
n_sectors=n_sectors,
)
hist = _update_var_attrs(hist, _HIS_STAB_ATTRS)
hist = _update_time_attributes_without_computing(hist, dss, previous_hist)
if finalize:
wv_count = _accumulate_wv_count(
dss,
hist,
wv_count,
hist_is_accumulated,
"stability",
wsbin_width,
)
res = _finalize_meso(
hist, wv_count["wv_count"], landmask, percentile=percentile
)
return _update_history(res)
return _update_history(hist)
def _stab_mean_std(hist, wv_count, percentile=0.5):
"""Finalize stability sums as sector means and standard deviations."""
counts = wv_count
cond_hist, freq = _conditional_histogram(hist["sum_pblh"], counts, percentile)
pblhmean = _sector_statistic("mean", cond_hist, freq)
cond_hist, freq = _conditional_histogram(hist["sum_temp_scale"], counts, percentile)
cond_hist2, freq = _conditional_histogram(
hist["sum_squared_temp_scale"], counts, percentile
)
tstarstd = _sector_statistic("std", cond_hist, cond_hist2, freq)
tstarmean = _sector_statistic("mean", cond_hist, freq)
result = xr.Dataset(
{
"mean_pblh_scale": 1 / pblhmean,
"rms_temp_scale": tstarstd,
"mean_temp_scale": tstarmean,
},
attrs=hist.attrs,
)
# The conditional-histogram gufunc does not always carry the crs coordinate
# through (older xarray/pyproj drop its attributes); re-set it from the input
# so the result's crs coordinate is identical across dependency versions.
if "crs" in hist.coords:
result = set_crs(result, get_crs(hist))
return result
def create_histogram_z0(
ds,
hist=None,
finalize=True,
revert_to_original=True,
wsbin_width=1,
n_wsbins=40,
n_sectors=12,
percentile=0.5,
wv_count=None,
landmask=None,
):
"""Add timeseries to existing histogram
Uses the stability parameters to derive
summed histogram of the surface layer temperature scale,
squared surface layer temperature scale, boundary
layer height and logarithm of roughness length. The time
dimension is copied so that one can see what was the
period that was used to generate the histogram.
.. warning::
This function is experimental and its signature may change.
Parameters
----------
ds: xr.Dataset
Dataset containing variables
['wind_speed', 'wind_direction', 'LANDMASK', 'ZNT']
Can have any pywasp spatial structure
hist: xr.Dataset
Histogram with dimensions point, wsbin, sector
finalize: boolean
Convert the stabilility histogram into a mean and
standard deviation?
revert_to_original: boolean
If True (default), operate on and return ds's own native spatial
structure directly; there is no promotion/reversion round trip. If
False, promote to point structure and return that; still useful when
composing with other point-structured accumulators, e.g. across
chunked calls.
wsbin_width: float
width of wind speed bins
n_wsbins: int
Number of wind speed bins
n_sectors: int
Number of sectors (wind direction bins)
kwargs:
Other keyword arguments passed on to _finalize_meso
Returns
-------
hist: xr.Dataset
Histogram with dimensions point, wsbin, sector
with the values from ds-timeseries added.
Raises
------
PywaspError
If a supplied hist is finalized without the matching accumulated
wv_count.
"""
# A supplied histogram may already hold accumulated sums from earlier chunks,
# so counts covering only the current chunk must never be used to finalize it.
hist_is_accumulated = hist is not None
ds_sub = ds[[WS, WD, "ZNT"]]
dss, previous_hist = _histogram_layout(ds_sub, hist, revert_to_original)
hist = _histogram_sums(
dss,
values={"z0sum": np.log(dss["ZNT"])},
hist=previous_hist,
wsbin_width=wsbin_width,
n_wsbins=n_wsbins,
n_sectors=n_sectors,
)
hist = _update_var_attrs(
hist,
{
"z0sum": _ALL_VARS_META["z0sum"],
"count_bins_exceeded": _ALL_VARS_META["count_bins_exceeded"],
},
)
hist = _update_time_attributes_without_computing(hist, dss, previous_hist)
hist = hist.drop_vars("point", errors="ignore")
if finalize:
wv_count = _accumulate_wv_count(
dss, hist, wv_count, hist_is_accumulated, "z0", wsbin_width
)
res = _z0_mean(hist, wv_count["wv_count"], landmask, percentile=percentile)
return _update_history(res)
return _update_history(hist)
def _z0_mean(hist, wv_count, landmask, percentile=0.5):
"""Finalize roughness sums as sector means."""
hist["wv_count"] = wv_count
cond_hist, wd_count = _conditional_histogram(
hist["z0sum"], hist["wv_count"], percentile
)
hist["logz0sec"] = _sector_statistic("mean", cond_hist, wd_count)
hist["wd_count"] = wd_count
if landmask is not None:
hist["landmask"] = landmask
# As in _finalize_baro, DataArray.where keeps mean_z0's own axis order
# rather than adopting the landmask's.
hist["mean_z0"] = mean_z0(hist).where(hist["landmask"] != 0, 0.0)
else:
hist["mean_z0"] = mean_z0(hist)
return hist
def mean_z0(hist):
"""Compute mean roughness
Compute mean roughness from log(z0) per wsbin and
sector
.. warning::
This function is experimental and its signature may change.
Parameters
----------
hist: xr.Dataset
Dataset with variables `logz0sec` and `wd_count`
and dimensions `wsbin`, `sector`
Returns
-------
ds: xr.DataArray
Dataarray with mean z0 per sector
"""
mean_z0 = hist["logz0sec"] * hist["wd_count"] / hist["wd_count"].sum(dim=["sector"])
return np.exp(mean_z0.sum(dim=["sector"]))
def _normalize_hist(hist):
"""Convert frequency to probability-density histogram
Convert frequency to probability density histogram in the
WAsP core structure, i.e. with an additional variables
containing wind speed sector frequencies.
Parameters
----------
hist: xr.Dataset
Dataset with the frequency of occurence in each bin. Can have any
pywasp spatial structure.
Returns
-------
hist: xr.Dataset
Binned wind climate with probability density in each bin, on the
same spatial structure as ``hist``.
"""
counts = hist["wv_count"]
# A sector with no samples divides by one rather than by zero, so it comes
# out as a zero density -- what the WAsP core's normalize_histogram_points
# produces, whether the sector alone or the whole location is empty.
ws_count = counts.sum("wsbin", skipna=False)
wd_count = ws_count
total_count = wd_count.sum("sector", skipna=False)
wsfreq = counts / ws_count.where(ws_count != 0, 1)
# A NaN anywhere in a location's wv_count -- e.g. a masked-out grid cell in
# a caller-supplied hist -- poisons the whole location in the core: every
# sector's wsfreq comes back 0, not just the sector containing the NaN.
# total_count (skipna=False over both wsbin and sector) is NaN exactly
# when that happens, so use it to zero the whole location the same way.
wsfreq = wsfreq.where(total_count.notnull(), 0)
wdfreq = (wd_count / total_count.where(total_count != 0, 1)).where(total_count != 0)
result = hist.drop_vars("wv_count")
result["wsfreq"] = wsfreq
result["wdfreq"] = wdfreq
return result
def _finalize_meso(total_stab_histogram, wv_count, landmask=None, percentile=0.5):
"""Finalize cumulative stability histograms, optionally by surface type."""
hist = _stab_mean_std(total_stab_histogram, wv_count, percentile=percentile)
if landmask is not None:
hist = _nearest_landsea(hist, landmask)
_update_var_attrs(hist, _MEAN_STAB_ATTRS)
return hist
def _nearest_landsea(ds, landmask):
"""Interpolate stability fields separately over land and sea.
Parameters
----------
ds : xr.Dataset
Dataset containing the stability variables.
landmask : xr.DataArray
Landmask with values 0 (sea) and 1 (land).
Returns
-------
xr.Dataset
Land and sea variants of each stability variable.
"""
source_names = ("mean_temp_scale", "rms_temp_scale", "mean_pblh_scale")
surfaces = (("land", 1, "land"), ("sea", 0, "water"))
result_names = [
f"{name}_{suffix}" for suffix, _, _ in surfaces for name in source_names
]
# Compatibility for xarray 2024.9.0; remove when that version is unsupported.
crs_attrs = ds["crs"].attrs.copy() or landmask["crs"].attrs.copy()
rounded_landmask = np.round(landmask)
for suffix, landflag, surface_type in surfaces:
targets = {
name: ds[name].where(rounded_landmask == landflag) for name in source_names
}
valid = {
name: to_point(target).dropna("point") for name, target in targets.items()
}
for source_name in source_names:
targets[source_name]["crs"].attrs = crs_attrs
valid[source_name]["crs"].attrs = crs_attrs
if valid[source_name].sizes["point"] == 0:
raise ValueError(
"The landmask you passed in does not contain any pixels "
f"with {landflag} ({surface_type}). Please include such "
"points to obtain a valid estimation of the stability "
f"conditions over {surface_type}"
)
reference = valid[source_names[0]]
# Reuse one spatial search when fields share valid points; retain
# per-field interpolation for variable-specific missing values.
shared_points = all(
reference["point"].equals(valid[name]["point"]) for name in source_names[1:]
)
if shared_points:
point_coords = {
name: coord
for name, coord in reference.coords.items()
if set(coord.dims).issubset({"point"})
}
source_index = xr.DataArray(
np.arange(reference.sizes["point"]),
dims="point",
coords=point_coords,
)
nearest = interp_unstructured_like(
source_index, targets[source_names[0]], method="nearest"
)
for source_name in source_names:
if shared_points:
indexer = xr.DataArray(nearest.data, dims=nearest.dims)
selected = valid[source_name].isel(point=indexer)
target_coords = {
name: coord
for name, coord in targets[source_name].coords.items()
if set(coord.dims).issubset(selected.dims)
}
target = xr.DataArray(
selected.data,
dims=selected.dims,
coords=target_coords,
attrs=selected.attrs,
)
else:
target = interp_unstructured_like(
valid[source_name], targets[source_name], method="nearest"
)
target["crs"].attrs = crs_attrs
if "sector" in target.dims:
target = target.transpose("sector", ...)
ds[f"{source_name}_{suffix}"] = target
result = ds[result_names]
result["crs"].attrs = crs_attrs
return result
def _finalize_baro(hist, wv_count, percentile=0.5):
"""Finalize shear sums as sector magnitudes and directions."""
cond_hist, freq = _conditional_histogram(hist["sum_dugdz"], wv_count, percentile)
mean_dugdz = _sector_statistic("mean", cond_hist, freq)
cond_hist, freq = _conditional_histogram(hist["sum_dvgdz"], wv_count, percentile)
mean_dvgdz = _sector_statistic("mean", cond_hist, freq)
north = np.arctan2(-mean_dvgdz, -mean_dugdz) * RTD
south = np.arctan2(mean_dvgdz, mean_dugdz) * RTD
# DataArray.where keeps the direction's axis order, unlike xr.where.
direction = south.where(hist["south_north"] < 0, north)
clim = xr.Dataset(
{
"mean_dgdz_dir": direction,
"mean_dgdz": np.hypot(mean_dugdz, mean_dvgdz),
},
attrs=hist.attrs,
)
clim = _update_var_attrs(clim, _MEAN_BARO_ATTRS)
# Older xarray/pyproj versions can drop the scalar CRS coordinate here.
if "crs" in hist.coords:
clim = set_crs(clim, get_crs(hist))
return clim
def _validate_baro_wv_count(hist, wv_count):
"""Validate that a wind-count histogram matches a baroclinicity histogram."""
if not np.array_equal(wv_count["wsbin"], hist["wsbin"]) or not np.array_equal(
wv_count["sector"], hist["sector"]
):
raise PywaspError(
"Baroclinicity histogram counts do not align with the supplied bins"
)
def baroclinicity_histogram(
ds,
hist=None,
finalize=True,
revert_to_original=True,
wsbin_width=1,
n_wsbins=40,
n_sectors=12,
percentile=0.5,
wv_count=None,
**kw,
):
"""Add timeseries to existing histogram
Uses the baroclinicity parameters to derive
summed histogram of the geostrophic shear components.
.. warning::
This function is experimental and its signature may change.
Parameters
----------
ds: xr.Dataset
Dataset containing the variables
['wind_speed', 'wind_direction', 'dugdz', 'dvgdz'].
Can have any pywasp spatial structure
hist: xr.Dataset
Histogram with dimensions point, wsbin, sector. When supplied together
with ``finalize=True``, the matching accumulated counts must be given as
``wv_count``; otherwise a `PywaspError` is raised, because counts
covering only the final chunk would silently inflate the means.
finalize: boolean
Convert the stabilility histogram into a mean and
standard deviation?
revert_to_original: boolean
If True (default), operate on and return ds's own native spatial
structure directly; there is no promotion/reversion round trip. If
False, promote to point structure and return that; still useful when
composing with other point-structured accumulators, e.g. across
chunked calls.
wsbin_width : float
Width of wind speed bins. Only used to initialize a new histogram; it is
ignored when ``hist`` is supplied, where the bins come from ``hist``.
n_wsbins : int
Number of wind speed bins. Init-only, like ``wsbin_width``.
n_sectors : int
Number of wind direction bins. Init-only, like ``wsbin_width``.
percentile : float
Percentile of the wind speed distribution used for finalization.
wv_count : xarray.Dataset, optional
Accumulating wind-vector count histogram covering every chunk already
added to ``hist``. When supplied during finalization, the current
dataset is added to this histogram before computing the conditional
means. This follows the same workflow as stability_histogram and avoids
rebuilding counts for every chunk.
**kwargs
Backward-compatible aliases nws and nwd for the bin counts.
Returns
-------
hist: xr.Dataset
Histogram with dimensions point, wsbin, sector
with the values from ds-timeseries added.
Raises
------
PywaspError
If an unknown keyword argument is given; if a supplied ``hist`` is
finalized without the matching accumulated ``wv_count``; if the
histogram has fewer than two wind-speed bins or bins that are not
uniformly spaced; or if the supplied counts are not on the same bin
grid as the histogram.
"""
n_wsbins = kw.pop("nws", n_wsbins)
n_sectors = kw.pop("nwd", n_sectors)
if kw:
unexpected = next(iter(kw))
raise PywaspError(f"Unexpected keyword argument: {unexpected}")
# A supplied histogram may already hold accumulated sums from earlier chunks,
# so counts covering only the current chunk must never be used to finalize it.
hist_is_accumulated = hist is not None
ds_sub = ds[[WS, WD, "dugdz", "dvgdz"]]
dss, previous_hist = _histogram_layout(ds_sub, hist, revert_to_original)
hist = _histogram_sums(
dss,
values={
"sum_dugdz": dss["dugdz"],
"sum_dvgdz": dss["dvgdz"],
},
hist=previous_hist,
wsbin_width=wsbin_width,
n_wsbins=n_wsbins,
n_sectors=n_sectors,
)
hist = _update_var_attrs(hist, _HIS_BARO_ATTRS)
hist = _update_time_attributes_without_computing(hist, dss, previous_hist)
if finalize:
wv_count = _accumulate_wv_count(
dss,
hist,
wv_count,
hist_is_accumulated,
"baroclinicity",
wsbin_width,
validate_bins=True,
)
_validate_baro_wv_count(hist, wv_count)
return _finalize_baro(hist, wv_count["wv_count"], percentile)
return _update_history(hist)
[docs]
@validate_input(_BWC_SCHEMA, _ANY_3D_SPATIAL, run_extra_checks=False)
def bwc_resample_like(source, target, conf=None):
"""
Resamples a histogram to different sector or wind speed bin structure.
Parameters
----------
source: xarray.Dataset
Binned wind climate xr.Dataset object that we want to resample
target: xarray.Dataset
Binned wind climate xr.Dataset object structure that we want to resample to
Returns
-------
xarray.Dataset
Binned Wind Climate
"""
if conf is None:
conf = Config()
result = bwc_resample_sectors(source, target.sizes["sector"])
result = bwc_resample_wsbins_like(result, target, conf=conf)
return result
[docs]
@validate_input(_BWC_SCHEMA, _ANY_3D_SPATIAL, run_extra_checks=False)
def bwc_resample_sectors(source, n_sectors=12, offset=None):
"""
Resamples a histogram to different sector structure.
Parameters
----------
source: xarray.Dataset
Binned wind climate xr.Dataset object that we want to resample
target: xarray.Dataset
Binned wind climate xr.Dataset object structure that we want to resample to
offset: xarray.DataArray or float
Offset in degrees with which the histogram are rotated (positive
for clockwise rotation). If it is a DataArray, needs to have the same spatial shape as
the binned wind climate.
Returns
-------
xarray.Dataset
Binned Wind Climate
"""
if source.sizes["sector"] == n_sectors and offset is None:
return source
if isinstance(offset, xr.DataArray):
if not are_spatially_equal(source, offset):
raise ValueError(
"Input source and offset must have the same spatial structure."
)
offset = offset.astype("float32")
elif offset is None:
offset = 0.0
elif np.isscalar(offset):
offset = np.float32(offset)
else:
raise PywaspError("offset must be a scalar or xarray.DataArray")
def _resample(wv_freq, doff, n_sectors):
return _RESAMPLE_HIST(wv_freq, doff, n_sectors)
wv_freq = xr.apply_ufunc(
_resample,
(source.wsfreq * source.wdfreq).transpose(..., "wsbin", "sector"),
offset,
input_core_dims=[["wsbin", "sector"], []],
output_core_dims=[["wsbin", "sector"]],
exclude_dims={"sector"},
kwargs={"n_sectors": n_sectors},
dask="allowed",
vectorize=True,
keep_attrs=True,
)
wdcenters = create_sector_coords(n_sectors)
wv_freq = wv_freq.assign_coords({**wdcenters.coords})
spatial_dims = [d for d in wv_freq.dims if d not in ("wsbin", "sector")]
wv_freq = wv_freq.transpose("wsbin", "sector", *spatial_dims)
result = _bwc_from_counts(wv_freq)
# Carry the input's global metadata (e.g. wind_dir_crs), then set our title.
result.attrs = {k: v for k, v in source.attrs.items() if not k.startswith("_pwio")}
result.attrs["title"] = "Binned wind climate"
return _update_history(result)
[docs]
@validate_input(_BWC_SCHEMA, _ANY_3D_SPATIAL, run_extra_checks=False)
def bwc_resample_wsbins_like(
source, target, fit_weibull=True, errors="raise", conf=None
):
"""
Resamples a histogram to different wind speed bin structure.
Parameters
----------
source: xarray.Dataset
Binned wind climate xr.Dataset object that we want to resample
target: xarray.Dataset
Binned wind climate xr.Dataset object structure that we want to resample to
fit_weibull: boolean, optional
Default True, when resampling the histogram find the sectorwise A and k and
use these to resample the histogram. If false, use the A and k provided in
the dataset instead.
errors: {'raise', 'ignore'}, optional
If 'raise' (default), raises a ValueError error if probability mass
is transfered beyond the highest wind speed bin.
If "ignore", all points including those that lose probability mass
are returned.
Returns
-------
xarray.Dataset
Binned Wind Climate
"""
if conf is None:
conf = Config()
# we reset coordinates with drop=True to make sure there is no user added
# coordinates that cause the dataarrays to be different
if (
not source["sector"]
.reset_coords(drop=True)
.equals(target["sector"].reset_coords(drop=True))
):
raise ValueError(
"Input and target histogram must have the same sector structure!"
)
# source and target wsbin structure are the same, so no resampling needed
same_wsbins = (
source["wsbin"]
.reset_coords(drop=True)
.equals(target["wsbin"].reset_coords(drop=True))
)
if same_wsbins and fit_weibull:
return source
if not fit_weibull:
if not all(v in target.data_vars for v in ["A", "k"]):
raise ValueError(
"Variables A and k are needed for resampling when fit_weibull is set to False"
)
if wksp.get_spatial_struct(source) != wksp.get_spatial_struct(target):
raise PywaspError(
"When fit_weibull=False, source and target must have the same "
"spatial structure so their per-point A/k values line up."
)
source = source.transpose(..., "wsbin", "sector")
target = target.transpose(..., "wsbin", "sector")
def _resample_wsbins(wsfreq, wsceil, target_wsceil, target_ws_floor, conf):
return _RESAMPLE_HIST_WSBINS(
wsfreq, wsceil, target_wsceil, target_ws_floor, conf.climate.param
)
def _resample_wsbins_ak(wsfreq, wsceil, A, k, target_wsceil, target_ws_floor, conf):
return _RESAMPLE_HIST_WSBINS_AK(
wsfreq, wsceil, A, k, target_wsceil, target_ws_floor, conf.climate.param
)
if fit_weibull:
wsfreq, ierror = xr.apply_ufunc(
_resample_wsbins,
source.wsfreq,
source.wsceil,
target.wsceil,
target.wsfloor,
input_core_dims=[["wsbin", "sector"], *[["wsbin"]] * 3],
output_core_dims=[["wsbin", "sector"], ["sector"]],
exclude_dims={"wsbin"},
kwargs={"conf": conf},
dask="allowed",
vectorize=True,
keep_attrs=True,
)
else:
wsfreq, ierror = xr.apply_ufunc(
_resample_wsbins_ak,
source.wsfreq,
source.wsceil,
target.A,
target.k,
target.wsceil,
target.wsfloor,
input_core_dims=[
["wsbin", "sector"],
["wsbin"],
*[["sector"]] * 2,
*[["wsbin"]] * 2,
],
output_core_dims=[["wsbin", "sector"], ["sector"]],
exclude_dims={"wsbin"},
kwargs={"conf": conf},
dask="allowed",
vectorize=True,
keep_attrs=True,
)
if ierror.sum() > 0 and errors == "raise":
bad = ierror.any(dim="sector")
i_points = np.flatnonzero(np.asarray(bad.values).ravel())
raise ValueError(
f"Resampling of the wsbins in {len(i_points)} out of {bad.size} points "
f"(i={','.join(map(str, i_points))}) exceeded the number of target bins!"
)
spatial_dims = [d for d in wsfreq.dims if d not in ("wsbin", "sector")]
wsfreq = wsfreq.transpose("wsbin", "sector", *spatial_dims).assign_coords(
wsbin=target["wsbin"],
wsceil=target["wsceil"].reset_coords(drop=True),
wsfloor=target["wsfloor"].reset_coords(drop=True),
)
wdfreq = source["wdfreq"].transpose("sector", *spatial_dims)
result = xr.Dataset({"wsfreq": wsfreq, "wdfreq": wdfreq})
result = result.assign_coords(sector=source["sector"])
result.attrs = {k: v for k, v in source.attrs.items() if not k.startswith("_pwio")}
return _update_history(result)
def _get_dgdz(z, sp, levs, lats):
"""
Calculate the gradients of the zonal and meridional geostrophic wind components with respect to height.
Parameters
----------
z : array_like
Geopotential height array.
sp : array_like
Surface pressure array.
levs : array_like
Pressure levels array.
lats : array_like
Latitudes array.
Returns
-------
tuple of array_like
Gradients of the zonal (dugdz) and meridional (dvgdz) geostrophic wind components with respect to height.
"""
dugdz, dvgdz = _DGDZ(z, sp, levs, lats)
return (dugdz, dvgdz)
def _get_dgdz_t(z, sp, levs, lats):
"""
Calculate the time-dependent gradients of the zonal and meridional geostrophic wind components with respect to height.
Parameters
----------
z : array_like
Geopotential height array.
sp : array_like
Surface pressure array.
levs : array_like
Pressure levels array.
lats : array_like
Latitudes array.
Returns
-------
tuple of array_like
Time-dependent gradients of the zonal (dugdz) and meridional (dvgdz) geostrophic wind components with respect to height.
"""
dugdz, dvgdz = _DGDZ_TIME(z, sp, levs, lats)
return (dugdz, dvgdz)
def _calc_dgdz(sfc):
"""
Calculate the vertical gradients of the zonal and meridional geostrophic wind components with respect to height for a given surface dataset.
Parameters
----------
sfc : xarray.Dataset
Surface dataset containing geopotential height, surface pressure, wind speed, and wind direction.
Returns
-------
xarray.Dataset
Dataset with added gradients of the zonal (dugdz) and meridional (dvgdz) geostrophic wind components with respect to height.
"""
dugdz, dvgdz = xr.apply_ufunc(
_get_dgdz,
sfc["z"]
.transpose("west_east", "south_north", "isobaricInhPa", "time")
.astype("float32"),
(sfc["sp"].transpose("west_east", "south_north", "time") / 100).astype(
"float32"
),
sfc["isobaricInhPa"],
sfc["south_north"],
input_core_dims=[
["west_east", "south_north", "isobaricInhPa"],
["west_east", "south_north"],
["isobaricInhPa"],
["south_north"],
],
output_core_dims=[["west_east", "south_north"], ["west_east", "south_north"]],
dask="allowed",
vectorize=True,
keep_attrs=True,
)
out = xr.Dataset(
{"dugdz": dugdz, "dvgdz": dvgdz},
)
out = set_crs(out, 4326)
return out.transpose("time", "south_north", "west_east")
def _calc_dgdz_t(sfc):
"""
Calculate vertical gradients of the zonal and meridional geostrophic wind components with respect to height for a given dataset.
Parameters
----------
sfc : xarray.Dataset
Dataset containing geopotential height, surface pressure, wind speed, and wind direction.
Returns
-------
xarray.Dataset
Dataset with added time-dependent gradients of the zonal (dugdz) and meridional (dvgdz) wind components with respect to height.
"""
dugdz, dvgdz = _get_dgdz_t(
sfc["z"].transpose("west_east", "south_north", "isobaricInhPa", "time"),
sfc["sp"].transpose("west_east", "south_north", "time") / 100,
sfc["isobaricInhPa"],
sfc["south_north"],
)
out = xr.Dataset(
{"dugdz": dugdz, "dvgdz": dvgdz},
)
out = set_crs(out, 4326)
return out.transpose("time", "south_north", "west_east")
