diff --git a/doc/conf.py b/doc/conf.py index 817ad1ce..6535a8e2 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -42,6 +42,7 @@ 'IPython.sphinxext.ipython_directive', 'IPython.sphinxext.ipython_console_highlighting'] +# apidoc_module_dir = '../xrft' # never execute notebooks: avoids lots of expensive imports on rtd # https://nbsphinx.readthedocs.io/en/0.2.14/never-execute.html nbsphinx_execute = 'never' diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c546dd2c..9e9dd6af 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -5,15 +5,26 @@ What's New .. _whats-new.0.2.0: -v0.2.0 (8 April 2019) +v0.2.0 (10 April 2019) ---------------------- Enhancements ~~~~~~~~~~~~ -- Allowed ``dft`` and ``power_spectrum`` function to support real Fourier transforms. (:issue:`57`) +- Allowed ``dft`` and ``power_spectrum`` functions to support real Fourier transforms. (:issue:`57`) By `Takaya Uchida `_ and `Tom Nicholas `_. - Implemented ``cross_phase`` function to calculate the phase difference between two signals as a function of frequency. By `Tom Nicholas `_. + +- Allowed ``isotropic_powerspectrum`` function to support arrays with up to four dimensions. (:issue:`9`) + By `Takaya Uchida `_ + +.. warning:: + + This is the last xrft release that will support Python 2.7. Future releases + will be Python 3 only. For the more details, see: + + - `Python 3 Statement `__ + - `Tips on porting to Python 3 `__ diff --git a/xrft/tests/test_xrft.py b/xrft/tests/test_xrft.py index 93afa7fd..974f1799 100644 --- a/xrft/tests/test_xrft.py +++ b/xrft/tests/test_xrft.py @@ -604,25 +604,28 @@ def test_isotropic_ps_slope(N=512, dL=1., amp=1e1, s=-3.): dims=['y', 'x'], coords={'y':range(N), 'x':range(N)}) iso_ps = xrft.isotropic_powerspectrum(theta, detrend='constant', - density=True) + 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:]) + iso_ps.values[4:]) npt.assert_allclose(a, s, atol=.1) 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)}) + 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']) - 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.) + 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_cs(): """Test isotropic cross spectrum""" @@ -632,14 +635,6 @@ def test_isotropic_cs(): da2 = xr.DataArray(np.random.rand(N,N), dims=['y','x'], coords={'y':range(N),'x':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) - iso_cs = xrft.isotropic_crossspectrum(da, da2, window=True) npt.assert_almost_equal(np.ma.masked_invalid(iso_cs[1:]).mask.sum(), 0.) @@ -649,21 +644,28 @@ def test_isotropic_cs(): with pytest.raises(ValueError): xrft.isotropic_crossspectrum(da, da2) - 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)}).chunk({'time':1}) - da2 = 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)}).chunk({'time':1}) + 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)}) + da2 = 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_crossspectrum(da, da2, dim=['y','x'], window=True) - iso_cs = np.zeros((5,int(20/4)+1)) - for t in range(5): - iso_cs[t] = xrft.isotropic_crossspectrum(da[t], da2[t], dim=['y','x'], - window=True).values - npt.assert_almost_equal(np.ma.masked_invalid(iso_cs[:,1:]).mask.sum(), 0.) + xrft.isotropic_crossspectrum(da, da2, dim=['y','x']) + da = da[:,0,:,:,:].drop(['zz']) + da2 = da2[:,0,:,:,:].drop(['zz']) + with pytest.raises(ValueError): + xrft.isotropic_crossspectrum(da, da2, dim=['z','y','x']) + + iso_cs = xrft.isotropic_crossspectrum(da, da2, dim=['y','x'], + window=True).values + npt.assert_almost_equal(np.ma.masked_invalid(iso_cs[:,:,1:]).mask.sum(), 0.) def test_spacing_tol(test_data_1d): da = test_data_1d diff --git a/xrft/xrft.py b/xrft/xrft.py index 3780671d..0915b26e 100644 --- a/xrft/xrft.py +++ b/xrft/xrft.py @@ -511,7 +511,7 @@ def cross_phase(da1, da2, spacing_tol=1e-3, dim=None, detrend=None, .. math:: da1' = da1 - \overline{da1};\ \ da2' = da2 - \overline{da2} .. math:: - cp = \text{Arg} \mathbb{F}(da1')^* {\mathbb{F}(da2')} + cp = \text{Arg} [\mathbb{F}(da1')^*, \mathbb{F}(da2')] Parameters ---------- @@ -573,14 +573,10 @@ def cross_phase(da1, da2, spacing_tol=1e-3, dim=None, detrend=None, return cp -def _azimuthal_avg(k, l, f, fftdim, N, nfactor): - """ - Takes the azimuthal average of a given field. - """ +def _azimuthal_wvnum(k, l, N, nfactor): k = k.values l = l.values - kk, ll = np.meshgrid(k, l) - K = np.sqrt(kk**2 + ll**2) + K = np.sqrt(k[np.newaxis,:]**2 + l[:,np.newaxis]**2) nbins = int(N/nfactor) if k.max() > l.max(): ki = np.linspace(0., l.max(), nbins) @@ -592,16 +588,33 @@ def _azimuthal_avg(k, l, f, fftdim, N, nfactor): kr = np.bincount(kidx, weights=K.ravel()) / area - if f.ndim == 2: - iso_f = np.ma.masked_invalid(np.bincount(kidx, - weights=f.data.ravel()) - / area) * kr + return kidx, area, kr + +def _azimuthal_avg(kidx, f, area, kr): + """ + Takes the azimuthal average of a given field. + """ + + iso_f = np.ma.masked_invalid(np.bincount(kidx, weights=f) + / area) * kr + + return iso_f + +def _azi_wrapper(M, kidx, f, area, kr): + iso = np.zeros(M) + if len(M) == 1: + iso = _azimuthal_avg(kidx, f, area, kr) + elif len(M) == 2: + for j in range(M[0]): + iso[j] = _azimuthal_avg(kidx, f[j], area, kr) + elif len(M) == 3: + for j in range(M[0]): + for i in range(M[1]): + iso[j,i] = _azimuthal_avg(kidx, f[j,i], area, kr) else: - raise ValueError('The data has too many or few dimensions. ' - 'The input should only have the two dimensions ' - 'to take the azimuthal averaging over.') + raise ValueError("Arrays with more than 4 dimensions are not supported.") - return kr, iso_f + return iso def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True, detrend=None, density=True, window=False, nfactor=4): @@ -611,7 +624,8 @@ def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True, azimuthal average. .. math:: - iso_{ps} = k_r N^{-1} \sum_{N} |\mathbb{F}(da')|^2 + \text{iso}_{ps} = k_r N^{-1} \sum_{N} |\mathbb{F}(da')|^2 + where :math:`N` is the number of azimuthal bins. Parameters @@ -661,8 +675,14 @@ def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True, axis_num = [da.get_axis_num(d) for d in dim] N = [da.shape[n] for n in axis_num] - kr, iso_ps = _azimuthal_avg(k, l, ps, fftdim, - np.asarray(N).min(), nfactor) + M = [da.shape[n] for n in [da.get_axis_num(d) for d in da.dims] if n not in axis_num] + shape = list(M) + + kidx, area, kr = _azimuthal_wvnum(k, l, np.asarray(N).min(), nfactor) + M.append(len(kr)) + shape.append(np.prod(N)) + f = ps.data.reshape(shape) + iso_ps = _azi_wrapper(M, kidx, f, area, kr) k_coords = {'freq_r': kr} @@ -687,7 +707,8 @@ def isotropic_crossspectrum(da1, da2, spacing_tol=1e-3, azimuthal average. .. math:: - iso_{cs} = k_r N^{-1} \sum_{N} (\mathbb{F}(da1') {\mathbb{F}(da2')}^*) + \text{iso}_{cs} = k_r N^{-1} \sum_{N} (\mathbb{F}(da1') {\mathbb{F}(da2')}^*) + where :math:`N` is the number of azimuthal bins. Parameters @@ -742,8 +763,14 @@ def isotropic_crossspectrum(da1, da2, spacing_tol=1e-3, axis_num = [da1.get_axis_num(d) for d in dim] N = [da1.shape[n] for n in axis_num] - kr, iso_cs = _azimuthal_avg(k, l, cs, fftdim, - np.asarray(N).min(), nfactor) + M = [da1.shape[n] for n in [da1.get_axis_num(d) for d in da1.dims] if n not in axis_num] + shape = list(M) + + kidx, area, kr = _azimuthal_wvnum(k, l, np.asarray(N).min(), nfactor) + M.append(len(kr)) + shape.append(np.prod(N)) + f = cs.data.reshape(shape) + iso_cs = _azi_wrapper(M, kidx, f, area, kr) k_coords = {'freq_r': kr}