def test_idft_dft():
"""
Testing idft(dft(s)) == s
"""
N = 40
dx = np.random.rand()
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="x",
coords={
"x":
dx * (np.arange(-N // 2, -N // 2 + N) +
np.random.randint(-N // 2, N // 2))
},
)
FTs = xrft.dft(s, true_phase=True, true_amplitude=True)
mean_lag = float(s["x"][{
"x": s.sizes["x"] // 2
}]) # lag ensure IFTs to be on the same coordinate range than s
# lag is set manually
IFTs = xrft.idft(FTs,
shift=True,
true_phase=True,
true_amplitude=True,
lag=mean_lag)
xrt.assert_allclose(s, IFTs)
# lag is set automatically
IFTs = xrft.idft(FTs, shift=True, true_phase=True, true_amplitude=True)
xrt.assert_allclose(s, IFTs)
def test_idft_centered_coordinates():
"""error should be raised if coordinates are not centered on zero in idft"""
N = 20
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="freq_x",
coords={"freq_x": np.arange(-N // 2, N // 2) + 2},
)
with pytest.raises(ValueError):
xrft.idft(s)
def test_constant_coordinates():
"""error should be raised if coordinates are constant"""
N = 20
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="freq_x",
coords={"freq_x": np.zeros(N)},
)
with pytest.raises(ValueError):
xrft.dft(s)
with pytest.raises(ValueError):
xrft.idft(s)
def fft_filter(data: xr.DataArray, stops):
"""
Applies a brick wall filter at region in ``stops`` in the Fourier transform of data. Use with care.
:param data:
:param stops:
:return:
"""
# This won't tolerate inverse inverse filtering ;)
kdata = xrft.dft(data)
# this will produce a fair amount of ringing, Scipy isn't very clear about how to butterworth filter in Nd
for stop in stops:
kstop = {'freq_' + k if 'freq_' not in k else k: v for k, v in stop.items()} # be nice
kdata.loc[kstop] = 0
kkdata = xrft.idft(kdata)
kkdata.values = np.real(kkdata.values)
kkdata.values = kkdata.values - np.min(kkdata.values) + np.mean(data.values)
filtered_arr = xr.DataArray(
kkdata,
data.coords,
data.dims,
attrs=data.attrs.copy()
)
if 'id' in filtered_arr:
del filtered_arr.attrs['id']
provenance(filtered_arr, data, {
'what': 'Apply a filter in frequency space by brick walling coordinate regions.',
'by': 'fft_filter',
'stops': stops,
})
return filtered_arr