isospectrum debugging / refactoring

.travis.yml

@@ -16,7 +16,6 @@ notifications:
 #     repo: rabernat/xrft
 
 python:
-  - 2.7
   - 3.6
 
 env:

xrft/tests/test_xrft.py

@@ -574,7 +574,7 @@ def test_parseval():
                             ), decimal=5
                            )
 
-def _synthetic_field(N, dL, amp, s):
+def synthetic_field(N, dL, amp, s):
     """
     Generate a synthetic series of size N by N
     with a spectral slope of s.
@@ -652,13 +652,72 @@ def _synthetic_field(N, dL, amp, s):
     theta = np.fft.ifft2(np.fft.ifftshift(F_theta))
     return np.real(theta)
 
+def synthetic_field_xr(N, dL, amp, s,
+                    other_dim_sizes=None, dim_order=True, 
+                    chunks=None):
+
+    theta = xr.DataArray(synthetic_field(N, dL, amp, s),
+                        dims=['y', 'x'],
+                        coords={'y':range(N), 'x':range(N)}
+                        )
+
+    if other_dim_sizes:
+        _da = xr.DataArray(np.ones(other_dim_sizes),
+                           dims=['d%d'%i for i in range(len(other_dim_sizes))])
+        if dim_order:
+            theta = theta + _da
+        else:
+            theta = _da + theta
+
+    if chunks:
+        theta = theta.chunk(chunks)
+
+    return theta
+
+def test_isotropize(N=512):
+    """Test the isotropization of a power spectrum."""
+
+    # generate synthetic 2D spectrum, isotropize and check values
+    dL, amp, s = 1., 1e1, -3.
+    dims = ['x','y']
+    fftdim = ['freq_x', 'freq_y']    
+    spacing_tol = 1e-3
+    nfactor = 4    
+    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)
+    # np data
+    theta = synthetic_field_xr(N, dL, amp, s)
+    _test_iso(theta)
+    # np with other dim
+    theta = synthetic_field_xr(N, dL, amp, s, 
+                                other_dim_sizes=[10],
+                                dim_order=True)
+    _test_iso(theta)
+    # da chunked, order 1
+    theta = synthetic_field_xr(N, dL, amp, s, 
+                                chunks={'y': None, 'x': None, 'd0': 2}, 
+                                other_dim_sizes=[10], 
+                                dim_order=True) 
+    _test_iso(theta)
+    # da chunked, order 2
+    theta = synthetic_field_xr(N, dL, amp, s, 
+                                chunks={'y': None, 'x': None, 'd0': 2}, 
+                                other_dim_sizes=[10], 
+                                dim_order=False) 
+    _test_iso(theta)
+
 def test_isotropic_ps_slope(N=512, dL=1., amp=1e1, s=-3.):
     """Test the spectral slope of isotropic power spectrum."""
 
-    theta = xr.DataArray(_synthetic_field(N, dL, amp, s),
+    theta = xr.DataArray(synthetic_field(N, dL, amp, s),
                         dims=['y', 'x'],
                         coords={'y':range(N), 'x':range(N)})
-    iso_ps = xrft.isotropic_powerspectrum(theta, detrend='constant',
+    iso_ps = xrft.isotropic_power_spectrum(theta, detrend='constant',
                                          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:],
@@ -674,13 +733,13 @@ def test_isotropic_ps():
                                           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'])
     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.)
+        xrft.isotropic_power_spectrum(da, dim=['z','y','x'])
+    iso_ps = xrft.isotropic_power_spectrum(da, dim=['y','x'])
+    npt.assert_almost_equal(
+            np.ma.masked_invalid(iso_ps.values[:,:,1:]).mask.sum(),
+            0.)
 
 def test_isotropic_cs():
     """Test isotropic cross spectrum"""
@@ -690,14 +749,14 @@ def test_isotropic_cs():
     da2 = xr.DataArray(np.random.rand(N,N),
                     dims=['y','x'], coords={'y':range(N),'x':range(N)})
 
-    iso_cs = xrft.isotropic_crossspectrum(da, da2, window=True)
+    iso_cs = xrft.isotropic_cross_spectrum(da, da2, window=True)
     npt.assert_almost_equal(np.ma.masked_invalid(iso_cs[1:]).mask.sum(), 0.)
 
     da2 = xr.DataArray(np.random.rand(N,N),
                     dims=['lat','lon'],
                     coords={'lat':range(N),'lon':range(N)})
     with pytest.raises(ValueError):
-        xrft.isotropic_crossspectrum(da, da2)
+        xrft.isotropic_cross_spectrum(da, da2)
 
     da = xr.DataArray(np.random.rand(2,5,5,16,32),
                   dims=['time','zz','z','y', 'x'],
@@ -711,16 +770,16 @@ def test_isotropic_cs():
                                           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'])
     da = da[:,0,:,:,:].drop(['zz'])
     da2 = da2[:,0,:,:,:].drop(['zz'])
     with pytest.raises(ValueError):
-        xrft.isotropic_crossspectrum(da, da2, dim=['z','y','x'])
+        xrft.isotropic_cross_spectrum(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.)
+    iso_cs = xrft.isotropic_cross_spectrum(da, da2, dim=['y','x'],
+                                         window=True)
+    npt.assert_almost_equal(
+            np.ma.masked_invalid(iso_cs.values[:,:,1:]).mask.sum(), 
+            0.)
 
 def test_spacing_tol(test_data_1d):
     da = test_data_1d

xrft/xrft.py

@@ -16,6 +16,8 @@
 
 __all__ = ["detrendn", "detrend_wrap",
            "dft","power_spectrum", "cross_spectrum", "cross_phase",
+           "isotropize",
+           "isotropic_power_spectrum", "isotropic_cross_spectrum",
            "isotropic_powerspectrum", "isotropic_crossspectrum",
            "fit_loglog"]
 
@@ -591,7 +593,12 @@ def cross_phase(da1, da2, spacing_tol=1e-3, dim=None, detrend=None,
     return cp
 
 
-def _azimuthal_wvnum(k, l, N, nfactor):
+def _radial_wvnum(k, l, N, nfactor):
+    """ Creates a radial wavenumber based on two horizontal wavenumbers
+    along with the appropriate index map
+    """
+
+    # compute target wavenumbers
     k = k.values
     l = l.values
     K = np.sqrt(k[np.newaxis,:]**2 + l[:,np.newaxis]**2)
@@ -601,40 +608,62 @@ def _azimuthal_wvnum(k, l, N, nfactor):
     else:
         ki = np.linspace(0., k.max(), nbins)
 
+    # compute bin index
     kidx = np.digitize(np.ravel(K), ki)
+    # compute number of points for each wavenumber
     area = np.bincount(kidx)
+    # compute the average radial wavenumber for each bin
+    kr = np.bincount(kidx, weights=K.ravel()) \
+            / np.ma.masked_where(area==0, area)
 
-    kr = np.bincount(kidx, weights=K.ravel()) / area
+    return ki, kr[1:-1]
 
-    return kidx, area, kr
 
-def _azimuthal_avg(kidx, f, area, kr):
-    """
-    Takes the azimuthal average of a given field.
+def isotropize(ps, fftdim, nfactor=4):
     """
+    Isotropize a 2D power spectrum or cross spectrum 
+    by taking an azimuthal average.
 
-    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("Arrays with more than 4 dimensions are not supported.")
+    .. math::
+        \text{iso}_{ps} = k_r N^{-1} \sum_{N} |\mathbb{F}(da')|^2
 
-    return iso
+    where :math:`N` is the number of azimuthal bins.
+
+    Parameters
+    ----------
+    ps : `xarray.DataArray`
+        The power spectrum or cross spectrum to be isotropized.
+    fftdim : list
+        The fft dimensions overwhich the isotropization must be performed.
+    nfactor : int, optional
+        Ratio of number of bins to take the azimuthal averaging with the
+        data size. Default is 4.        
+    """    
 
-def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True,
+    # compute radial wavenumber bins
+    k = ps[fftdim[1]]
+    l = ps[fftdim[0]]
+    N = [k.size, l.size]
+    ki, kr = _radial_wvnum(k, l, min(N), nfactor)
+
+    # average azimuthally
+    ps = ps.assign_coords(freq_r=np.sqrt(k**2+l**2))
+    iso_ps = (ps.groupby_bins('freq_r', bins=ki, labels=kr).mean()
+              .rename({'freq_r_bins': 'freq_r'})
+             )
+    return iso_ps * iso_ps.freq_r
+
+def isotropic_powerspectrum(*args, **kwargs): # pragma: no cover
+    """ 
+    Deprecated function. See isotropic_power_spectrum doc
+    """
+    import warnings
+    msg = "This function has been renamed and will disappear in the future."\
+          +" Please use isotropic_power_spectrum instead"
+    warnings.warn(msg, Warning)
+    return isotropic_power_spectrum(*args, **kwargs)
+
+def isotropic_power_spectrum(da, spacing_tol=1e-3, dim=None, shift=True,
                            detrend=None, density=True, window=False, nfactor=4):
     """
     Calculates the isotropic spectrum from the
@@ -646,6 +675,8 @@ def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True,
 
     where :math:`N` is the number of azimuthal bins.
 
+    Note: the method is not lazy does trigger computations.
+
     Parameters
     ----------
     da : `xarray.DataArray`
@@ -688,35 +719,20 @@ def isotropic_powerspectrum(da, spacing_tol=1e-3, dim=None, shift=True,
                        window=window)
 
     fftdim = ['freq_' + d for d in dim]
-    k = ps[fftdim[1]]
-    l = ps[fftdim[0]]
-
-    axis_num = [da.get_axis_num(d) for d in dim]
-    N = [da.shape[n] for n in axis_num]
-    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)
+    return isotropize(ps, fftdim, nfactor=nfactor)
 
-    k_coords = {'freq_r': kr}
-
-    newdims = [d for d in da.dims if d not in dim]
-    newdims.append('freq_r')
-
-    newcoords = {}
-    for d in newdims:
-        if d in da.coords:
-            newcoords[d] = da.coords[d].values
-        else:
-            newcoords[d] = k_coords[d]
-
-    return xr.DataArray(iso_ps, dims=newdims, coords=newcoords)
+def isotropic_crossspectrum(*args, **kwargs): # pragma: no cover
+    """ 
+    Deprecated function. See isotropic_cross_spectrum doc
+    """
+    import warnings
+    msg = "This function has been renamed and will disappear in the future."\
+          +" Please use isotropic_cross_spectrum instead"
+    warnings.warn(msg, Warning)
+    return isotropic_cross_spectrum(*args, **kwargs)
 
-def isotropic_crossspectrum(da1, da2, spacing_tol=1e-3,
+def isotropic_cross_spectrum(da1, da2, spacing_tol=1e-3,
                            dim=None, shift=True, detrend=None,
                            density=True, window=False, nfactor=4):
     """
@@ -728,6 +744,8 @@ def isotropic_crossspectrum(da1, da2, spacing_tol=1e-3,
         \text{iso}_{cs} = k_r N^{-1} \sum_{N} (\mathbb{F}(da1') {\mathbb{F}(da2')}^*)
 
     where :math:`N` is the number of azimuthal bins.
+
+    Note: the method is not lazy does trigger computations.
 
     Parameters
     ----------
@@ -776,33 +794,8 @@ def isotropic_crossspectrum(da1, da2, spacing_tol=1e-3,
                        window=window)
 
     fftdim = ['freq_' + d for d in dim]
-    k = cs[fftdim[1]]
-    l = cs[fftdim[0]]
-
-    axis_num = [da1.get_axis_num(d) for d in dim]
-    N = [da1.shape[n] for n in axis_num]
-    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}
-
-    newdims = [d for d in da1.dims if d not in dim]
-    newdims.append('freq_r')
-
-    newcoords = {}
-    for d in newdims:
-        if d in da1.coords:
-            newcoords[d] = da1.coords[d].values
-        else:
-            newcoords[d] = k_coords[d]
 
-    return xr.DataArray(iso_cs, dims=newdims, coords=newcoords)
+    return isotropize(cs, fftdim, nfactor=nfactor)
 
 def fit_loglog(x, y):
     """