for loop for azimuthal averaging in isotropic functions

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'

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 <https://github.com/roxyboy>`_ and
   `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 - Implemented ``cross_phase`` function to calculate the phase difference between two signals as a function of frequency.
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+
+- Allowed ``isotropic_powerspectrum`` function to support arrays with up to four dimensions. (:issue:`9`)
+  By `Takaya Uchida <https://github.com/roxyboy>`_
+
+.. 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 <http://www.python3statement.org/>`__
+  - `Tips on porting to Python 3 <https://docs.python.org/3/howto/pyporting.html>`__

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

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}