def test_cross_spectrum(self, dask):
"""Test the cross spectrum function"""
N = 16
dim = ['x','y']
da = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'x':range(N), 'y':range(N)})
da2 = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'x':range(N), 'y':range(N)})
if dask:
da = da.chunk({'time': 1})
da2 = da2.chunk({'time': 1})
daft = xrft.dft(da, dim=dim, shift=True, detrend='constant',
window=True)
daft2 = xrft.dft(da2, dim=dim, shift=True, detrend='constant',
window=True)
cs = xrft.cross_spectrum(da, da2, dim=dim, window=True, density=False,
detrend='constant')
npt.assert_almost_equal(cs.values, np.real(daft*np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
cs = xrft.cross_spectrum(da, da2, dim=dim, shift=True, window=True,
detrend='constant')
test = (daft * np.conj(daft2)).real.values/N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 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_cross_spectrum(self, dask):
"""Test the cross spectrum function"""
N = 16
dim = ['x','y']
da = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'x':range(N), 'y':range(N)})
da2 = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
coords={'time':np.array(['2019-04-18', '2019-04-19'],
dtype='datetime64'),
'x':range(N), 'y':range(N)})
if dask:
da = da.chunk({'time': 1})
da2 = da2.chunk({'time': 1})
daft = xrft.dft(da, dim=dim, shift=True, detrend='constant',
window=True)
daft2 = xrft.dft(da2, dim=dim, shift=True, detrend='constant',
window=True)
cs = xrft.cross_spectrum(da, da2, dim=dim, window=True, density=False,
detrend='constant')
npt.assert_almost_equal(cs.values, np.real(daft*np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
cs = xrft.cross_spectrum(da, da2, dim=dim, shift=True, window=True,
detrend='constant')
test = (daft * np.conj(daft2)).real.values/N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
def test_cross_spectrum(self, dask):
"""Test the cross spectrum function"""
N = 16
dim = ["x", "y"]
da = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "x", "y"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"x": range(N),
"y": range(N),
},
)
da2 = xr.DataArray(
np.random.rand(2, N, N),
dims=["time", "x", "y"],
coords={
"time": np.array(["2019-04-18", "2019-04-19"],
dtype="datetime64"),
"x": range(N),
"y": range(N),
},
)
if dask:
da = da.chunk({"time": 1})
da2 = da2.chunk({"time": 1})
daft = xrft.dft(da,
dim=dim,
shift=True,
detrend="constant",
window=True)
daft2 = xrft.dft(da2,
dim=dim,
shift=True,
detrend="constant",
window=True)
cs = xrft.cross_spectrum(da,
da2,
dim=dim,
window=True,
density=False,
detrend="constant")
npt.assert_almost_equal(cs.values, np.real(daft * np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.0)
cs = xrft.cross_spectrum(da,
da2,
dim=dim,
shift=True,
window=True,
detrend="constant")
test = (daft * np.conj(daft2)).real.values / N**4
dk = np.diff(np.fft.fftfreq(N, 1.0))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).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_cross_spectrum():
"""Test the cross spectrum function"""
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)
})
cs = xrft.cross_spectrum(da,
da2,
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da, dim=None, shift=True, detrend='constant', window=True)
daft2 = xrft.dft(da2,
dim=None,
shift=True,
detrend='constant',
window=True)
npt.assert_almost_equal(cs.values, np.real(daft * np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
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_cross_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})
da2 = 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})
cs = xrft.cross_spectrum(da, da2, dim=dim, density=False)
daft = xrft.dft(da, dim=dim)
daft2 = xrft.dft(da2, dim=dim)
npt.assert_almost_equal(cs.values, (daft * np.conj(daft2)).real.values)
cs = xrft.cross_spectrum(da,
da2,
dim=dim,
shift=True,
window=True,
detrend='constant')
daft = xrft.dft(da, dim=dim, shift=True, window=True, detrend='constant')
daft2 = xrft.dft(da2, dim=dim, shift=True, window=True, detrend='constant')
coord = list(daft.coords)
test = (daft * np.conj(daft2)).real.values / N**4
# for i in dim:
# test /= daft['freq_' + i + '_spacing']
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
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_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_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 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,
)
