# (c) 2026 DTU Wind Energy
"""
Public interface to all of the spatial tools in WindKit
"""
__all__ = [
"are_spatially_equal",
"clip",
"clip_with_margin",
"create_dataset",
"create_point",
"create_stacked_point",
"create_cuboid",
"create_raster",
"count_spatial_points",
"equal_spatial_shape",
"mask",
"grid_convergence",
"reproject",
"warp",
"nearest_points",
"covers",
"gdf_to_ds",
"ds_to_gdf",
"add_projected_wrf_coordinates",
]
import logging
import warnings
from dataclasses import dataclass
from enum import Enum, auto
import geopandas as gpd
import numpy as np
import pyproj
import xarray as xr
from scipy.spatial import KDTree
from ..xarray_structures.metadata import (
_GLOBAL_CONVENTIONS,
_ALL_VARS_META,
_update_var_attrs,
_update_history,
)
from ._crs import set_crs, get_crs
from ._bbox import BBox
from ._dimensions import _point_dim, _stacked_point_dim, _vertical_dim, _xy_dims
from ._point import (
_mask_point,
_stack_to_point,
is_point,
is_stacked_point,
to_point,
to_stacked_point,
)
from ._struct import _apply_height_contract, _drop_pwio_attrs
from ._raster import _bounds as _raster_bounds
from ._raster import (
_can_be_raster,
_clip_to_bbox_raster,
_has_raster_dims,
_mask_raster,
)
from ._raster import _resolution as _raster_resolution
from ._raster import (
_warp_raster,
to_raster,
)
from ._struct import (
_from_scalar,
_save_stacked_index,
get_spatial_struct,
is_cuboid,
is_raster,
_is_vertical,
)
from ._vector import _clip_vector
from ._vertical import has_height_coord, has_height_dim
logger = logging.getLogger(__name__)
_STACK_ATTRS = ("_pwio_was_stacked_point", "_pwio_was_cuboid", "_pwio_orig_srs_wkt")
[docs]
def count_spatial_points(obj):
"""Get the number of spatial points for a dataset or DataArray
Parameters
----------
obj: xarray.Dataset, xarray.DataArray
WindKit xarray dataset or dataarray containing spatial dimensions.
Raises
------
ValueError
Undetectable spatial structure
"""
obj = _from_scalar(obj)
dims = obj.sizes
# Get dimension names
xdim, ydim = _xy_dims()
pt_dim = _point_dim()
vert_dim = _vertical_dim()
spt_dim = _stacked_point_dim()
if is_cuboid(obj):
return dims[vert_dim] * dims[ydim] * dims[xdim]
elif is_raster(obj):
return dims[ydim] * dims[xdim]
elif is_stacked_point(obj):
return dims[vert_dim] * dims[spt_dim]
elif is_point(obj):
return dims[pt_dim]
else:
raise ValueError("Unknown data structure cannot count points.")
def _crs_from_crs_like(crs_like):
"""Return a pyproj CRS from a WindKit spatial object or a CRS-like object.
Parameters
----------
crs_like : xarray.Dataset, xarray.DataArray, BBox, geopandas.GeoDataFrame, geopandas.GeoSeries, or any CRS-like input
Spatial object containing a CRS or an object that can be converted
to a pyproj CRS (e.g. EPSG code, dict, string, or CRS object).
Returns
-------
pyproj.crs.CRS
The pyproj CRS object.
"""
if isinstance(
crs_like, (xr.Dataset, xr.DataArray, BBox, gpd.GeoDataFrame, gpd.GeoSeries)
):
return get_crs(crs_like)
return pyproj.CRS.from_user_input(crs_like)
[docs]
def grid_convergence(ds, source=None):
"""Return grid convergence for a WindKit spatial object.
Parameters
----------
ds : xarray.Dataset or xarray.DataArray
WindKit xarray object containing spatial coordinates and a CRS.
source : xarray.Dataset, xarray.DataArray, geopandas.GeoDataFrame,
geopandas.GeoSeries, BBox, int, dict, str or pyproj.crs.CRS, optional
Reference CRS to rotate from. When provided, returns the effective
grid convergence needed to rotate directions from the ``source``
frame into the frame of ``ds`` at the locations of ``ds``.
Returns
-------
xarray.DataArray
Grid convergence in degrees, in the range (-180, 180].
A positive value means grid north of ``ds`` is rotated clockwise from
the reference north: subtract it from a wind direction expressed in
the reference frame to express that direction in the grid frame of
``ds``.
Without ``source``, the reference is true (geographic) north: the
returned angle is the classic geodetic grid convergence of the
projection of ``ds``.
With ``source``, the reference is the grid-north of ``source``: the
returned angle is the rotation needed to convert a direction expressed
in the ``source`` CRS frame into the ``ds`` CRS frame.
Raises
------
ValueError
If no CRS is found on ``ds``.
"""
ds_crs = get_crs(ds)
proj = pyproj.Proj(ds_crs)
south_north, west_east = xr.broadcast(ds["south_north"], ds["west_east"])
lon, lat = proj(west_east.values, south_north.values, inverse=True)
mc = proj.get_factors(lon, lat, False, True).meridian_convergence
if source is not None:
source_proj = pyproj.Proj(_crs_from_crs_like(source))
source_mc = source_proj.get_factors(lon, lat, False, True).meridian_convergence
mc = (mc - source_mc + 180.0) % 360.0 - 180.0
coords = {
name: coord
for name, coord in ds.coords.items()
if name == "crs" or set(coord.dims).issubset(south_north.dims)
}
mc = xr.DataArray(
mc,
dims=south_north.dims,
coords=coords,
name="grid_convergence",
)
return _update_var_attrs(
mc,
{"grid_convergence": _ALL_VARS_META["grid_convergence"]},
)
[docs]
def are_spatially_equal(obj_a, obj_b):
"""Checks that the spatial points are equivalent for both datasets
Parameters
----------
obj_a, obj_b : xarray.Dataset, xarray.DataArray
WindKit xarray dataset or dataarray containing spatial
dimensions and CRS variable
Returns
-------
bool
True if spatial coords are numpy.allclose
"""
x_dim, y_dim = _xy_dims()
if obj_a.coords[x_dim].size != obj_b.coords[x_dim].size:
return False
if obj_a.coords[y_dim].size != obj_b.coords[y_dim].size:
return False
return np.allclose(obj_a[x_dim], obj_b[x_dim]) and np.allclose(
obj_a[y_dim], obj_b[y_dim]
)
[docs]
def create_dataset(west_east, south_north, height, crs, struct="auto", thresh=1e-9):
"""Create a WindKit dataset given a set of locations.
Parameters
----------
west_east : array_like
1D array of west_east locations of interest
south_north : array_like
1D array of south_north locations of interest
height : array_like
Array of heights to create in output WindKit xarray dataset
crs : int, dict, str or pyproj.crs.CRS
Value to initialize :py:class:`pyproj.crs.CRS`
struct : str
Used to specify the type of dataset that is desired
Valid values are 'auto', 'cuboid', 'point' and
'stacked_point'. The default value of 'auto' tries to create a cuboid,
but falls back to point depending on the values specified.
thresh : float
Threshold value for raster detection
Default set to 1e-9 m
Returns
-------
xarray.Dataset
WindKit formated zero-filled Dataset with one variable "output"
and with the grid dimensions specified
Notes
-----
This function will create a WindKit formatted dataset including all of the
geospatial information that is desired.
If ds_fmt is set to "auto", the function will attempt to create a cuboid,
if the deviation between the largest and smallest interval along west_east
and south_north are lower than value in thresh (west_east and south north
can be different lengths). Otherwise, the function will attempt to make a
point dataset. Specifically, it will create a 3D point object,
if west_east, south_north and heights have the same lengths.
If west_east and south_north have the same lengths but heights length is
different the function will create a 2D point object.
Alternatively, ds_fmt can be set to desired output format, that is 'cuboid',
'stacked_point' and 'point', which accordingly requires specific structures
of west_east, south_north and heights (as described above):
* ``cuboid``:
west_east, south_north, and heights must contain unique points along their
dimension, and west_east and south_north are evenly spaced
* ``stacked_point``:
west_east and south_north must have the same length, but heights will have
unique values representing the dimension values
* ``point``:
west_east, south_north and heights must have same lengths
"""
# Ensure that all input arrays are 1d
x, y, z = np.atleast_1d(west_east, south_north, height)
if struct.lower() == "auto":
if _has_raster_dims(x, y):
struct = "cuboid"
elif (y.size == x.size) and (x.size != z.size):
struct = "stacked_point"
elif (y.size == x.size) and (y.size == z.size):
struct = "point"
else:
raise ValueError("Cannot identify struct of input data.")
# Get names of dimensions
x_dim, y_dim = _xy_dims()
z_dim = _vertical_dim()
pt_dim = _point_dim()
stacked_pt_dim = _stacked_point_dim()
if struct.lower() == "cuboid":
if not _has_raster_dims(x, y):
raise ValueError("Data cannot be converted to raster or cuboid dataset.")
# Create spatial data arrays
shape = z.size, y.size, x.size
dims = (z_dim, y_dim, x_dim)
z_coord = ((z_dim,), z)
y_coord = ((y_dim,), y)
x_coord = ((x_dim,), x)
elif struct.lower() == "stacked_point":
# Raise error if point array cannot be made
if y.size != x.size:
err_str = "south_north and west_east sizes do not match."
raise ValueError(err_str)
# make sure the heights in the z array are unique
z = np.unique(z)
# Create spatial data arrays
shape = z.size, x.size
dims = (z_dim, stacked_pt_dim)
z_coord = ((z_dim,), z)
y_coord = ((stacked_pt_dim,), y)
x_coord = ((stacked_pt_dim,), x)
elif struct.lower() == "point":
# Raise error if point array cannot be made
if (y.size != x.size) | (y.size != z.size):
err_str = (
"point dataset cannot be made with input " + "arrays of differing sizes"
)
raise ValueError(err_str)
# Create spatial data arrays
shape = (x.size,)
dims = (pt_dim,)
z_coord = ((pt_dim,), z)
y_coord = ((pt_dim,), y)
x_coord = ((pt_dim,), x)
else:
raise ValueError("Unknown struct provided.")
# Build DataArray
data = np.zeros(shape)
coords = {z_dim: z_coord, y_dim: y_coord, x_dim: x_coord}
out_da = xr.DataArray(data=data, dims=dims, coords=coords)
# Build dataset with crs for storing results
out_ds = out_da.to_dataset(name="output")
out_ds = set_crs(out_ds, crs)
out_ds[z_dim].attrs = _ALL_VARS_META[z_dim]
out_ds.attrs = _GLOBAL_CONVENTIONS
return _update_history(out_ds)
[docs]
def create_point(west_east, south_north, height, crs):
"""Create a WindKit point dataset given a set of locations.
Parameters
----------
west_east : array_like
1D array of west_east locations of interest
south_north : array_like
1D array of south_north locations of interest
height : array_like
Array of heights to create in output WindKit xarray dataset
crs : int, dict, str or pyproj.crs.CRS
Value to initialize :py:class:`pyproj.crs.CRS`
Returns
-------
xarray.Dataset
WindKit formated zero-filled Dataset with one variable "output"
and with the grid dimensions specified
Notes
-----
This function will create a WindKit formatted dataset including all of the
geospatial information that is desired.
"""
return create_dataset(west_east, south_north, height, crs, struct="point")
[docs]
def create_stacked_point(west_east, south_north, height, crs):
"""Create a WindKit stacked point dataset given a set of locations.
Parameters
----------
west_east : array_like
1D array of west_east locations of interest
south_north : array_like
1D array of south_north locations of interest
height : array_like
Array of heights to create in output WindKit xarray dataset
crs : int, dict, str or pyproj.crs.CRS
Value to initialize `pyproj.crs.CRS`
Returns
-------
xarray.Dataset
WindKit formated zero-filled Dataset with one variable "output"
and with the grid dimensions specified
Notes
-----
This function will create a WindKit formatted dataset including all of the
geospatial information that is desired.
"""
return create_dataset(west_east, south_north, height, crs, struct="stacked_point")
[docs]
def create_cuboid(west_east, south_north, height, crs):
"""Create a WindKit cuboid dataset given a set of locations.
Parameters
----------
west_east : array_like
1D array of west_east locations of interest
south_north : array_like
1D array of south_north locations of interest
height : array_like
Array of heights to create in output WindKit xarray dataset
crs : int, dict, str or pyproj.crs.CRS
Value to initialize `pyproj.crs.CRS`
Returns
-------
xarray.Dataset
WindKit formated zero-filled Dataset with one variable "output"
and with the grid dimensions specified
Notes
-----
This function will create a WindKit formatted dataset including all of the
geospatial information that is desired.
"""
return create_dataset(west_east, south_north, height, crs, struct="cuboid")
[docs]
def create_raster(west_east, south_north, crs):
"""Create a WindKit raster dataset given a set of locations.
Parameters
----------
west_east : array_like
1D array of west_east locations of interest
south_north : array_like
1D array of south_north locations of interest
crs : int, dict, str or pyproj.crs.CRS
Value to initialize `pyproj.crs.CRS`
Returns
-------
xarray.Dataset
WindKit formated zero-filled Dataset with one variable "output"
and with the grid dimensions specified
Notes
-----
This function will create a WindKit formatted dataset including all of the
geospatial information that is desired.
"""
return create_dataset(west_east, south_north, [0], crs, struct="cuboid")
def _replace_close(arr, thresh=1e-9): # pragma:no cover internal
"""Replace all values that are close to each other with an identical value
Parameters
----------
arr : array_like
Array to be examined
thresh : float
Threshold value for checking spatial separation between coordinates
Default set to 1e-9 m
Returns
-------
array_like
Array of the same shape and dtype as arr,
but with close values replaced
"""
unq = np.sort(np.unique(arr))
diff_unq = np.diff(unq)
if any(diff_unq <= thresh):
for i, diff in enumerate(diff_unq):
if diff <= thresh:
val = unq[i]
arr = np.where(np.isclose(arr, val), val, arr)
return arr
[docs]
def reproject(obj, to_crs, copy=True):
"""Reprojects WindKit object a new CRS without changing the data.
If the input is a xarray.Dataset or xarray.DataArray with
a 'cuboid' spatial structure, the spatial
structure will be changed to 'stacked_point', since the coordiates of the new
dataset will no longer be regularly spaced.
Parameters
----------
obj : geopandas.GeoDataFrame, xarray.DataArray, xarray.Dataset, or BBox
WindKit object that will be reprojected.
crs : int, dict, str or pyproj.crs.CRS
Value to initialize :py:class:`pyproj.crs.CRS`
copy : bool
If true, the object is copied. Defaults to True.
Returns
-------
geopandas.GeoDataFrame, xarray.DataArray,xarray.Dataset, or BBox
WindKit object with new projection
See Also
--------
warp
Notes
-----
This script reprojects the coordinates of the data, and reshapes the data
from ``cuboid`` or ``raster`` to ``stacked_point``. This is done so that none of
the data is interpolated, but rather the coordinates are just changed to match
those of the new projection. If you want to keep your data in ``cuboid`` or
``raster`` format, use the function ``warp`` instead, which however does do
interpolation of the data.
"""
if isinstance(obj, gpd.GeoDataFrame) or isinstance(obj, gpd.GeoSeries):
# pyproj triggers a numpy deprecation warning about ndim>0-to-scalar conversion.
# See: https://github.com/pyproj4/pyproj/issues/1307
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore", message="Conversion of an array with ndim > 0"
)
return obj.to_crs(to_crs)
elif isinstance(obj, BBox):
return obj.reproject(to_crs)
# From here we know that the input object is
# a xarray.Dataset or xarray.DataArray
if copy:
obj = obj.copy()
# Get input CRS object
try:
from_crs = get_crs(obj)
except ValueError:
raise ValueError(
"No CRS found on object! Please set with:" + " obj = set_crs(obj, crs)!"
)
# Get spatial dim names
x_dim, y_dim = _xy_dims()
# Get output CRS object
to_crs = pyproj.CRS.from_user_input(to_crs)
# Reproject dimensions
transformer = pyproj.Transformer.from_crs(from_crs, to_crs, always_xy=True)
struct = get_spatial_struct(obj)
if struct in ["raster", "cuboid"]:
obj = to_stacked_point(obj)
struct = "stacked_point"
elif struct is None:
try:
obj = _from_scalar(obj)
struct = get_spatial_struct(obj)
except Exception:
raise RuntimeError("Could not determine spatial structure of your data.")
# pyproj internally converts ndim>0 arrays to scalars, which numpy >= 1.25
# deprecates. squeeze() makes single-element (1,) arrays 0-d preemptively;
# atleast_1d() restores the 1-d shape expected downstream.
# See: https://github.com/pyproj4/pyproj/issues/1307
x_new, y_new = transformer.transform(
obj[x_dim].values.squeeze(), obj[y_dim].values.squeeze()
)
x_new = np.atleast_1d(x_new)
y_new = np.atleast_1d(y_new)
x_new = _replace_close(x_new)
y_new = _replace_close(y_new)
obj = obj.drop_vars([x_dim, y_dim])
obj = obj.assign_coords({x_dim: ((struct), x_new), y_dim: ((struct), y_new)})
obj = set_crs(obj, to_crs)
return obj
[docs]
def warp(
obj, to_crs, resolution=None, method="nearest", nodata=None, coerce_to_float=True
):
"""Warp cuboid WindKit object to another in a new CRS using data interpolation.
Parameters
----------
obj : xarray.DataArray or xarray.Dataset
WindKit object to warp.
to_crs : int, dict, str or CRS
Value to create CRS object or an existing CRS object
resolution : tuple (x resolution, y resolution) or float, optional
Target resolution, in units of target coordinate reference
system. Default: None calculates a resolution in the target crs similar
to the resolution of the original crs.
method : str
Interpolation method, passed to rasterio.warp.reproject. Defaults to "nearest".
nodata : scalar
Initial data to fill output arrays, passed to rasterio.warp.reproject. Defaults to
numpy.nan.
Returns
-------
xarray.DataArray or xarray.Dataset
Warped WindKit object.
See Also
--------
reproject
"""
kwargs = {
"resolution": resolution,
"method": method,
"nodata": nodata,
"coerce_to_float": coerce_to_float,
}
if isinstance(obj, (xr.Dataset, xr.DataArray)):
struct = get_spatial_struct(obj)
if struct in ["raster", "cuboid"]:
obj = _warp_raster(obj, to_crs, **kwargs)
if isinstance(obj, xr.Dataset):
return _update_history(obj)
else:
return obj
else:
raise ValueError(
"Only 'raster' and 'cuboid' objects currently" + " supports warping!"
)
else:
raise ValueError(
"Only xarray Dataset and DataArray objects "
+ "currently supported for warping"
)
[docs]
def mask(obj, mask, all_touched=False, invert=False, nodata=None, **kwargs):
"""Mask WindKit object with geometric mask.
Masking an object returns the same object with values outside of the masked
region filled with NaN.
Parameters
----------
obj : xarray.Dataset,xarray.DataArray
WindKit object to mask.
mask : geopandas.GeoDataFrame or BBox
Mask to mask object by.
all_touched : bool
*raster* or *cuboid* only: Include all pixels touched by the mask? False
includes only those that pass through the center. Passed to
rasterio.features.geometry_mask. Defaults to False.
invert : bool
*raster* or *cuboid* only: If True, values inside the mask are nulled;
if False, values outside the mask are nulled. The opposite is passed to
rio.features.geometry_mask. Defaults to False.
nodata : float
*raster* or *cuboid* only: If no data is not None, all masked data will be filled
with this value. Default, masked data is set to NaN.
Returns
-------
same as obj
Clipped WindKit object.
See Also
--------
clip
Note
----
This function behaves the opposite of rasterio.features.geometry_mask by default, in
that it nulls areas outside of the area of interest rather than inside.
For rasters, when the mask edges intersects with the cell centers they are not guaranteed to be
included. It is recommend to use a buffer or all_touched=True to be sure.
"""
kwargs = {
**kwargs,
"all_touched": all_touched,
"invert": invert,
"nodata": nodata,
}
if isinstance(obj, (xr.Dataset, xr.DataArray)):
struct = get_spatial_struct(obj)
if struct in ["raster", "cuboid"]:
obj = _mask_raster(obj, mask, drop=False, **kwargs)
if isinstance(obj, xr.Dataset):
return _update_history(obj)
else:
return obj
elif struct in ["point", "stacked_point"]:
obj = _mask_point(obj, mask, drop=False, **kwargs)
if isinstance(obj, xr.Dataset):
return _update_history(obj)
else:
return obj
else:
raise ValueError("Spatial structure not supported!")
else:
raise ValueError("Object not supported!")
[docs]
def clip(obj, mask, all_touched=False, invert=False, nodata=None, **kwargs):
"""Clip object to mask.
Clipping returns an object that has been reduced to the requested shape. Dropping
data that falls outside of the masked region.
Parameters
----------
obj : geopandas.GeoDataFrame, xarray.DataArray or xarray.Dataset
Object with raster-like dimensions to clip with mask.
mask : geopandas.GeoDataFrame or BBox
Geometric features to clip out of object.
all_touched : bool
*raster* or *cuboid* only: Include all pixels touched by the mask? False
inclodes only those that pass through the center. Passed to
rasterio.features.geometry_mask. Defaults to False.
invert : bool
*raster* or *cuboid* only: If True, values inside the mask are nulled;
if False, values outside the mask are nulled. The opposite is passed to
rio.features.geometry_mask. Ignored when ``mask`` is a BBox. Defaults
to False.
nodata : float
*raster* or *cuboid* only: If no data is not None, all masked data will be filled
with this value. Default, masked data is set to NaN.
kwargs : dict
Other keyword-arguments are passed to the underlying function, depending
on the type of object.
Returns
-------
geopandas.GeoDataFrame, xarray.DataArray,xarray.Dataset
Object of the same type as obj clipped by geometric features.
See Also
--------
mask
Notes
-----
This function behaves the opposite of rasterio.features.geometry_mask by default, in
that it nulls areas outside of the area of interest rather than inside.
For rasters, when the mask edges intersects with the cell centers they are not guaranteed to be
included. It is recommend to use a buffer or all_touched=True to be sure.
"""
kwargs = {**kwargs, "all_touched": all_touched, "invert": invert, "nodata": nodata}
if isinstance(obj, gpd.GeoDataFrame) or isinstance(obj, gpd.GeoSeries):
return _clip_vector(obj, mask, **kwargs)
elif isinstance(obj, (xr.Dataset, xr.DataArray)):
struct = get_spatial_struct(obj)
if struct in ["raster", "cuboid"]:
if isinstance(mask, BBox):
obj = _clip_to_bbox_raster(obj, mask)
else:
obj = _mask_raster(obj, mask, drop=True, **kwargs)
if isinstance(obj, xr.Dataset):
return _update_history(obj)
else:
return obj
elif struct in ["point", "stacked_point"]:
obj = _mask_point(obj, mask, drop=True, **kwargs)
if isinstance(obj, xr.Dataset):
return _update_history(obj)
else:
return obj
else:
raise ValueError("Spatial structure not supported!")
else:
raise ValueError("Object not supported!")
def _to_point_with_saved_index(source):
"""Convert to point structure, saving the stacked index for reversal.
Parameters
----------
source : xarray.Dataset or xarray.DataArray
WindKit object in any spatial structure.
Returns
-------
xarray.Dataset or xarray.DataArray
Object in point structure with stacked index saved in attributes.
"""
source = _from_scalar(source)
if is_point(source):
return source
obj_out, mindex = _stack_to_point(source)
_save_stacked_index(obj_out, mindex)
if isinstance(obj_out, xr.Dataset):
obj_out = _update_history(obj_out)
return obj_out
def _spatial_stack(
source, target_crs=None, revertable=True, copy=True, remove_height=False
):
"""Returns source in a revertable version of the "point" format
This routine can be used to ensure a consistent input form to external routines, by
always returning in "point" format. It can also do reprojection, and remove the
height field to make the result a 2D spatial object.
Parameters
----------
source : xarray.DataSet
WindKit dataset containing spatial dimensions and CRS variable to convert
target_crs : int, dict, str or pyproj.crs.CRS
Value to initialize `pyproj.crs.CRS`
(Default is to not reproject.)
revertable : bool
Should we retain information about the original datastructure so we can revert
this process? This is typically True, but should be False when interpolating
the data to a new projection.
copy : bool
Should we make a copy of the initial dataset? This is typically true as we don't
want to manipulate the original object, but work on a new version.
remove_height : bool
Is the resulting object always 2D? This is typically false, but can be useful in
some instances.
Returns
-------
stacked : xarray.DataSet
WindKit formated xr.DataSet as a point object on the new projection and
with additional that allow it to be converted back to its original form
using _spatial_unstack.
Notes
-----
This routine serves two purposes:
1. Convert source to a point object, storing its former structure.
2. Reproject to the target_crs if provided, and not already met.
"""
if copy:
source = source.copy(deep=True)
# For some routines, we don't want to pass a height dimesion, so remove it if asked
if remove_height:
if has_height_dim(source):
source = source.isel(height=0, drop=True)
if has_height_coord(source):
del source["height"]
# Remove attributes from previous call if found
if not is_point(source):
for attr in _STACK_ATTRS:
if attr in source.attrs:
del source.attrs[attr]
# Label original format and get 2D variables if not "point"
orig_format = get_spatial_struct(source)
if revertable and orig_format in ("stacked_point", "cuboid", "raster"):
logger.debug("Converting from original format %s", orig_format)
source.attrs[f"_pwio_was_{orig_format}"] = 1
# Identify 2D variables
if has_height_dim(source):
hgt_dim = _vertical_dim()
if isinstance(source, xr.Dataset):
for var in source.data_vars:
logger.debug("Flagging variable %s as 2d variable", var)
source[var].attrs["_pwio_data_is_2d"] = int(
hgt_dim not in source[var].dims
)
elif isinstance(source, xr.DataArray):
logger.debug("Flagging variable as 2d variable")
source.attrs["_pwio_data_is_2d"] = int(hgt_dim not in source.dims)
# Reproject if requested
if target_crs is not None:
source_crs = get_crs(source)
target_crs = pyproj.CRS.from_user_input(target_crs)
if source_crs != target_crs:
if revertable:
source.attrs["_pwio_orig_srs_wkt"] = source_crs.to_wkt()
return reproject(_to_point_with_saved_index(source), target_crs)
ds = _to_point_with_saved_index(source)
return _update_history(ds)
def _spatial_unstack(source):
"""Unstacks a point object that was created using _spatial_stack
Parameters
----------
source : xarray.DataSet
WindKit dataset containing spatial dimensions and CRS variable to revert
Returns
-------
WindKit dataset
Source dataset structured as it was originally.
Notes
-----
This is the companion function to _spatial_stack and does 2 things:
1. Converts 3D-point object to either raster or 2D-point object
2. Reprojects converted object back to its original projection
"""
# Check that source was previously reprojected
if "_pwio_orig_srs_wkt" in source.attrs:
source = reproject(source, source.attrs["_pwio_orig_srs_wkt"])
del source.attrs["_pwio_orig_srs_wkt"]
# Convert back to original format; the converters may already have
# stripped the `_pwio*` attrs on their stale-metadata fallback paths
if "_pwio_was_raster" in source.attrs:
source = to_raster(source, True)
source.attrs.pop("_pwio_was_raster", None)
elif "_pwio_was_cuboid" in source.attrs:
source = to_raster(source, True)
source.attrs.pop("_pwio_was_cuboid", None)
elif "_pwio_was_stacked_point" in source.attrs:
source = to_stacked_point(source)
source.attrs.pop("_pwio_was_stacked_point", None)
# Transpose the coordinates to match netCDF best practices
hgt_dim = (_vertical_dim(),) if _is_vertical(source) else ()
if is_stacked_point(source):
geo_dims_order = hgt_dim + ("stacked_point",)
elif is_raster(source): # also a cuboid is a raster, if it includes a hgt_dim
geo_dims_order = hgt_dim + ("south_north", "west_east")
else:
geo_dims_order = ("point",)
# Removing _pwio_data_is_2d attr from dataset or dataarray
if isinstance(source, xr.Dataset):
for var in source.data_vars.keys():
if "_pwio_data_is_2d" in source[var].attrs:
source[var].attrs.pop("_pwio_data_is_2d", None)
elif isinstance(source, xr.DataArray):
if "_pwio_data_is_2d" in source.attrs:
source.attrs.pop("_pwio_data_is_2d", None)
# Drop any leftover internal attrs from the fallback path where stale
# stacked-index metadata could not be restored safely.
source = _drop_pwio_attrs(source)
ds = source.transpose(..., *geo_dims_order)
return _update_history(ds)
def _isel_nearest(ds_ref, source_idx, ref_spatial_dims, ref_shape, out_dims):
"""Lazily select ``ds_ref`` at flat ``source_idx``, laying results along ``out_dims``.
All work is done on the integer index arrays returned by the KDTree query
together with a single :meth:`xarray.Dataset.isel`; the data variables of
``ds_ref`` are never loaded into memory.
Parameters
----------
ds_ref : xr.Dataset
Reference dataset to select from.
source_idx : numpy.ndarray
Flat indices into the raveled reference spatial space, shaped like ``out_dims``.
ref_spatial_dims : list of str
Reference spatial dimensions, in ravel order.
ref_shape : tuple of int
Sizes of ``ref_spatial_dims``.
out_dims : tuple of str
Output dimension names to lay the selection along.
Returns
-------
xr.Dataset
Lazily selected points.
"""
source_idx = np.asarray(source_idx)
unraveled = tuple(
component.reshape(source_idx.shape)
for component in np.unravel_index(source_idx.ravel(), ref_shape)
)
indexer = {
dim: xr.DataArray(unraveled[i], dims=out_dims)
for i, dim in enumerate(ref_spatial_dims)
}
nearest = ds_ref.isel(indexer)
for dim in ("point", "stacked_point"):
if dim in nearest.coords and dim not in out_dims:
nearest = nearest.drop_vars(dim)
return nearest
def _normalize_point_coords(nearest):
"""Make scalar spatial coordinates conform to the point layout."""
if "point" not in nearest.dims:
return nearest
new_coords = {}
for name in ("west_east", "south_north", "height"):
if name not in nearest.coords:
continue
coord = nearest[name]
if coord.dims == ():
values = np.full(nearest.sizes["point"], coord.item())
new_coords[name] = (("point",), values, coord.attrs)
if new_coords:
nearest = nearest.assign_coords(new_coords)
return nearest
def _spatial_dims(obj):
"""Return spatial dimensions in a deterministic internal order."""
return [
dim
for dim in ("point", "stacked_point", "south_north", "west_east", "height")
if dim in obj.dims
]
def _lookup_view(obj, dims):
"""Slice structured axes not used by lookup coordinates to their first cell."""
spatial_dims = _spatial_dims(obj)
used_dims = {dim for name in dims for dim in obj[name].dims if dim in spatial_dims}
omitted_dims = [dim for dim in spatial_dims if dim not in used_dims]
view = obj.isel({dim: 0 for dim in omitted_dims})
return view, _spatial_dims(view)
def _flatten_lookup_coords(obj, dims, spatial_dims):
"""Broadcast lookup coordinates and return a flat coordinate matrix."""
coords = [obj[name] for name in dims]
dummy = xr.DataArray(
0,
coords={dim: obj[dim] for dim in spatial_dims},
dims=spatial_dims,
)
broadcasted = list(xr.broadcast(dummy, *coords))[1:]
broadcasted = [
coord.transpose(*[dim for dim in spatial_dims if dim in coord.dims])
for coord in broadcasted
]
shape = tuple(obj.sizes[dim] for dim in spatial_dims)
arrays = np.stack([coord.values.ravel() for coord in broadcasted], axis=-1)
return arrays, shape
def _lookup_is_unique(obj, dims, spatial_dims):
"""Return whether a structured lookup is provably unique without deduplication."""
if get_spatial_struct(obj) not in ("cuboid", "raster"):
return False
for spatial_dim in spatial_dims:
matching_coords = [
obj[name] for name in dims if obj[name].dims == (spatial_dim,)
]
if not matching_coords:
return False
if not any(
np.unique(coord.values).size == coord.size for coord in matching_coords
):
return False
return True
def _broadcast_query_result(values, reduced_dims, reduced_shape, full_dims, full_shape):
"""Broadcast a reduced structured query result over omitted target axes."""
values = np.asarray(values)
trailing_shape = values.shape[1:]
values = values.reshape(reduced_shape + trailing_shape)
reshape = tuple(
size if dim in reduced_dims else 1 for dim, size in zip(full_dims, full_shape)
)
return np.broadcast_to(
values.reshape(reshape + trailing_shape), full_shape + trailing_shape
)
def _collapse_coords(nearest, coord_dims):
"""Reduce selected source coordinates to a target coordinate dimensionality.
For each named coordinate, drop any dimension not present in the requested
``coord_dims`` after verifying the coordinate is invariant along it. Returns
``None`` when a coordinate genuinely varies along a dropped dimension, i.e.
the selection cannot be represented with the requested structure. Only
coordinate metadata is touched, so laziness of the data variables is kept.
Parameters
----------
nearest : xr.Dataset
Selected dataset whose spatial coordinates may carry extra dimensions.
coord_dims : dict of str to tuple of str
Target dimensions for each spatial coordinate.
Returns
-------
xr.Dataset or None
Dataset with collapsed coordinates, or ``None`` if not representable.
"""
new_coords = {}
for name, target_dims in coord_dims.items():
if name not in nearest.coords:
continue
coord = nearest[name]
drop_dims = [d for d in coord.dims if d not in target_dims]
if drop_dims:
reduced = coord.isel({d: 0 for d in drop_dims})
# NaN-tolerant invariance check: a constant coordinate (including one
# holding NaN) must not be judged as varying, since ``NaN != NaN``.
matches = (coord == reduced) | (coord.isnull() & reduced.isnull())
if not bool(matches.all()):
return None
else:
# Already at the target dimensionality; keep it (and its attributes).
reduced = coord
# A coordinate that ends up as its own dimension coordinate (e.g. ``height``
# on a stacked_point output) must stay unique; repeated values -- when
# several target levels select the same source level -- give an ambiguous
# index, so the selection cannot form a clean structure and we return None
# (the caller raises for a structured target). Checked regardless of
# whether a dimension was dropped.
if reduced.dims == (name,):
vals = reduced.values
if np.unique(vals).size != vals.size:
return None
if drop_dims:
new_coords[name] = (reduced.dims, reduced.values, reduced.attrs)
if new_coords:
nearest = nearest.assign_coords(new_coords)
return nearest
def _target_coord_dims(ds_target):
"""Return the coordinate layout required by a target spatial structure."""
struct = get_spatial_struct(ds_target)
if struct == "point":
return {
"west_east": ("point",),
"south_north": ("point",),
"height": ("point",),
}
if struct == "stacked_point":
return {
"west_east": ("stacked_point",),
"south_north": ("stacked_point",),
"height": ("height",),
}
if struct == "cuboid":
return {
"west_east": ("west_east",),
"south_north": ("south_north",),
"height": ("height",),
}
if struct == "raster":
return {
"west_east": ("west_east",),
"south_north": ("south_north",),
"height": (),
}
if "height" in ds_target.dims:
return {
"west_east": (),
"south_north": (),
"height": ("height",),
}
return {"west_east": (), "south_north": (), "height": ()}
def _match_target_structure(nearest, ds_target):
"""Return ``nearest`` in the target layout when its coordinates are separable."""
coord_dims = _target_coord_dims(ds_target)
matched = _collapse_coords(nearest, coord_dims)
if matched is None:
return None
for name, expected_dims in coord_dims.items():
if name in matched.coords and matched[name].dims != expected_dims:
return None
target_struct = get_spatial_struct(ds_target)
if target_struct is not None and get_spatial_struct(matched) != target_struct:
return None
if target_struct in ("cuboid", "raster") and not _can_be_raster(
matched["west_east"].values,
matched["south_north"].values,
):
return None
return _apply_height_contract(matched, _vertical_dim())
def _restore_structure(nearest, ds_target):
"""Lay unique nearest points back onto the target's structure if they form it.
Converts the point-structured ``nearest`` selection to the target's
``stacked_point``/``raster``/``cuboid`` structure and validates it by comparing
spatial sizes (metadata only, so no data variables are loaded here). Returns
``None`` when the points do not form that structure, so the caller can keep a
point output rather than silently mis-shaping the result.
Parameters
----------
nearest : xr.Dataset
Point-structured selection of unique nearest points.
ds_target : xr.Dataset
Target whose structure is being restored.
Returns
-------
xr.Dataset or None
The restored dataset, or ``None`` if the points do not form the structure.
"""
n_points = nearest.sizes.get("point", 0)
try:
if is_stacked_point(ds_target):
restored = to_stacked_point(nearest)
elif is_cuboid(ds_target) or is_raster(ds_target):
restored = to_raster(nearest, True)
if is_raster(ds_target) and "height" in restored.dims:
if restored.sizes["height"] != 1:
return None
restored = restored.isel(height=0)
else:
return None
except (ValueError, KeyError):
return None
# Metadata-only completeness check: if the restored structure spans more grid
# cells than there were unique points, an incomplete grid was NaN-padded, so
# the restore mis-shaped the result and we reject it.
spatial_dims = [
d
for d in ("stacked_point", "south_north", "west_east", "height")
if d in restored.dims
]
restored_size = int(np.prod([restored.sizes[d] for d in spatial_dims]))
if restored_size != n_points:
return None
if (is_cuboid(restored) or is_raster(restored)) and not _can_be_raster(
restored["west_east"].values,
restored["south_north"].values,
):
return None
return restored
@dataclass(frozen=True)
class _NearestLookup:
"""Reference data and index used by a nearest-neighbour query."""
source: xr.Dataset
spatial_dims: tuple[str, ...]
shape: tuple[int, ...]
representative_indices: np.ndarray
tree: KDTree
@dataclass(frozen=True)
class _NearestQuery:
"""Nearest-neighbour indices and distances in the target layout."""
source_indices: np.ndarray
distances: np.ndarray
target_dims: tuple[str, ...]
target_shape: tuple[int, ...]
class _NearestOutput(Enum):
"""Ways an exact nearest-neighbour selection can be represented."""
POINT = auto()
TARGET = auto()
RESTORE_UNIQUE = auto()
def _validate_nearest_points(
ds_ref,
ds_target,
dims,
n_nearest,
return_rank,
return_distance,
keep_duplicates,
):
"""Validate nearest-neighbour arguments and lookup coordinates."""
ref_crs = get_crs(ds_ref)
target_crs = get_crs(ds_target)
if ref_crs.is_geographic or target_crs.is_geographic:
warnings.warn(
"You are doing nearest neighbour lookup in non metric coordinate systems!",
UserWarning,
stacklevel=3,
)
if ref_crs != target_crs:
raise ValueError("Datasets must have the same coordinate system!")
if not dims:
raise ValueError("dims must contain at least one coordinate name")
if n_nearest < 1:
raise ValueError("n_nearest must be at least 1")
if not keep_duplicates and n_nearest > 1:
raise ValueError("n_nearest must be 1 when keep_duplicates is False")
if return_distance and "distance" in ds_ref.variables:
raise ValueError(
"Cannot return distance because ds_ref already contains 'distance'."
)
if return_rank and n_nearest > 1 and "rank" in ds_ref.variables:
raise ValueError("Cannot return rank because ds_ref already contains 'rank'.")
for name in dims:
if name not in ds_ref.coords:
raise ValueError(
f"Dimension '{name}' not found in reference dataset coordinates."
)
if name not in ds_target.coords:
raise ValueError(
f"Dimension '{name}' not found in target dataset coordinates."
)
for label, obj in (("reference", ds_ref), ("target", ds_target)):
for name in dims:
try:
is_finite = np.isfinite(obj[name].values)
except TypeError as exc:
raise ValueError(
f"Lookup coordinate '{name}' in the {label} dataset must be numeric."
) from exc
if not bool(np.all(is_finite)):
raise ValueError(
f"Lookup coordinate '{name}' in the {label} dataset contains "
"non-finite values."
)
omitted_coords = [
name
for name in ("west_east", "south_north", "height")
if name in ds_ref.coords and name not in dims and ds_ref[name].size > 1
]
if omitted_coords:
warnings.warn(
f"Lookup dims omit source spatial coordinate(s) {omitted_coords}; "
"nearest_points will deterministically use the first source "
"representative. Include these coordinates in dims for "
"multidimensional matching.",
UserWarning,
stacklevel=3,
)
def _prepare_nearest_lookup(ds_ref, dims):
"""Build a KDTree over distinct reference locations in lookup space."""
source = _drop_pwio_attrs(_from_scalar(ds_ref))
lookup_source, spatial_dims = _lookup_view(source, dims)
arrays, shape = _flatten_lookup_coords(lookup_source, dims, spatial_dims)
if _lookup_is_unique(source, dims, spatial_dims):
representative_indices = np.arange(arrays.shape[0])
else:
_, representative_indices = np.unique(arrays, axis=0, return_index=True)
representative_indices.sort()
tree = KDTree(
arrays[representative_indices],
balanced_tree=False,
compact_nodes=False,
)
return _NearestLookup(
source=lookup_source,
spatial_dims=tuple(spatial_dims),
shape=shape,
representative_indices=representative_indices,
tree=tree,
)
def _query_nearest(lookup, ds_target, dims, n_nearest, workers):
"""Query the lookup and broadcast results to the complete target layout."""
if n_nearest > lookup.tree.n:
raise ValueError(
f"n_nearest ({n_nearest}) cannot exceed the number of reference "
f"locations in lookup space ({lookup.tree.n})."
)
target_dims = tuple(_spatial_dims(ds_target))
target_shape = tuple(ds_target.sizes[dim] for dim in target_dims)
target_lookup, reduced_dims = _lookup_view(ds_target, dims)
arrays, reduced_shape = _flatten_lookup_coords(target_lookup, dims, reduced_dims)
distances, source_indices = lookup.tree.query(
arrays,
k=n_nearest,
workers=workers,
)
source_indices = lookup.representative_indices[source_indices]
distances = _broadcast_query_result(
distances,
reduced_dims,
reduced_shape,
target_dims,
target_shape,
)
source_indices = _broadcast_query_result(
source_indices,
reduced_dims,
reduced_shape,
target_dims,
target_shape,
)
return _NearestQuery(source_indices, distances, target_dims, target_shape)
def _structure_error(reason):
"""Raise a structured-output error with the lossless point remedy."""
reason = reason.rstrip(".")
raise ValueError(
f"{reason}. Convert ds_target with to_point(ds_target) to request "
"lossless point output."
)
def _nearest_output_policy(source_indices, target_struct, n_nearest, keep_duplicates):
"""Choose how exact selections can satisfy the public output contract."""
positions = np.arange(source_indices.size)
if n_nearest > 1:
if target_struct != "point":
_structure_error(
"Multiple neighbours per target cannot represent the target "
"spatial structure"
)
return _NearestOutput.POINT, positions, False
flat_indices = source_indices.ravel()
duplicates_removed = False
if not keep_duplicates:
_, positions = np.unique(flat_indices, return_index=True)
positions.sort()
duplicates_removed = positions.size != flat_indices.size
if target_struct == "point":
return _NearestOutput.POINT, positions, duplicates_removed
if duplicates_removed:
return _NearestOutput.RESTORE_UNIQUE, positions, True
if flat_indices.size != np.unique(flat_indices).size:
_structure_error(
"Multiple target cells map to the same source cell, so the target "
"structure cannot represent every mapping"
)
return _NearestOutput.TARGET, positions, False
def _select_nearest(lookup, indices, distances, out_dims, return_distance):
"""Select exact source cells and optionally attach query distances."""
nearest = _isel_nearest(
lookup.source,
indices,
lookup.spatial_dims,
lookup.shape,
out_dims,
)
if return_distance:
distances = np.array(distances, copy=True, order="C")
nearest["distance"] = xr.DataArray(distances, dims=out_dims)
return nearest
def _select_nearest_point(lookup, indices, distances, return_distance):
"""Select exact source cells into a normalized point structure."""
nearest = _select_nearest(
lookup,
np.asarray(indices).ravel(),
np.asarray(distances).ravel(),
("point",),
return_distance,
)
return _normalize_point_coords(nearest)
def _format_nearest_result(
lookup,
query,
ds_target,
policy,
positions,
duplicates_removed,
n_nearest,
return_rank,
return_distance,
):
"""Select and format query results according to the chosen output policy."""
if duplicates_removed:
warnings.warn(
"You have multiple points in your target dataset mapping to the "
"same point in your source ds. The resulting duplicates will be "
"removed from the output! Use keep_duplicates=True to retain "
"every mapping.",
UserWarning,
stacklevel=3,
)
flat_indices = query.source_indices.ravel()
flat_distances = query.distances.ravel()
if policy is _NearestOutput.POINT:
nearest = _select_nearest_point(
lookup,
flat_indices[positions],
flat_distances[positions],
return_distance,
)
if n_nearest > 1 and return_rank:
n_targets = int(np.prod(query.target_shape) or 1)
rank = np.tile(np.arange(n_nearest), n_targets)
nearest = nearest.assign_coords(rank=("point", rank))
elif policy is _NearestOutput.RESTORE_UNIQUE:
nearest = _select_nearest_point(
lookup,
flat_indices[positions],
flat_distances[positions],
return_distance,
)
nearest = _restore_structure(nearest, ds_target)
if nearest is None:
_structure_error(
"The unique exact source selections cannot form the target "
"spatial structure"
)
else:
nearest = _select_nearest(
lookup,
query.source_indices,
query.distances,
query.target_dims,
return_distance,
)
nearest = _match_target_structure(nearest, ds_target)
if nearest is None:
_structure_error(
"The exact source selections cannot form the target spatial structure "
"without changing coordinates or adding cells"
)
return _update_history(_drop_pwio_attrs(nearest))
[docs]
def nearest_points(
ds_ref: xr.Dataset,
ds_target: xr.Dataset,
dims=["west_east", "south_north", "height"],
n_nearest=1,
return_rank=False,
return_distance=False,
keep_duplicates=True,
workers=1,
):
"""Select the nearest source point(s) for each target location.
For every location in ``ds_target``, find the nearest location(s) in
``ds_ref`` by Euclidean distance over ``dims`` and return the exact selected
source coordinates and data. No source values are interpolated.
Parameters
----------
ds_ref : xr.Dataset
Source dataset. Any WindKit spatial structure is accepted.
ds_target : xr.Dataset
Dataset whose locations and spatial structure define the result.
dims : list of str, default ["west_east", "south_north", "height"]
Coordinate names used to measure distance. When a source spatial
coordinate is omitted, its first representative is selected
deterministically and a warning is issued.
n_nearest : int, default 1
Number of distinct source locations in lookup space to return per target.
Values greater than one require a point target.
return_rank : bool, default False
Add a ``rank`` coordinate when ``n_nearest > 1``.
return_distance : bool, default False
Add a ``distance`` data variable containing query distances.
keep_duplicates : bool, default True
Keep repeated source selections. When False, repeated selections are
removed and ``n_nearest`` must be one.
workers : int, default 1
Number of workers passed to :meth:`scipy.spatial.KDTree.query`. Use ``-1``
to use all available CPU threads.
Returns
-------
xr.Dataset
Exact selected source cells represented with the target spatial structure.
Raises
------
ValueError
If coordinate systems differ; lookup coordinates are missing or
non-finite; arguments are invalid; ``distance`` or ``rank`` would collide
with an existing source variable; or the exact source selections cannot
represent the target spatial structure.
Warns
-----
UserWarning
For geographic coordinate lookup, omitted source spatial coordinates, or
removed duplicate mappings.
Notes
-----
Dimension order is not part of the return contract. Equal-distance neighbour
ordering follows SciPy and is unspecified; selected coordinates, data,
distance, and rank always remain paired. Structured output is accepted only
when the exact selected coordinates are separable and unique, and grid output
additionally satisfies WindKit's regular-raster requirements. Pass
``to_point(ds_target)`` to request lossless point output explicitly.
"""
_validate_nearest_points(
ds_ref,
ds_target,
dims,
n_nearest,
return_rank,
return_distance,
keep_duplicates,
)
lookup = _prepare_nearest_lookup(ds_ref, dims)
query = _query_nearest(lookup, ds_target, dims, n_nearest, workers)
policy, positions, duplicates_removed = _nearest_output_policy(
query.source_indices,
get_spatial_struct(ds_target),
n_nearest,
keep_duplicates,
)
return _format_nearest_result(
lookup,
query,
ds_target,
policy,
positions,
duplicates_removed,
n_nearest,
return_rank,
return_distance,
)
[docs]
def covers(obj_a, obj_b):
"""
Checks if obj_a covers obj_b
Parameters
----------
obj_a : xarray.Dataset, xarray.DataArray
WindKit xarray dataset or dataarray containing spatial dimensions.
Only cuboid structure is currently supported for obj_a.
obj_b : xarray.Dataset, xarray.DataArray
WindKit xarray dataset or dataarray containing spatial dimensions.
Returns
-------
bool
True if obj_a covers obj_b
"""
if not is_cuboid(obj_a):
raise ValueError(
"obj_a must be cuboid structure, got:", get_spatial_struct(obj_a)
)
x_dim, y_dim = _xy_dims()
z_dim = _vertical_dim()
x_min, x_max = obj_a[x_dim].min(), obj_a[x_dim].max()
y_min, y_max = obj_a[y_dim].min(), obj_a[y_dim].max()
z_min, z_max = obj_a[z_dim].min(), obj_a[z_dim].max()
x, y, z = obj_b[x_dim], obj_b[y_dim], obj_b[z_dim]
return (
(x_min <= x).all()
& (x <= x_max).all()
& (y_min <= y).all()
& (y <= y_max).all()
& (z_min <= z).all()
& (z <= z_max).all()
)
[docs]
def clip_with_margin(
obj_to_clip, obj_clipper, margin_dx_factor=5.0, cap_style=1, join_style=1
):
"""
Clip a 'raster' or 'cuboid' dataset to the bounding box of another 'raster', 'cuboid', 'stacked_point',
or 'point' dataset including a margin around that dataset to ensure several points are avaiable around
the clipper dataset from the clipped one.
Parameters
----------
obj_to_clip : xarray Dataset or DataArray
The dataset to clip
obj_clipper : xarray Dataset or DataArray
The dataset to clip to
margin_dx_factor : float, optional
The margin to add to the bounding box of the clipper dataset, by default 5.0,
i.e. 5 times the grid spacing of the dataset to clip
cap_style : int, optional
The cap style to use for buffering the bounding box of the clipper dataset,
by default 1 (round buffering)
join_style : int, optional
The join style to use for buffering the bounding box of the clipper dataset,
by default 1 (round buffering)
Returns
-------
xarray Dataset or DataArray
The clipped dataset
"""
obj_to_clip = obj_to_clip.copy()
if not isinstance(obj_to_clip, (xr.Dataset, xr.DataArray)):
raise TypeError("obj_to_clip must be an xarray Dataset or DataArray")
if not isinstance(obj_clipper, (xr.Dataset, xr.DataArray)):
raise TypeError("obj_clipper must be an xarray Dataset or DataArray")
if not is_raster(obj_to_clip):
raise ValueError("obj_to_clip must be a cuboid")
# Get the CRS of the dataset to clip and the grid spacing resolution
# to use for buffering of the bounding buffer
# (5 times the resolution of the climate dataset)
# We will use the bbox to clip the climate dataset before interpolation
crs_clip = get_crs(obj_to_clip)
spacing_clip, _ = _raster_resolution(obj_to_clip)
bbox_buffer = margin_dx_factor * spacing_clip
# Create a bounding box around the output locations buffered with enough edge points
# for the interpolation to work¨
# cap_style=1, join_style=1 is used (round buffering) to avoid issues with
# spatial datasets of points forming a straigth line (default buffer is square
# and wont buffer the points in this case)
bbox = (
BBox.from_ds(obj_clipper)
.reproject(crs_clip, use_bounds=False)
.buffer(bbox_buffer, cap_style=cap_style, join_style=join_style)
.envelope()
)
# check if bbox is outside bounds of the dataset
bounds_to_clip = _raster_bounds(obj_to_clip)
bounds_bbox = bbox.bounds()
# Use min/max to handle datasets where coordinate bounds may be
# reversed (e.g., ERA5/CFSR with south_north going from 90 to -90).
clip_x_min = min(bounds_to_clip[0], bounds_to_clip[2])
clip_x_max = max(bounds_to_clip[0], bounds_to_clip[2])
clip_y_min = min(bounds_to_clip[1], bounds_to_clip[3])
clip_y_max = max(bounds_to_clip[1], bounds_to_clip[3])
x_size = clip_x_max - clip_x_min
if (
bbox.crs.is_geographic
and x_size >= 360
and (bounds_bbox[0] < clip_x_min or bounds_bbox[2] > clip_x_max)
):
raise ValueError(
"Bounding box is crossing the date line, this may result in missing data!"
)
elif (
bounds_bbox[0] < clip_x_min
or bounds_bbox[1] < clip_y_min
or bounds_bbox[2] > clip_x_max
or bounds_bbox[3] > clip_y_max
):
warnings.warn(
"Clipping bbox including margin is outside the bounds of the dataset to clip!",
stacklevel=2,
)
# Clip the dataset to the bounding box
return clip(obj_to_clip, bbox)
[docs]
def equal_spatial_shape(obj_a, obj_b):
"""Check if two spatial objects have the same shape.
Parameters
----------
obj_a : xr.Dataset or xr.DataArray
Spatial object.
obj_b : xr.Dataset or xr.DataArray
Spatial object.
Returns
-------
bool
True if the objects have the same shape, False otherwise.
"""
struct_a = get_spatial_struct(obj_a)
struct_b = get_spatial_struct(obj_b)
if struct_a != struct_b:
return False
if struct_a == "point":
return obj_a.point.size == obj_b.point.size
elif struct_a == "stacked_point":
return (
obj_a.stacked_point.size == obj_b.stacked_point.size
and obj_a.height.size == obj_b.height.size
)
elif struct_a == "cuboid":
return (
obj_a.west_east.size == obj_b.west_east.size
and obj_a.south_north.size == obj_b.south_north.size
and obj_a.height.size == obj_b.height.size
)
elif struct_a == "raster":
return (
obj_a.west_east.size == obj_b.west_east.size
and obj_a.south_north.size == obj_b.south_north.size
)
else:
raise ValueError(f"Unknown spatial structure '{struct_a}'.")
[docs]
def ds_to_gdf(ds, include_height=False):
"""Convert windkit spatial structure to geopandas dataframe.
Parameters
----------
ds : xr.Dataset or xr.DataArray
Spatial object.
include_height : boolean, optional
Default False, do not include the height dimension in the geopandas dataframe
Returns
-------
gdf : gpd.GeoDataFrame
GeoDataFrame with columns ``x``, ``y`` and optionally ``z`` if include_height is True
Notes
-----
If a coordinate ``name`` is present on the dataset, it will be included
in the geodataframe as well. This can be convenient when you have mast
locations that have a name/label. Note that if your dataset is in cuboid
or stacked_point format, it will be flattened to a point structure.
"""
point_ds = to_point(ds)
in_crs = get_crs(ds)
struct = (point_ds.coords["west_east"], point_ds.coords["south_north"])
if include_height:
struct = struct + (point_ds.coords["height"],)
try:
gdf = gpd.GeoDataFrame(
{"name": point_ds["name"]}, geometry=gpd.points_from_xy(*struct), crs=in_crs
)
except Exception:
gdf = gpd.GeoDataFrame(geometry=gpd.points_from_xy(*struct), crs=in_crs)
return gdf.drop_duplicates()
def _get_coords_from_gdf(gdf, height=None):
"""
Extracts the x, y, and z coordinates from a GeoDataFrame. If the GeoDataFrame
does not contain z values, a specified height is used for all points.
Parameters
----------
gdf : GeoDataFrame
The GeoDataFrame from which to extract the coordinates. The GeoDataFrame
should contain a 'geometry' column with Point geometries. If the points
do not have a z coordinate, the 'height' parameter must be provided.
height : float or None
The height to use for all points if the GeoDataFrame does not contain
z values. If the GeoDataFrame does contain z values, this parameter is
ignored. If the GeoDataFrame does not contain z values and this parameter
is None, a ValueError is raised.
Returns
-------
x, y, z : ndarray
The x, y, and z coordinates of the points in the GeoDataFrame. The
coordinates are returned as separate 1D numpy arrays.
Raises
------
ValueError
If the GeoDataFrame does not contain z values and 'height' is None.
"""
if all(gdf.geometry.has_z):
df = gdf.geometry.get_coordinates(include_z=True)
else:
df = gdf.geometry.get_coordinates(include_z=False)
if height is None:
raise ValueError(
"You must provide a height for the dataset, because the geodataframe only contains x and y!"
)
df["z"] = height
if "name" in gdf.columns:
df["name"] = gdf["name"]
return df
[docs]
def gdf_to_ds(gdf, height=None, struct=None):
"""Convert geopandas dataframe to windkit spatial structure.
Parameters
----------
gdf : gpd.GeoDataFrame
Geopandas dataframe with geometry column
height : float, optional
The value of the height coordinate if it is not included in the gpd.GeoDataFrame.
struct : {"point", "stacked_point", "cuboid"}
Default None, which means detect the best fitting windkit spatial structure.
Returns
-------
ds : xr.Dataset
dataset with a windkit spatial structure.
Notes
-----
If a column ``name`` is present on the GeoDataFrame and your data is
in point or stacked_point structure, it will be included in the xr.Dataset as well.
This can be convenient when you have mast locations that have a name/label.
"""
df = _get_coords_from_gdf(gdf, height=height)
if struct is None:
z = df.set_index("z").index.unique()
# check if dataset can be a raster
if _can_be_raster(df.x, df.y):
struct = "cuboid"
# in horizontal plane we can have a cuboid, but height dimension is varying, so return point
if np.mod(gdf.size, z.size) != 0:
struct = "point"
elif np.mod(gdf.size, z.size) == 0:
struct = "stacked_point"
else:
struct = "point"
if struct == "cuboid":
x = sorted(set(df.x))
y = sorted(set(df.y))
z = sorted(set(df.z))
elif struct == "stacked_point":
df = df[~df.set_index(["x", "y"]).index.duplicated(keep="first")]
x, y = df[["x", "y"]].values.T
z = df.set_index(["z"]).index.unique()
ds = create_dataset(x, y, z, struct=struct, crs=gdf.crs)
if "name" in gdf.columns and (struct == "point" or struct == "stacked_point"):
ds.coords["name"] = (struct, df["name"])
return ds
[docs]
def add_projected_wrf_coordinates(ds):
"""Add the west_east, south_north coordinates to a WRF xarray.Dataset output.
Parameters
----------
ds: xarray.Dataset
Dataset that is the output of WRF, with TRUELAT1, TRUELAT2, MOAD_CEN_LAT, STAND_LON
coordinates.
Returns
-------
ds : xr.Dataset
dataset with the windkit west_east, south_north coordinates.
"""
if ds.MAP_PROJ == 1:
wrf_proj = pyproj.Proj(
proj="lcc",
lat_1=ds.TRUELAT1,
lat_2=ds.TRUELAT2,
lat_0=ds.MOAD_CEN_LAT,
lon_0=ds.STAND_LON,
a=6370000,
b=6370000,
)
# I have not verified if this works, easiest way to check is just write out
# the wrfout and check that you can obtain the given lat/lon after reprojection
# as described here:
# https://fabienmaussion.info/2018/01/06/wrf-projection/
# elif ds.MAP_PROJ == 2:
# wrf_proj = pyproj.Proj(proj='stere',
# lat_1=ds.TRUELAT1, lat_2=ds.TRUELAT2,
# lat_0=ds.MOAD_CEN_LAT, lon_0=ds.STAND_LON,
# a=6370000, b=6370000)
# elif ds.MAP_PROJ == 0:
# wrf_proj = pyproj.Proj(proj='eqc',
# lat_1=ds.TRUELAT1, lat_2=ds.TRUELAT2,
# lat_0=ds.MOAD_CEN_LAT, lon_0=ds.STAND_LON,
# a=6370000, b=6370000)
elif ds.MAP_PROJ == 3:
# radius earth has to be specified, see here:
# https://gis.stackexchange.com/questions/203923/using-pyproj-and-mercator-x-0-0-and-y-0-0
rad_earth = 6370000 * np.cos(np.radians(ds.TRUELAT1))
wrf_proj = pyproj.Proj(
proj="merc",
lat_1=ds.TRUELAT1,
lat_2=ds.TRUELAT2,
lat_0=ds.MOAD_CEN_LAT,
lon_0=ds.STAND_LON,
a=rad_earth,
b=rad_earth,
)
else:
raise NotImplementedError(
"Only llc and mercator (MAP_PROJ=1 or 3) supported at the moment!"
)
wgs_proj = pyproj.Proj(proj="latlong", datum="WGS84")
trans = pyproj.Transformer.from_proj(wgs_proj, wrf_proj)
e, n = trans.transform(ds.CEN_LON, ds.CEN_LAT)
dx, dy = ds.DX, ds.DY
nx, ny = ds.dims["west_east"], ds.dims["south_north"]
x0 = -(nx - 1) / 2.0 * dx + e
y0 = -(ny - 1) / 2.0 * dy + n
x = np.arange(nx) * dx + x0
y = np.arange(ny) * dy + y0
ds = ds.assign_coords(
west_east=(("west_east",), x), south_north=(("south_north"), y)
)
ds = set_crs(ds, wrf_proj.crs)
return ds
