def power_spectrum(da, spacing_tol=1e-3, dim=None, shift=True, detrend=None, density=True,
window=False):
"""
Calculates the power spectrum of da.
.. math::
da' = da - \overline{da}
ps = \mathbb{F}(da') * {\mathbb{F}(da')}^*
Parameters
----------
da : `xarray.DataArray`
The data to be transformed
spacing_tol: float (default)
Spacing tolerance. Fourier transform should not be applied to uneven grid but
this restriction can be relaxed with this setting. Use caution.
dim : list (optional)
The dimensions along which to take the transformation. If `None`, all
dimensions will be transformed.
shift : bool (optional)
Whether to shift the fft output.
detrend : str (optional)
If `constant`, the mean across the transform dimensions will be
subtracted before calculating the Fourier transform (FT).
If `linear`, the linear least-square fit will be subtracted before
the FT.
density : list (optional)
If true, it will normalize the spectrum to spectral density
window : bool (optional)
Whether to apply a Hann window to the data before the Fourier
transform is taken
Returns
-------
ps : `xarray.DataArray`
Two-dimensional power spectrum
"""
if dim is None:
dim = da.dims
# the axes along which to take ffts
axis_num = [da.get_axis_num(d) for d in dim]
N = [da.shape[n] for n in axis_num]
daft = dft(da, spacing_tol,
dim=dim, shift=shift, detrend=detrend,
window=window)
coord = list(daft.coords)
ps = (daft * np.conj(daft)).real
if density:
ps /= (np.asarray(N).prod())**2
for i in dim:
ps /= daft['freq_' + i + '_spacing']
return ps
评论列表
文章目录