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)