def test_spectrum_dim(self):
N = 16
da = xr.DataArray(np.random.rand(2,N,N), dims=['time','y','x'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'y':range(N),'x':range(N)}
)
ps = xrft.power_spectrum(da, dim='y', real='x', window=True,
density=False, detrend='constant')
npt.assert_array_equal(ps.values,
xrft.power_spectrum(da, dim=['y'],
real='x', window=True,
density=False,
detrend='constant').values)
da2 = xr.DataArray(np.random.rand(2,N,N), dims=['time','y','x'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'y':range(N), 'x':range(N)})
cs = xrft.cross_spectrum(da, da2, dim='y',
shift=True, window=True,
detrend='constant')
npt.assert_array_equal(xrft.cross_spectrum(da, da2, dim=['y'],
shift=True, window=True,
detrend='constant').values,
cs.values
)
assert ps.dims == ('time','freq_y','freq_x')
assert cs.dims == ('time','freq_y','x')
def compute_spectra_ll_month(yymm):
dsc = xr.open_dataset('../data/llc4320/gridded/llc_T10.3_20%s_cur.nc' %yymm)
dsh = xr.open_dataset('../data/llc4320/gridded/llc_T10.3_20%s_hs.nc' %yymm)
ni = np.where(dsc.x==200000)[0][0]
nf = np.where(dsc.x==600000)[0][0]
u_mix = dsc.ucur[:, ni:nf, ni:nf]
v_mix = dsc.vcur[:, ni:nf, ni:nf]
hs = dsh.hs[:, ni:nf, ni:nf]
dx = (dsh.x.values[1] - dsh.x.values[0])
u_phi, v_phi, u_psi, v_psi = helm_decompose_flow(u_mix, v_mix, dx)
us_hat2 = xrft.power_spectrum(u_psi, dim=['x','y'], detrend='linear', window=True)
vs_hat2 = xrft.power_spectrum(v_psi, dim=['x','y'], detrend='linear', window=True)
ekes_hat2 = 0.5*(us_hat2 + vs_hat2)
nk = us_hat2.shape[-1]
ekes_y = ekes_hat2.sum(dim='freq_x')[:,nk//2:]*ekes_hat2.freq_x_spacing
up_hat2 = xrft.power_spectrum(u_phi, dim=['x','y'], detrend='linear', window=True)
vp_hat2 = xrft.power_spectrum(v_phi, dim=['x','y'], detrend='linear', window=True)
ekep_hat2 = 0.5*(up_hat2 + vp_hat2)
ekep_y = ekep_hat2.sum(dim='freq_x')[:,nk//2:]*ekep_hat2.freq_x_spacing
hs_hat2 = xrft.power_spectrum(hs, dim=['x','y'], detrend='linear', window=True)
nk = hs_hat2.shape[-1]
hsy = hs_hat2.sum(dim='freq_x')[:,nk//2:]* hs_hat2.freq_x_spacing
N = hsy.time.size
ky = hsy.freq_y*1e3
Es = ekes_y.mean(dim='time')*1e-3
Ep = ekep_y.mean(dim='time')*1e-3
Ehs = hsy.mean(dim='time')*1e-3
Es, Ep, Ehs = dask.compute(Es, Ep, Ehs)
return N, ky, Es, Ep, Ehs
def test_power_spectrum(self, dask):
"""Test the power spectrum function"""
N = 16
da = xr.DataArray(np.random.rand(2, N, N),
dims=['time', 'y', 'x'],
coords={
'time':
np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'y':
range(N),
'x':
range(N)
})
if dask:
da = da.chunk({'time': 1})
ps = xrft.power_spectrum(da,
dim=['y', 'x'],
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da, dim=['y', 'x'], detrend='constant', window=True)
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
ps = xrft.power_spectrum(da,
dim=['y'],
real='x',
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da,
dim=['y'],
real='x',
detrend='constant',
window=True)
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
### Normalized
ps = xrft.power_spectrum(da,
dim=['y', 'x'],
window=True,
detrend='constant')
daft = xrft.dft(da, dim=['y', 'x'], window=True, detrend='constant')
test = np.real(daft * np.conj(daft)) / N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove least-square fit
ps = xrft.power_spectrum(da,
dim=['y', 'x'],
window=True,
density=False,
detrend='linear')
daft = xrft.dft(da, dim=['y', 'x'], window=True, detrend='linear')
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
def test_power_spectrum():
"""Test the power spectrum function"""
N = 16
da = xr.DataArray(np.random.rand(N, N),
dims=['x', 'y'],
coords={
'x': range(N),
'y': range(N)
})
ps = xrft.power_spectrum(da,
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da, dim=None, shift=True, detrend='constant', window=True)
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Normalized
dim = da.dims
ps = xrft.power_spectrum(da, window=True, detrend='constant')
daft = xrft.dft(da, window=True, detrend='constant')
coord = list(daft.coords)
test = np.real(daft * np.conj(daft)) / N**4
for i in range(len(dim)):
test /= daft[coord[-i - 1]].values
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove mean
da = xr.DataArray(np.random.rand(5, 20, 30),
dims=['time', 'y', 'x'],
coords={
'time': np.arange(5),
'y': np.arange(20),
'x': np.arange(30)
})
ps = xrft.power_spectrum(da,
dim=['y', 'x'],
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da, dim=['y', 'x'], window=True, detrend='constant')
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove least-square fit
da_prime = np.zeros_like(da.values)
for t in range(5):
da_prime[t] = numpy_detrend(da[t].values)
da_prime = xr.DataArray(da_prime, dims=da.dims, coords=da.coords)
ps = xrft.power_spectrum(da_prime,
dim=['y', 'x'],
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da_prime, dim=['y', 'x'], window=True, detrend='constant')
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
def test_power_spectrum(self, dask):
"""Test the power spectrum function"""
N = 16
da = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "y", "x"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"y": range(N),
"x": range(N),
},
)
if dask:
da = da.chunk({"time": 1})
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window=True,
density=False,
detrend="constant")
daft = xrft.dft(da, dim=["y", "x"], detrend="constant", window=True)
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
ps = xrft.power_spectrum(da,
dim=["y"],
real="x",
window=True,
density=False,
detrend="constant")
daft = xrft.dft(da,
dim=["y"],
real="x",
detrend="constant",
window=True)
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
### Normalized
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window=True,
detrend="constant")
daft = xrft.dft(da, dim=["y", "x"], window=True, detrend="constant")
test = np.real(daft * np.conj(daft)) / N**4
dk = np.diff(np.fft.fftfreq(N, 1.0))[0]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
### Remove least-square fit
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window=True,
density=False,
detrend="linear")
daft = xrft.dft(da, dim=["y", "x"], window=True, detrend="linear")
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
def test_spectrum_dim(self):
N = 16
da = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "y", "x"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"y": range(N),
"x": range(N),
},
)
ps = xrft.power_spectrum(da,
dim="y",
real="x",
window=True,
density=False,
detrend="constant")
npt.assert_array_equal(
ps.values,
xrft.power_spectrum(da,
dim=["y"],
real="x",
window=True,
density=False,
detrend="constant").values,
)
da2 = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "y", "x"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"y": range(N),
"x": range(N),
},
)
cs = xrft.cross_spectrum(da,
da2,
dim="y",
shift=True,
window=True,
detrend="constant")
npt.assert_array_equal(
xrft.cross_spectrum(da,
da2,
dim=["y"],
shift=True,
window=True,
detrend="constant").values,
cs.values,
)
assert ps.dims == ("time", "freq_y", "freq_x")
assert cs.dims == ("time", "freq_y", "x")
def test_window_single_dim():
# Julius' example
# https://github.com/rabernat/xrft/issues/16
data = xr.DataArray(np.random.random([20,30,100]),
dims=['time','lat','lon'],
coords={'time':range(20),'lat':range(30),'lon':range(100)})
ps = xrft.power_spectrum(data, dim=['time'], window=True)
# make sure it works with dask data
ps = xrft.power_spectrum(data.chunk(), dim=['time'], window=True)
ps.load()
def test_window_single_dim():
# Julius' example
# https://github.com/rabernat/xrft/issues/16
data = xr.DataArray(np.random.random([20,30,100]),
dims=['time','lat','lon'],
coords={'time':range(20),'lat':range(30),'lon':range(100)})
ps = xrft.power_spectrum(data, dim=['time'], window=True)
# make sure it works with dask data
ps = xrft.power_spectrum(data.chunk(), dim=['time'], window=True)
ps.load()
def test_parseval():
"""Test whether the Parseval's relation is satisfied."""
N = 16
da = xr.DataArray(np.random.rand(N, N),
dims=['x', 'y'],
coords={
'x': range(N),
'y': range(N)
})
da2 = xr.DataArray(np.random.rand(N, N),
dims=['x', 'y'],
coords={
'x': range(N),
'y': range(N)
})
dim = da.dims
delta_x = []
for d in dim:
coord = da[d]
diff = np.diff(coord)
delta = diff[0]
delta_x.append(delta)
window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
ps = xrft.power_spectrum(da, window=True, detrend='constant')
da_prime = da.values - da.mean(dim=dim).values
npt.assert_almost_equal(ps.values.sum(), (np.asarray(delta_x).prod() *
((da_prime * window)**2).sum()),
decimal=5)
cs = xrft.cross_spectrum(da, da2, window=True, detrend='constant')
da2_prime = da2.values - da2.mean(dim=dim).values
npt.assert_almost_equal(cs.values.sum(), (np.asarray(delta_x).prod() *
((da_prime * window) *
(da2_prime * window)).sum()),
decimal=5)
d3d = xr.DataArray(np.random.rand(N, N, N),
dims=['time', 'y', 'x'],
coords={
'time': range(N),
'y': range(N),
'x': range(N)
}).chunk({'time': 1})
ps = xrft.power_spectrum(d3d,
dim=['x', 'y'],
window=True,
detrend='linear')
npt.assert_almost_equal(
ps[0].values.sum(),
(np.asarray(delta_x).prod() *
((numpy_detrend(d3d[0].values) * window)**2).sum()),
decimal=5)
def test_window_single_dim():
# Julius' example
# https://github.com/rabernat/xrft/issues/16
data = xr.DataArray(
np.random.random([20, 30, 100]),
dims=["time", "lat", "lon"],
coords={
"time": range(20),
"lat": range(30),
"lon": range(100)
},
)
ps = xrft.power_spectrum(data, dim=["time"], window="hann")
# make sure it works with dask data
ps = xrft.power_spectrum(data.chunk(), dim=["time"], window="hann")
ps.load()
def _test_iso(theta):
ps = xrft.power_spectrum(theta, spacing_tol, dim=dims)
ps = np.sqrt(ps.freq_x**2 + ps.freq_y**2)
ps_iso = xrft.isotropize(ps, fftdim, nfactor=nfactor)
assert len(ps_iso.dims) == 1
assert ps_iso.dims[0] == "freq_r"
npt.assert_allclose(ps_iso, ps_iso.freq_r**2, atol=0.02)
def test_dft_nocoords():
# Julius' example
# https://github.com/rabernat/xrft/issues/17
data = xr.DataArray(np.random.random([20, 30, 100]),
dims=["time", "lat", "lon"])
dft = xrft.dft(data, dim=["time"])
ps = xrft.power_spectrum(data, dim=["time"])
def test_dft_2d(self):
"""Test the discrete Fourier transform on 2D data"""
N = 16
da = xr.DataArray(np.random.rand(N, N),
dims=["x", "y"],
coords={
"x": range(N),
"y": range(N)
})
ft = xrft.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
ft = xrft.dft(da, shift=False, window=True, detrend="constant")
dim = da.dims
window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
da_prime = (da - da.mean(dim=dim)).values
npt.assert_almost_equal(ft.values, np.fft.fftn(da_prime * window))
da = xr.DataArray(
np.random.rand(N, N),
dims=["x", "y"],
coords={
"x": range(N, 0, -1),
"y": range(N, 0, -1)
},
)
assert (xrft.power_spectrum(da, shift=False, density=True) >=
0.0).all()
def test_parseval():
"""Test whether the Parseval's relation is satisfied."""
N = 16
da = xr.DataArray(np.random.rand(N,N),
dims=['x','y'], coords={'x':range(N), 'y':range(N)})
da2 = xr.DataArray(np.random.rand(N,N),
dims=['x','y'], coords={'x':range(N), 'y':range(N)})
dim = da.dims
delta_x = []
for d in dim:
coord = da[d]
diff = np.diff(coord)
delta = diff[0]
delta_x.append(delta)
window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
ps = xrft.power_spectrum(da, window=True, detrend='constant')
da_prime = da.values - da.mean(dim=dim).values
npt.assert_almost_equal(ps.values.sum(),
(np.asarray(delta_x).prod()
* ((da_prime*window)**2).sum()
), decimal=5
)
cs = xrft.cross_spectrum(da, da2, window=True, detrend='constant')
da2_prime = da2.values - da2.mean(dim=dim).values
npt.assert_almost_equal(cs.values.sum(),
(np.asarray(delta_x).prod()
* ((da_prime*window)
* (da2_prime*window)).sum()
), decimal=5
)
d3d = xr.DataArray(np.random.rand(N,N,N),
dims=['time','y','x'],
coords={'time':range(N), 'y':range(N), 'x':range(N)}
).chunk({'time':1})
ps = xrft.power_spectrum(d3d, dim=['x','y'], window=True, detrend='linear')
npt.assert_almost_equal(ps[0].values.sum(),
(np.asarray(delta_x).prod()
* ((numpy_detrend(d3d[0].values)*window)**2).sum()
), decimal=5
)
def psd_based_scores(ds_oi, ds_ref):
logging.info(' Compute PSD-based scores...')
with ProgressBar():
# Compute error = SSH_reconstruction - SSH_true
err = (ds_oi['sossheig'] - ds_ref['sossheig'])
err = err.chunk({
"lat": 1,
'time': err['time'].size,
'lon': err['lon'].size
})
# make time vector in days units
err['time'] = (err.time - err.time[0]) / numpy.timedelta64(1, 'D')
# Rechunk SSH_true
signal = ds_ref['sossheig'].chunk({
"lat": 1,
'time': ds_ref['time'].size,
'lon': ds_ref['lon'].size
})
# make time vector in days units
signal['time'] = (signal.time - signal.time[0]) / numpy.timedelta64(
1, 'D')
# Compute PSD_err and PSD_signal
psd_err = xrft.power_spectrum(err,
dim=['time', 'lon'],
detrend='linear',
window=True).compute()
psd_signal = xrft.power_spectrum(signal,
dim=['time', 'lon'],
detrend='linear',
window=True).compute()
# Averaged over latitude
mean_psd_signal = psd_signal.mean(dim='lat').where(
(psd_signal.freq_lon > 0.) & (psd_signal.freq_time > 0), drop=True)
mean_psd_err = psd_err.mean(dim='lat').where(
(psd_err.freq_lon > 0.) & (psd_err.freq_time > 0), drop=True)
# return PSD-based score
return (1.0 - mean_psd_err / mean_psd_signal)
def test_dim_format(self, dim, window_correction, detrend):
"""Check that can deal with dim in various formats"""
data = xr.DataArray(
np.random.random([10]),
dims=[dim],
coords={dim: range(10)},
)
ps = xrft.power_spectrum(
data,
dim=dim,
window="hann",
window_correction=window_correction,
detrend=detrend,
)
ps = xrft.power_spectrum(
data,
dim=[dim],
window="hann",
window_correction=window_correction,
detrend=detrend,
)
def test_power_spectrum_dask():
"""Test the power spectrum function on dask data"""
N = 16
dim = ['x', 'y']
da = xr.DataArray(np.random.rand(2, N, N),
dims=['time', 'x', 'y'],
coords={
'time': range(2),
'x': range(N),
'y': range(N)
}).chunk({'time': 1})
ps = xrft.power_spectrum(da, dim=dim, density=False)
daft = xrft.dft(da, dim=['x', 'y'])
npt.assert_almost_equal(ps.values, (daft * np.conj(daft)).real.values)
ps = xrft.power_spectrum(da, dim=dim, window=True, detrend='constant')
daft = xrft.dft(da, dim=dim, window=True, detrend='constant')
coord = list(daft.coords)
test = (daft * np.conj(daft)).real / N**4
for i in dim:
test /= daft['freq_' + i + '_spacing']
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
def test_power_spectrum(self, dask):
"""Test the power spectrum function"""
N = 16
da = xr.DataArray(np.random.rand(2,N,N), dims=['time','y','x'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'y':range(N),'x':range(N)}
)
if dask:
da = da.chunk({'time': 1})
ps = xrft.power_spectrum(da, dim=['y','x'], window=True, density=False,
detrend='constant')
daft = xrft.dft(da, dim=['y','x'], detrend='constant', window=True)
npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
ps = xrft.power_spectrum(da, dim=['y'], real='x', window=True,
density=False, detrend='constant')
daft = xrft.dft(da, dim=['y'], real='x', detrend='constant', window=True)
npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
### Normalized
ps = xrft.power_spectrum(da, dim=['y','x'], window=True, detrend='constant')
daft = xrft.dft(da, dim=['y','x'], window=True, detrend='constant')
test = np.real(daft*np.conj(daft))/N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove least-square fit
ps = xrft.power_spectrum(da, dim=['y','x'],
window=True, density=False, detrend='linear'
)
daft = xrft.dft(da, dim=['y','x'], window=True, detrend='linear')
npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
def varweighted_mean_period(
da: Union[xr.Dataset, xr.DataArray],
dim: Union[str, List[str]] = "time",
**kwargs: Any,
) -> Union[xr.Dataset, xr.DataArray]:
r"""Calculate the variance weighted mean period of time series.
.. math::
P_{x} = \frac{\sum_k V(f_k,x)}{\sum_k f_k \cdot V(f_k,x)}
Args:
da: input data including dim.
dim: Name of time dimension.
for **kwargs see xrft.power_spectrum
Reference:
* Branstator, Grant, and Haiyan Teng. “Two Limits of Initial-Value
Decadal Predictability in a CGCM." Journal of Climate 23, no. 23
(August 27, 2010): 6292-6311. https://doi.org/10/bwq92h.
See also:
https://xrft.readthedocs.io/en/latest/api.html#xrft.xrft.power_spectrum
"""
if power_spectrum is None:
raise ImportError(
"xrft is not installed; see"
"https://xrft.readthedocs.io/en/latest/installation.html")
if isinstance(da, xr.Dataset):
return da.map(varweighted_mean_period, dim=dim, **kwargs)
da = da.fillna(0.0)
# dim should be list
if isinstance(dim, str):
dim = [dim]
assert isinstance(dim, list)
ps = power_spectrum(da, dim=dim, **kwargs)
# take pos
for d in dim:
ps = ps.where(ps[f"freq_{d}"] > 0)
# weighted average
vwmp = ps
for d in dim:
vwmp = vwmp.sum(f"freq_{d}") / (
(vwmp * vwmp[f"freq_{d}"]).sum(f"freq_{d}"))
for d in dim:
if f"freq_{d}_spacing" in vwmp.coords:
del vwmp.coords[f"freq_{d}_spacing"]
return vwmp
def varweighted_mean_period(da, dim="time", **kwargs):
"""Calculate the variance weighted mean period of time series based on
xrft.power_spectrum.
.. math::
P_{x} = \\frac{\\sum_k V(f_k,x)}{\\sum_k f_k \\cdot V(f_k,x)}
Args:
da (xarray object): input data including dim.
dim (optional str): Name of time dimension.
for **kwargs see xrft.power_spectrum
Reference:
* Branstator, Grant, and Haiyan Teng. “Two Limits of Initial-Value
Decadal Predictability in a CGCM." Journal of Climate 23, no. 23
(August 27, 2010): 6292-6311. https://doi.org/10/bwq92h.
See also:
https://xrft.readthedocs.io/en/latest/api.html#xrft.xrft.power_spectrum
"""
# set nans to 0
if isinstance(da, xr.Dataset):
raise ValueError("require xr.Dataset")
da = da.fillna(0.0)
# dim should be list
if isinstance(dim, str):
dim = [dim]
assert isinstance(dim, list)
ps = power_spectrum(da, dim=dim, **kwargs)
# take pos
for d in dim:
ps = ps.where(ps[f"freq_{d}"] > 0)
# weighted average
vwmp = ps
for d in dim:
vwmp = vwmp.sum(f"freq_{d}") / (
(vwmp * vwmp[f"freq_{d}"]).sum(f"freq_{d}"))
for d in dim:
del vwmp[f"freq_{d}_spacing"]
# try to copy coords
try:
vwmp = copy_coords_from_to(da.drop(dim), vwmp)
except ValueError:
warnings.warn("Couldn't keep coords.")
return vwmp
def hs_spectral_slope(var, data, ni, nf):
""" Computes spectral slope of the Hs wavenumber spectrum between 10km and 100km
"""
output = {}
hs = data[var][ni:nf, ni:nf]
hs_hat2 = xrft.power_spectrum(hs, dim=['x','y'], detrend='linear', window=True)
nk = hs_hat2.shape[-1]
hsy = hs_hat2.sum(dim='freq_x')[nk//2:]*hs_hat2.freq_x_spacing
hsx = hs_hat2.sum(dim='freq_y')[nk//2:]*hs_hat2.freq_y_spacing
q1=1/100e3
q2 = 1/10e3
ind = ((hsy.freq_y>=q1) & (hsy.freq_y<=q2)).values
sy, intercepty, r_valuey, p_valuey, std_erry = stats.linregress(
np.log10(hsy.freq_y[ind].values), np.log10(hsy[ind].values))
sx, _, _, _, _ = stats.linregress(np.log10(hsx.freq_x[ind].values), np.log10(hsx[ind].values))
output['hs_yslope'] = abs(sy)
output['hs_xslope'] = abs(sx)
return output
def test_chunks_to_segments():
N = 32
da = xr.DataArray(np.random.rand(N,N,N),
dims=['time','y','x'],
coords={'time':range(N),'y':range(N),'x':range(N)}
)
with pytest.raises(ValueError):
xrft.dft(da.chunk(chunks=((20,N,N),(N-20,N,N))), dim=['time'],
detrend='linear', chunks_to_segments=True)
ft = xrft.dft(da.chunk({'time':16}), dim=['time'], shift=False,
chunks_to_segments=True)
assert ft.dims == ('time_segment','freq_time','y','x')
data = da.chunk({'time':16}).data.reshape((2,16,N,N))
npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[1]),
decimal=7)
ft = xrft.dft(da.chunk({'y':16,'x':16}), dim=['y','x'], shift=False,
chunks_to_segments=True)
assert ft.dims == ('time','y_segment','freq_y','x_segment','freq_x')
data = da.chunk({'y':16,'x':16}).data.reshape((N,2,16,2,16))
npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[2,4]),
decimal=7)
ps = xrft.power_spectrum(da.chunk({'y':16,'x':16}), dim=['y','x'],
shift=False, density=False,
chunks_to_segments=True)
npt.assert_almost_equal(ps.values,
(ft*np.conj(ft)).real.values,
)
da2 = xr.DataArray(np.random.rand(N,N,N),
dims=['time','y','x'],
coords={'time':range(N),'y':range(N),'x':range(N)}
)
ft2 = xrft.dft(da2.chunk({'y':16,'x':16}), dim=['y','x'], shift=False,
chunks_to_segments=True)
cs = xrft.cross_spectrum(da.chunk({'y':16,'x':16}),
da2.chunk({'y':16,'x':16}),
dim=['y','x'], shift=False, density=False,
chunks_to_segments=True)
npt.assert_almost_equal(cs.values,
(ft*np.conj(ft2)).real.values,
)
def test_chunks_to_segments():
N = 32
da = xr.DataArray(np.random.rand(N,N,N),
dims=['time','y','x'],
coords={'time':range(N),'y':range(N),'x':range(N)}
)
with pytest.raises(ValueError):
xrft.dft(da.chunk(chunks=((20,N,N),(N-20,N,N))), dim=['time'],
detrend='linear', chunks_to_segments=True)
ft = xrft.dft(da.chunk({'time':16}), dim=['time'], shift=False,
chunks_to_segments=True)
assert ft.dims == ('time_segment','freq_time','y','x')
data = da.chunk({'time':16}).data.reshape((2,16,N,N))
npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[1]),
decimal=7)
ft = xrft.dft(da.chunk({'y':16,'x':16}), dim=['y','x'], shift=False,
chunks_to_segments=True)
assert ft.dims == ('time','y_segment','freq_y','x_segment','freq_x')
data = da.chunk({'y':16,'x':16}).data.reshape((N,2,16,2,16))
npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[2,4]),
decimal=7)
ps = xrft.power_spectrum(da.chunk({'y':16,'x':16}), dim=['y','x'],
shift=False, density=False,
chunks_to_segments=True)
npt.assert_almost_equal(ps.values,
(ft*np.conj(ft)).real.values,
)
da2 = xr.DataArray(np.random.rand(N,N,N),
dims=['time','y','x'],
coords={'time':range(N),'y':range(N),'x':range(N)}
)
ft2 = xrft.dft(da2.chunk({'y':16,'x':16}), dim=['y','x'], shift=False,
chunks_to_segments=True)
cs = xrft.cross_spectrum(da.chunk({'y':16,'x':16}),
da2.chunk({'y':16,'x':16}),
dim=['y','x'], shift=False, density=False,
chunks_to_segments=True)
npt.assert_almost_equal(cs.values,
(ft*np.conj(ft2)).real.values,
)
def test_chunks_to_segments():
N = 32
da = xr.DataArray(
np.random.rand(N, N, N),
dims=["time", "y", "x"],
coords={
"time": range(N),
"y": range(N),
"x": range(N)
},
)
with pytest.raises(ValueError):
xrft.dft(
da.chunk(chunks=((20, N, N), (N - 20, N, N))),
dim=["time"],
detrend="linear",
chunks_to_segments=True,
)
ft = xrft.dft(da.chunk({"time": 16}),
dim=["time"],
shift=False,
chunks_to_segments=True)
assert ft.dims == ("time_segment", "freq_time", "y", "x")
data = da.chunk({"time": 16}).data.reshape((2, 16, N, N))
npt.assert_almost_equal(ft.values,
dsar.fft.fftn(data, axes=[1]),
decimal=7)
ft = xrft.dft(
da.chunk({
"y": 16,
"x": 16
}),
dim=["y", "x"],
shift=False,
chunks_to_segments=True,
)
assert ft.dims == ("time", "y_segment", "freq_y", "x_segment", "freq_x")
data = da.chunk({"y": 16, "x": 16}).data.reshape((N, 2, 16, 2, 16))
npt.assert_almost_equal(ft.values,
dsar.fft.fftn(data, axes=[2, 4]),
decimal=7)
ps = xrft.power_spectrum(
da.chunk({
"y": 16,
"x": 16
}),
dim=["y", "x"],
shift=False,
density=False,
chunks_to_segments=True,
)
npt.assert_almost_equal(
ps.values,
(ft * np.conj(ft)).real.values,
)
da2 = xr.DataArray(
np.random.rand(N, N, N),
dims=["time", "y", "x"],
coords={
"time": range(N),
"y": range(N),
"x": range(N)
},
)
ft2 = xrft.dft(
da2.chunk({
"y": 16,
"x": 16
}),
dim=["y", "x"],
shift=False,
chunks_to_segments=True,
)
cs = xrft.cross_spectrum(
da.chunk({
"y": 16,
"x": 16
}),
da2.chunk({
"y": 16,
"x": 16
}),
dim=["y", "x"],
shift=False,
density=False,
chunks_to_segments=True,
)
npt.assert_almost_equal(
cs.values,
(ft * np.conj(ft2)).real.values,
)
else:
div=0.5
std = 0.02
ls='--'
utype = r'$\tilde E^\psi + \tilde E^\phi$'
hstype = r'$H_s^\psi$ + $H_s^\phi$'
cur_filename = "K%sA%sT%sS%s_cur.nc" %(slope, div, period, std)
cur_data = os.path.join(path, cur_filename)
dsc = data = xr.open_dataset(cur_data)
ni = np.where(dsc.x==200000)[0][0]
nf = np.where(dsc.x==600000)[0][0]
u_mix = dsc['ucur'][:, ni:nf, ni:nf]
v_mix = dsc['vcur'][:, ni:nf, ni:nf]
um_hat2 = xrft.power_spectrum(u_mix, dim=['x','y'], detrend='linear', window=True)
vm_hat2 = xrft.power_spectrum(v_mix, dim=['x','y'], detrend='linear', window=True)
nk = um_hat2.shape[-1]
ekem_hat2 = 0.5*(um_hat2 + vm_hat2)
ekem_y = ekem_hat2.sum(dim='freq_x')[:,nk//2:]*ekem_hat2.freq_x_spacing
N = ekem_y.seed.size
ky = ekem_y.freq_y.values*1e3
Em = ekem_y.mean(dim='seed').values*1e-3
Em_l, Em_u = spec_error(Em, N, ci=0.95)
hs_filename = "K%sA%sT%sS%s_hs.nc" %(slope, div, period, std)
hs_data = os.path.join(path, hs_filename)
dsh = xr.open_dataset(hs_data)
hs = dsh.hs[:, ni:nf, ni:nf]
hs_hat2 = xrft.power_spectrum(hs, dim=['x','y'], detrend='linear', window=True)
def test_dft_nocoords():
# Julius' example
# https://github.com/rabernat/xrft/issues/17
data = xr.DataArray(np.random.random([20,30,100]),dims=['time','lat','lon'])
dft = xrft.dft(data,dim=['time'])
ps = xrft.power_spectrum(data,dim=['time'])
def test_parseval(chunks_to_segments):
"""Test whether the Parseval's relation is satisfied."""
N = 16 # Must be divisible by n_segments (below)
da = xr.DataArray(np.random.rand(N, N),
dims=["x", "y"],
coords={
"x": range(N),
"y": range(N)
})
da2 = xr.DataArray(np.random.rand(N, N),
dims=["x", "y"],
coords={
"x": range(N),
"y": range(N)
})
if chunks_to_segments:
n_segments = 2
# Chunk da and da2 into n_segments
da = da.chunk({"x": N / n_segments, "y": N / n_segments})
da2 = da2.chunk({"x": N / n_segments, "y": N / n_segments})
else:
n_segments = 1
dim = da.dims
fftdim = [f"freq_{d}" for d in da.dims]
delta_x = []
for d in dim:
coord = da[d]
diff = np.diff(coord)
delta = diff[0]
delta_x.append(delta)
delta_xy = np.asarray(delta_x).prod() # Area of the spacings
### Test Parseval's theorem for power_spectrum with `window=False` and detrend=None
ps = xrft.power_spectrum(da, chunks_to_segments=chunks_to_segments)
# If n_segments > 1, use xrft._stack_chunks() to stack each segment along a new dimension
da_seg = xrft.xrft._stack_chunks(
da, dim).squeeze() if chunks_to_segments else da
da_prime = da_seg
# Check that the (rectangular) integral of the spectrum matches the energy
npt.assert_almost_equal(
(1 / delta_xy) * ps.mean(fftdim).values,
(da_prime**2).mean(dim).values,
decimal=5,
)
### Test Parseval's theorem for power_spectrum with `window=True` and detrend='constant'
# Note that applying a window weighting reduces the energy in a signal and we have to account
# for this reduction when testing Parseval's theorem.
ps = xrft.power_spectrum(da,
window=True,
detrend="constant",
chunks_to_segments=chunks_to_segments)
# If n_segments > 1, use xrft._stack_chunks() to stack each segment along a new dimension
da_seg = xrft.xrft._stack_chunks(
da, dim).squeeze() if chunks_to_segments else da
da_prime = da_seg - da_seg.mean(dim=dim)
# Generate the window weightings for each segment
window = xr.DataArray(
np.tile(
np.hanning(N / n_segments) *
np.hanning(N / n_segments)[:, np.newaxis],
(n_segments, n_segments),
),
dims=dim,
coords=da.coords,
)
# Check that the (rectangular) integral of the spectrum matches the windowed variance
npt.assert_almost_equal(
(1 / delta_xy) * ps.mean(fftdim).values,
((da_prime * window)**2).mean(dim).values,
decimal=5,
)
### Test Parseval's theorem for cross_spectrum with `window=True` and detrend='constant'
cs = xrft.cross_spectrum(da,
da2,
window=True,
detrend="constant",
chunks_to_segments=chunks_to_segments)
# If n_segments > 1, use xrft._stack_chunks() to stack each segment along a new dimension
da2_seg = xrft.xrft._stack_chunks(
da2, dim).squeeze() if chunks_to_segments else da2
da2_prime = da2_seg - da2_seg.mean(dim=dim)
# Check that the (rectangular) integral of the cross-spectrum matches the windowed co-variance
npt.assert_almost_equal(
(1 / delta_xy) * cs.mean(fftdim).values,
((da_prime * window) * (da2_prime * window)).mean(dim).values,
decimal=5,
)
### Test Parseval's theorem for a 3D case with `window=True` and `detrend='linear'`
if not chunks_to_segments:
d3d = xr.DataArray(
np.random.rand(N, N, N),
dims=["time", "y", "x"],
coords={
"time": range(N),
"y": range(N),
"x": range(N)
},
).chunk({"time": 1})
dim = ["x", "y"]
ps = xrft.power_spectrum(d3d, dim=dim, window=True, detrend="linear")
npt.assert_almost_equal(
(1 / delta_xy) * ps[0].values.mean(),
((xrft.detrend(d3d, dim, detrend_type="linear")[0].values *
window)**2).mean(),
decimal=5,
)
def psd_based_scores(ds_oi, ds_ref):
logging.info(' Compute PSD-based scores...')
with ProgressBar():
# Compute error = SSH_reconstruction - SSH_true
err = (ds_oi['sossheig'] - ds_ref['sossheig'])
err = err.chunk({
"lat": 1,
'time': err['time'].size,
'lon': err['lon'].size
})
# make time vector in days units
err['time'] = (err.time - err.time[0]) / numpy.timedelta64(1, 'D')
# Rechunk SSH_true
signal = ds_ref['sossheig'].chunk({
"lat": 1,
'time': ds_ref['time'].size,
'lon': ds_ref['lon'].size
})
# make time vector in days units
signal['time'] = (signal.time - signal.time[0]) / numpy.timedelta64(
1, 'D')
# Compute PSD_err and PSD_signal
psd_err = xrft.power_spectrum(err,
dim=['time', 'lon'],
detrend='linear',
window=True).compute()
psd_signal = xrft.power_spectrum(signal,
dim=['time', 'lon'],
detrend='linear',
window=True).compute()
# Averaged over latitude
mean_psd_signal = psd_signal.mean(dim='lat').where(
(psd_signal.freq_lon > 0.) & (psd_signal.freq_time > 0), drop=True)
mean_psd_err = psd_err.mean(dim='lat').where(
(psd_err.freq_lon > 0.) & (psd_err.freq_time > 0), drop=True)
# return PSD-based score
psd_based_score = (1.0 - mean_psd_err / mean_psd_signal)
# Find the key metrics: shortest temporal & spatial scales resolved based on the 0.5 contour criterion of the PSD_score
level = [0.5]
cs = plt.contour(1. / psd_based_score.freq_lon.values,
1. / psd_based_score.freq_time.values,
psd_based_score, level)
x05, y05 = cs.collections[0].get_paths()[0].vertices.T
plt.close()
shortest_spatial_wavelength_resolved = numpy.min(x05)
shortest_temporal_wavelength_resolved = numpy.min(y05)
logging.info(
' => Leaderboard Spectral score = %s (degree lon)',
numpy.round(shortest_spatial_wavelength_resolved, 2))
logging.info(
' => shortest temporal wavelength resolved = %s (days)',
numpy.round(shortest_temporal_wavelength_resolved, 2))
return (1.0 - mean_psd_err / mean_psd_signal), numpy.round(
shortest_spatial_wavelength_resolved,
2), numpy.round(shortest_temporal_wavelength_resolved, 2)
def test_power_spectrum(self, dask):
"""Test the power spectrum function"""
N = 16
da = xr.DataArray(
np.random.rand(N),
dims=["x"],
coords={
"x": range(N),
},
)
f_scipy, p_scipy = sps.periodogram(da.values,
window="rectangular",
return_onesided=True)
ps = xrft.power_spectrum(da, dim="x", real_dim="x", detrend="constant")
npt.assert_almost_equal(ps.values, p_scipy)
A = 20
fs = 1e4
n_segments = int(fs // 10)
fsig = 300
ii = int(fsig * n_segments // fs) # Freq index of fsig
tt = np.arange(fs) / fs
x = A * np.sin(2 * np.pi * fsig * tt)
for window_type in ["hann", "bartlett", "tukey", "flattop"]:
# see https://github.com/scipy/scipy/blob/master/scipy/signal/tests/test_spectral.py#L485
x_da = xr.DataArray(x, coords=[tt],
dims=["t"]).chunk({"t": n_segments})
ps = xrft.power_spectrum(
x_da,
dim="t",
window=window_type,
chunks_to_segments=True,
window_correction=True,
).mean("t_segment")
# Check the energy correction
npt.assert_allclose(
np.sqrt(np.trapz(ps.values, ps.freq_t.values)),
A * np.sqrt(2) / 2,
rtol=1e-3,
)
ps = xrft.power_spectrum(
x_da,
dim="t",
window=window_type,
chunks_to_segments=True,
scaling="spectrum",
window_correction=True,
).mean("t_segment")
# Check the amplitude correction
# The factor of 0.5 is there because we're checking the two-sided spectrum
npt.assert_allclose(ps.sel(freq_t=fsig), 0.5 * A**2 / 2.0)
da = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "y", "x"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"y": range(N),
"x": range(N),
},
)
if dask:
da = da.chunk({"time": 1})
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window="hann",
density=False,
detrend="constant")
daft = xrft.fft(da, dim=["y", "x"], detrend="constant", window="hann")
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
ps = xrft.power_spectrum(
da,
dim=["y"],
real_dim="x",
window="hann",
density=False,
detrend="constant",
)
daft = xrft.fft(da,
dim=["y"],
real_dim="x",
detrend="constant",
window="hann")
ps_test = np.real(daft * np.conj(daft))
f = np.full(ps_test.sizes["freq_x"], 2.0)
f[0], f[-1] = 1.0, 1.0
ps_test = ps_test * xr.DataArray(f, dims="freq_x")
npt.assert_almost_equal(ps.values, ps_test.values)
### Normalized
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window="hann",
detrend="constant")
daft = xrft.fft(da, dim=["y", "x"], window="hann", detrend="constant")
test = np.real(daft * np.conj(daft)) / N**4
dk = np.diff(np.fft.fftfreq(N, 1.0))[0]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
### Remove least-square fit
ps = xrft.power_spectrum(da,
dim=["y", "x"],
window="hann",
density=False,
detrend="linear")
daft = xrft.fft(da, dim=["y", "x"], window="hann", detrend="linear")
npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft)))
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
with pytest.raises(ValueError):
xrft.power_spectrum(da,
dim=["y", "x"],
window=None,
window_correction=True)
