""""""
"""
This file is part of WaveBreaking.
WaveBreaking provides indices to detect, classify
and track Rossby Wave Breaking (RWB) in climate and weather data.
The tool was developed during my master thesis at the University of Bern.
Link to thesis: https://occrdata.unibe.ch/students/theses/msc/406.pdf
---
Spatial pre-processing functions
"""
__author__ = "Severin Kaderli"
__license__ = "MIT"
__email__ = "severin.kaderli@unibe.ch"
# import modules
import xarray as xr
import numpy as np
from scipy import ndimage
from wavebreaking.utils.data_utils import check_argument_types, get_dimension_attributes
[docs]@check_argument_types(["u", "v"], [xr.DataArray, xr.DataArray])
@get_dimension_attributes("u")
def calculate_momentum_flux(u, v, *args, **kwargs):
"""
Calculate the momentum flux derived from the product of the deviations
of both wind components from the zonal mean.
Dimension names ("time_name", "lon_name", "lat_name"), size ("ntime", "nlon", "nlat")
and resolution ("dlon", "dlat") can be passed as key=value arguments.
Parameters
----------
u : xarray.DataArray
zonal (x) component of the wind
v : xarray.DataArray
meridional (y) component of the wind
Returns
-------
momentum flux: xarray.DataArray
Data containing the momentum flux
"""
# calculate deviations from the zonal mean
u_prime = u - u.mean(kwargs["lon_name"])
v_prime = v - v.mean(kwargs["lon_name"])
# mflux is given by the product of the deviation of both wind components
mflux = u_prime * v_prime
mflux.name = "mflux"
return mflux
[docs]@check_argument_types(["data"], [xr.DataArray])
@get_dimension_attributes("data")
def calculate_smoothed_field(
data,
passes,
weights=np.array([[0, 1, 0], [1, 2, 1], [0, 1, 0]]),
mode="wrap",
*args,
**kwargs
):
"""
Calculate smoothed field based on a two-dimensional weight kernel
and multiple smoothing passes. Default weight kernel is a 3x3
5-point smoothing with double-weighted centre. The arguments
"weight" and "mode" must be accepted by scipy.ndimage.convolve.
Values at the latitude border are always set to NaN.
Dimension names ("time_name", "lon_name", "lat_name"), size ("ntime", "nlon", "nlat")
and resolution ("dlon", "dlat") can be passed as key=value arguments.
Parameters
----------
data : xarray.DataArray
data to smooth
passes : int or float
number of smoothing passes of the 5-point smoothing
weigths : array_like, optional
array of weight, two-dimensional
(see scipy.ndimage.convolve function)
mode : string, optional
defines how the array is extended at boundaries
(see scipy.ndimage.convolve function)
Returns
-------
smoothed data: xarray.DataArray
Data containing the smoothed field
"""
# perform smoothing
smoothed = []
for step in data[kwargs["time_name"]]:
temp = data.sel({kwargs["time_name"]: step})
for p in range(passes):
temp = ndimage.convolve(temp, weights=weights, mode=mode) / np.sum(weights)
# set latitude border values to nan
border_size = int(weights.shape[0] / 2 + 0.5)
temp[np.arange(-border_size, border_size), :] = np.nan
smoothed.append(temp)
# define DataArray
da = xr.DataArray(
smoothed,
coords=[
data[kwargs["time_name"]],
data[kwargs["lat_name"]],
data[kwargs["lon_name"]],
],
)
# set name
da.name = "smooth_" + data.name
# assign attributes
da = da.assign_attrs(data.attrs)
da.attrs["smooth_passes"] = passes
return da