def test_isotropic_ps():
"""Test data with extra coordinates"""
da = xr.DataArray(np.random.rand(2, 5, 5, 16, 32),
dims=['time', 'zz', 'z', 'y', 'x'],
coords={
'time':
np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'zz':
np.arange(5),
'z':
np.arange(5),
'y':
np.arange(16),
'x':
np.arange(32)
})
with pytest.raises(ValueError):
xrft.isotropic_powerspectrum(da, dim=['y', 'x'])
da = da[:, 0, :, :, :].drop(['zz'])
with pytest.raises(ValueError):
xrft.isotropic_powerspectrum(da, dim=['z', 'y', 'x'])
iso_ps = xrft.isotropic_powerspectrum(da, dim=['y', 'x']).values
npt.assert_almost_equal(
np.ma.masked_invalid(iso_ps[:, :, 1:]).mask.sum(), 0.)
def eke_spectral_slope(u, v):
""" Computes spectral slope of the KE wavenumber spectrum between 10km and 100km
"""
u_iso = xrft.isotropic_powerspectrum(u, dim=['x','y'], detrend='linear', window=True)
v_iso = xrft.isotropic_powerspectrum(v, dim=['x','y'], detrend='linear', window=True)
Ei = 0.5*(u_iso + v_iso)
df = np.diff(Ei.freq_r[1:]).mean()
diffusivity = np.sum(Ei.freq_r*Ei*df)
q1=1/100e3
q2 = 1/10e3
ind = ((Ei.freq_r>=q1) & (Ei.freq_r<=q2)).values
s, _, _, _, _ = stats.linregress(np.log10(Ei.freq_r[ind].values), np.log10(Ei[ind].values))
return abs(s), float(diffusivity)
def test_isotropic_ps():
"""Test data with extra coordinates"""
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)
})
with pytest.raises(ValueError):
xrft.isotropic_powerspectrum(da, dim=['y', 'x'])
iso_ps = np.zeros((5, int(20 / 4) + 1))
for t in range(5):
iso_ps[t] = xrft.isotropic_powerspectrum(da[t], dim=['y', 'x']).values
npt.assert_almost_equal(np.ma.masked_invalid(iso_ps[:, 1:]).mask.sum(), 0.)
def test_isotropic_ps():
"""Test data with extra coordinates"""
da = xr.DataArray(np.random.rand(2,5,5,16,32),
dims=['time','zz','z','y', 'x'],
coords={'time': np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'zz': np.arange(5), 'z': np.arange(5),
'y': np.arange(16), 'x': np.arange(32)})
with pytest.raises(ValueError):
xrft.isotropic_powerspectrum(da, dim=['y','x'])
da = da[:,0,:,:,:].drop(['zz'])
with pytest.raises(ValueError):
xrft.isotropic_powerspectrum(da, dim=['z','y','x'])
iso_ps = xrft.isotropic_powerspectrum(da, dim=['y','x']).values
npt.assert_almost_equal(np.ma.masked_invalid(iso_ps[:,:,1:]).mask.sum(), 0.)
def test_isotropic_ps_slope(N=512, dL=1., amp=1e1, s=-3.):
"""Test the spectral slope of isotropic power spectrum."""
theta = xr.DataArray(_synthetic_field(N, dL, amp, s),
dims=['y', 'x'],
coords={'y':range(N), 'x':range(N)})
iso_ps = xrft.isotropic_powerspectrum(theta, detrend='constant',
density=True)
npt.assert_almost_equal(np.ma.masked_invalid(iso_ps[1:]).mask.sum(), 0.)
y_fit, a, b = xrft.fit_loglog(iso_ps.freq_r.values[4:],
iso_ps.values[4:])
npt.assert_allclose(a, s, atol=.1)
def test_isotropic_ps_slope(N=512, dL=1., amp=1e1, s=-3.):
"""Test the spectral slope of isotropic power spectrum."""
theta = xr.DataArray(_synthetic_field(N, dL, amp, s),
dims=['y', 'x'],
coords={
'y': range(N),
'x': range(N)
})
iso_ps = xrft.isotropic_powerspectrum(theta,
detrend='constant',
density=True)
npt.assert_almost_equal(np.ma.masked_invalid(iso_ps[1:]).mask.sum(), 0.)
y_fit, a, b = xrft.fit_loglog(iso_ps.freq_r.values[4:], iso_ps.values[4:])
npt.assert_allclose(a, s, atol=.1)
alpha=0.5)
x = np.arange(0, L, dx) * 1e-3
y = np.arange(0, L, dx) * 1e-3
u_psi = xr.DataArray(u_psi, dims=('y', 'x'), coords={'y': y, 'x': x})
v_psi = xr.DataArray(v_psi, dims=('y', 'x'), coords={'y': y, 'x': x})
u_phi = xr.DataArray(u_phi, dims=('y', 'x'), coords={'y': y, 'x': x})
v_phi = xr.DataArray(v_phi, dims=('y', 'x'), coords={'y': y, 'x': x})
u_mix = xr.DataArray(u_mix, dims=('y', 'x'), coords={'y': y, 'x': x})
v_mix = xr.DataArray(v_mix, dims=('y', 'x'), coords={'y': y, 'x': x})
Us = (u_psi**2 + v_psi**2)**.5
Up = (u_phi**2 + v_phi**2)**.5
Um = (u_mix**2 + v_mix**2)**.5
um_iso = xrft.isotropic_powerspectrum(u_mix,
dim=['x', 'y'],
detrend='linear',
window=True)
vm_iso = xrft.isotropic_powerspectrum(v_mix,
dim=['x', 'y'],
detrend='linear',
window=True)
Ei_mix = 0.5 * (um_iso + vm_iso)
axpsi = axes_rot[i]
axs = axes_s[i]
axphi = axes_div[i]
axs.loglog(Ei_mix.freq_r,
Ei_mix,
color='k',
label='q$^{-%s}$' % st['%s' % slope])