import warnings
from collections.abc import Sequence
import numpy as np
import xarray as xr
import xrft # type: ignore
from sgs_tools.util.dask_adapt_chunking import chunk_ds
[docs]
def radial_spectrum(
ps: xr.DataArray,
fftdim: Sequence[str],
radial_bin_width: float,
bin_anchor: str = "center",
truncate: bool = True,
scaling: str = "spectrum",
prefix: str = "freq_",
) -> xr.DataArray:
r"""Isotropize a 2D power spectrum or cross spectrum
by taking an "spherical" average over the specified dimensions.
.. math::
\mathbb{F}_\text{iso}(k_r) &= \sum_{k_r: |k_r| \in [kr_0, kr_1]}
|\mathbb{F}(\mathbf{k})|^2 * w(\mathbf{k}) \\
k_r &= \langle k \rangle_{|k| \in [kr_0, kr_1]}
where :math:`k_r` is the radial wavenumber and the weights :math:`w`
are defined implicitly through the ``scaling``.
Always :math:`\sum_\mathbf{k} |\mathbb{F}(\mathbf{k})|^2 =
\sum \mathbb{F}_\text{iso}(k_r) * w(k_r)`.
This satisfies Parseval assuming
:math:`\sum_\mathbf{k} |\mathbb{F}(\mathbf{k})|^2 =
\sum \mathbb{F}(\mathbf{x})^2 * dx * dy` .
:param ps:
The power spectrum or cross spectrum to be isotropized.
:param fftdim:
The fft dimensions overwhich the isotropization must be performed.
:param radial_bin_width:
Width of radial bins in units of inverse length
(or whatever the 2d power spectrum is in)
:param truncate:
If True, the spectrum will be truncated for
wavenumbers larger than min(max(ps[fftdim].size)).
:param bin_anchor:
Where to place the radial wavenumber within the bin.
Choices ``left``, ``right``, ``centre``, ``com``.
If ``com``: compute as the centre-of-mass radius
:math:`k_r = \sum_{|k| \in [kr_0, kr_1]} (\mathbb{F} * |\mathbf{k}|)
/ \mathbb{F}_\text{iso}(k_r)` before rescaling.
Default: ``com``.
:param scaling: Rescale the power spectrum to satisfy
:math:`\sum ps = \sum \mathbb{F}_\text{iso} * w(k_r)`
* ``density``:
set :math:`w(k_r) = \pi * ((k_r^{top})^2 - (k_r^{bottom})^2)`,
where :math:`k_r^{top}` and :math:`k_r^{bottom}` are the bin edges.
* ``spectrum``:
set :math:`w(k_r) = 1`
:param prefix:
Prefix for the name of the new spectral dimension.
:return: an `xarray.DataArray` with the isotropic spectrum
coordinates for the radial wavenumber ( `prefix` r), bin width ( `prefix` dr)
and corresponding weigths ( `prefix` dA), i.e. :math:`w(k_r)`.
"""
# name of new spectral dimension
dim_name = prefix + "r"
# compute radial wavenumber bins
fftcoords = xr.Dataset(
{f"d{i}": (d, ps.coords[d].values) for i, d in enumerate(fftdim)}
)
freq_r = ((fftcoords**2).to_dataarray().sum("variable") ** 0.5).rename(dim_name)
max_linear_freq_space = max([ps.coords[d].attrs["spacing"] for d in fftdim])
if radial_bin_width <= max_linear_freq_space:
msg = (
f"radial_bin_width {radial_bin_width} <= "
f"max linear frequency spacing {max_linear_freq_space}"
", likely to have empty bins with nan values"
)
warnings.warn(UserWarning(msg), stacklevel=2)
# select radial bins
if truncate:
last_bin_edge = min([abs(fftcoords[x]).max().item() for x in fftcoords])
else:
last_bin_edge = freq_r.max().item()
# last_bin_edge *= 1.001 # add tolerance for floating point comparison
kr_bins = np.arange(0, last_bin_edge + radial_bin_width, radial_bin_width)
kr_delta = kr_bins[1:] - kr_bins[:-1]
# total spectral power in annulus
iso_ps = (
ps.groupby_bins(freq_r, bins=kr_bins, right=True, include_lowest=True)
.sum()
.rename({f"{dim_name}_bins": dim_name})
.drop_vars(dim_name)
)
# select reference wave number
if bin_anchor == "center":
kr_ref = (kr_bins[1:] + kr_bins[:-1]) / 2
elif bin_anchor == "left":
kr_ref = kr_bins[:-1]
elif bin_anchor == "right":
kr_ref = kr_bins[1:]
elif bin_anchor == "com":
kr_ref = (
(freq_r * ps)
.groupby_bins(freq_r, bins=kr_bins, right=True, include_lowest=True)
.sum()
/ iso_ps
).data
else:
raise ValueError(
f"Unrecognised bin_anchor {bin_anchor}. "
"Choose from 'center', 'left', 'right', 'com'."
)
# add a bin coordinates
iso_ps.coords[dim_name] = kr_ref
iso_ps[dim_name].attrs["anchor"] = bin_anchor
iso_ps[dim_name].attrs["spacing"] = radial_bin_width
iso_ps = iso_ps.assign_coords({prefix + "dr": (dim_name, kr_delta)})
# rescale amplitude
if scaling == "density":
annulus_area = np.pi * (kr_bins[1:] ** 2 - kr_bins[:-1] ** 2)
iso_ps = iso_ps / annulus_area
iso_ps = iso_ps.assign_coords({prefix + "dA": (dim_name, annulus_area)})
iso_ps[prefix + "dA"].attrs["description"] = "pi * (rmax^2 - rmin^2)"
if not truncate:
msg = (
"Energy density scaling is inconsistent beyond"
"the min(max(linear frequency)). Interpete with caution!"
)
warnings.warn(msg, stacklevel=2)
elif scaling == "spectrum":
iso_ps = iso_ps.assign_coords({prefix + "dA": (dim_name, np.ones_like(iso_ps))})
iso_ps[prefix + "dA"].attrs["description"] = "trivial"
else:
raise ValueError(
f"Unrecognised scaling {scaling}. Choose from 'spectrum', 'density'."
)
return iso_ps
[docs]
def spectra_1d_radial(
simulation: xr.Dataset,
hdims: Sequence[str],
power_spectra_fields: Sequence[str],
cross_spectra_fields: Sequence[tuple[str, str]],
radial_smooth_factor: int = 2,
radial_truncation: bool = False,
fillnan: float = 0.0,
) -> xr.Dataset:
r"""
Compute a 1d directional and radial power and cross spectra of all fields
in the `simultation` Dataset.
Notes: resulting spectral cooordinates are in units of inverse length,
not radians/length.
:param simulation: xarray Dataset of multidimensional fields.
must contain the set of `power_spectra_fields` and `cross_spectra_fields`
:param hdims: horizonal dimensions along which to compute linear spectra.
The radial spectrum is computed along the Euclidean radius
along vector spanned by these dimensions.
:param power_spectra_fields: sequence of fields whose power spectrum to compute
:param cross_spectra_fields: sequence of tuples of fields
whose cross-spectrum to compute
:param radial_smooth_factor: smoothing factor for radial spectral bins.
If 2 will have radial bin widht is 2*linear wavenumber.
:param radial_truncation: switch to include/exclude aliased wavenumber beyond max.
1d freqeuncy in the radial spectrum.
If **True** the result won't respect Parseval exactly.
:param fillnan: value to fill nans with in order to compute spectrum of dirty date.
If set to a nan value (e.g. np.nan) will produce nan spectra globally.
Defaults to 0.
:return: an xarray Dataset with the requested 1d and radial power and co-spectra.
The number of physical-space grid size and cell size along hdims is included
in the attributes along with the fourier normalization convention.
"""
spec = {}
extra_coords = []
for x in hdims:
extra_coords += [
d for d in simulation.coords if d != x and x in simulation[d].dims
]
sim = simulation.drop_vars(extra_coords, errors="ignore")
assert "r" not in hdims, (
"'r' dimension found in hdims, but it is reserved for radial spectra"
)
all_fields = set(power_spectra_fields)
for cfields in cross_spectra_fields:
all_fields.update(cfields)
# xrft doesn't play well with nans and chunking in to-be-spectral directions
prepped_data = chunk_ds(sim[list(all_fields)], dict.fromkeys(hdims, -1)).fillna(
fillnan
)
# power spectra
for field in power_spectra_fields:
data = prepped_data[field]
for x in hdims:
spec[f"{field}_F{x}"] = xrft.power_spectrum(
data,
dim=x,
real_dim=x,
scaling="density",
detrend=None,
prefix="k_",
true_phase=True,
)
nd_spectrum = xrft.power_spectrum(
data,
dim=hdims,
scaling="density",
prefix="k_",
detrend=None,
)
spec[f"{field}_Fr"] = radial_spectrum(
nd_spectrum,
fftdim=[f"k_{x}" for x in hdims],
radial_bin_width=min([nd_spectrum[f"k_{d}"].spacing for d in hdims]),
bin_anchor="left",
truncate=radial_truncation,
scaling="density",
prefix="k_",
)
# cross spectra
for field1, field2 in cross_spectra_fields:
data1 = prepped_data[field1]
data2 = prepped_data[field2]
for x in hdims:
spec[f"{field1}_{field2}_F{x}"] = xrft.cross_spectrum(
data1,
data2,
dim=x,
real_dim=x,
scaling="density",
prefix="k_",
true_phase=True,
detrend=None,
)
nd_spectrum = xrft.cross_spectrum(
data1, data2, dim=hdims, scaling="density", prefix="k_", true_phase=True
)
spec[f"{field1}_{field2}_Fr"] = radial_spectrum(
nd_spectrum,
fftdim=[f"k_{x}" for x in hdims],
radial_bin_width=min([nd_spectrum[f"k_{d}"].spacing for d in hdims]),
bin_anchor="left",
truncate=radial_truncation,
scaling="density",
prefix="k_",
)
# reduce along non-spectral hdims
for f in spec:
if not f.endswith("_Fr"):
non_spec_hdims = [x for x in hdims if x in spec[f].dims]
spec[f] = spec[f].sum(non_spec_hdims)
spec[f].attrs["statistic"] = f"sum_along_{non_spec_hdims}"
spec_ds = xr.Dataset(spec)
# add clarifying attributes for linear spectra
for x in hdims:
spec_ds.attrs[f"N{x}"] = sim[x].size
# the fourier spectrum call ensures that the hdims coordinates
# are regularly spaced
spec_ds.attrs[f"d{x}"] = (sim[x][1] - sim[x][0]).item()
# add header clarifying spectral conventions
spec_ds.attrs["convention"] = """Total signal power =
= (g**2).sum([<x>,...]) * d<x> * ...
= g_F<x>.sum(k_<x>) * d<x>
= (g_Fr*F<k_dA>).sum(k_dr)* """
return spec_ds