Added Linear detrend function for data that has more than one dimension (xgcm#15)

Takaya Uchida · rabernat · commit fa6466b9cf6c · 2017-07-03T14:38:28.000-04:00
* Added Linear detrend function for data that has more than one dimension

* Travis needs to install scipy

* Updated the requirements for xrft

* Added the shift argument for the test function of cross spectrum of dsar

* Let's see if this works

* updated appveyor.yml

* I think I have something working now. Added a wrapper function for detrending. After detrending though, the array no longer becomes a dask array

* Added one more test for Parseval

* Cleaned up the reshape argument in `detrend`

* Added test to improve codecov

* Added another test to improve codecov
diff --git a/.travis.yml b/.travis.yml
@@ -20,7 +20,7 @@ python:
   - 3.6
 
 env:
-  - CONDA_DEPS="pip flake8 pytest coverage numpy pandas xarray dask" PIP_DEPS="codecov pytest-cov"
+  - CONDA_DEPS="pip flake8 pytest coverage numpy scipy pandas xarray dask" PIP_DEPS="codecov pytest-cov"
 
 before_install:
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
diff --git a/appveyor.yml b/appveyor.yml
@@ -7,21 +7,11 @@ environment:
       PYTHON_ARCH: "32"
       MINICONDA: C:\Miniconda
 
-    - PYTHON: "C:\\Python33"
-      PYTHON_VERSION: "3.3.5"
+    - PYTHON: "C:\\Python36"
+      PYTHON_VERSION: "3.6.1"
       PYTHON_ARCH: "32"
       MINICONDA: C:\Miniconda3
 
-    - PYTHON: "C:\\Python34"
-      PYTHON_VERSION: "3.4.1"
-      PYTHON_ARCH: "32"
-      MINICONDA: C:\Miniconda3
-
-    - PYTHON: "C:\\Python35"
-      PYTHON_VERSION: "3.5.1"
-      PYTHON_ARCH: "32"
-      MINICONDA: C:\Miniconda35
-
 init:
   - "ECHO %PYTHON% %PYTHON_VERSION% %PYTHON_ARCH% %MINICONDA%"
 
diff --git a/xrft/tests/test_xrft.py b/xrft/tests/test_xrft.py
@@ -2,6 +2,8 @@
 import pandas as pd
 import xarray as xr
 import numpy.testing as npt
+import scipy.signal as sps
+import scipy.linalg as spl
 import pytest
 import xrft
 
@@ -15,6 +17,96 @@ def sample_data_1d():
     """Create one dimensional test DataArray."""
     pass
 
+def numpy_detrend(da):
+    """
+    Detrend a 2D field by subtracting out the least-square plane fit.
+
+    Parameters
+    ----------
+    da : `numpy.array`
+        The data to be detrended
+
+    Returns
+    -------
+    da : `numpy.array`
+        The detrended input data
+    """
+    N = da.shape
+
+    G = np.ones((N[0]*N[1],3))
+    for i in range(N[0]):
+        G[N[1]*i:N[1]*i+N[1], 1] = i+1
+        G[N[1]*i:N[1]*i+N[1], 2] = np.arange(1, N[1]+1)
+
+    d_obs = np.reshape(da.copy(), (N[0]*N[1],1))
+    m_est = np.dot(np.dot(spl.inv(np.dot(G.T, G)), G.T), d_obs)
+    d_est = np.dot(G, m_est)
+
+    lin_trend = np.reshape(d_est, N)
+
+    return da - lin_trend
+
+def test_detrend_1d():
+    N = 16
+    x = np.arange(N+1)
+    y = np.arange(N-1)
+    t = np.linspace(-int(N/2), int(N/2), N-6)
+    z = np.arange(int(N/2))
+    d4d = (t[:,np.newaxis,np.newaxis,np.newaxis]
+            + z[np.newaxis,:,np.newaxis,np.newaxis]
+            + y[np.newaxis,np.newaxis,:,np.newaxis]
+            + x[np.newaxis,np.newaxis,np.newaxis,:]
+          )
+    da4d = xr.DataArray(d4d, dims=['time','z','y','x'],
+                     coords={'time':range(len(t)),'z':range(len(z)),'y':range(len(y)),
+                             'x':range(len(x))}
+                     )
+    func = xrft._detrend_wrap(xrft.detrend2)
+    da = da4d.chunk({'time': 1})
+    da_prime = func(da.data, axes=[2]).compute()
+    npt.assert_allclose(da_prime[0,0], sps.detrend(d4d[0,0], axis=0))
+
+def test_detrend_2d():
+    N = 16
+    x = np.arange(N+1)
+    y = np.arange(N-1)
+    t = np.linspace(-int(N/2), int(N/2), N-6)
+    z = np.arange(int(N/2))
+    d4d = (t[:,np.newaxis,np.newaxis,np.newaxis]
+            + z[np.newaxis,:,np.newaxis,np.newaxis]
+            + y[np.newaxis,np.newaxis,:,np.newaxis]
+            + x[np.newaxis,np.newaxis,np.newaxis,:]
+          )
+    da4d = xr.DataArray(d4d, dims=['time','z','y','x'],
+                     coords={'time':range(len(t)),'z':range(len(z)),
+                             'y':range(len(y)),'x':range(len(x))}
+                       )
+    func = xrft._detrend_wrap(xrft.detrend2)
+    da = da4d.chunk({'time':1})
+    with pytest.raises(ValueError):
+        func(da.data, axes=[0]).compute()
+    da = da4d.chunk({'time':1, 'z':1})
+    with pytest.raises(ValueError):
+        func(da.data, axes=[1,2]).compute()
+    with pytest.raises(ValueError):
+        func(da.data, axes=[2,2]).compute()
+    da_prime = func(da.data, axes=[2,3]).compute()
+    npt.assert_allclose(da_prime[0,0], numpy_detrend(d4d[0,0]))
+
+    s = np.arange(2)
+    d5d = d4d[np.newaxis,:,:,:,:] + s[:,np.newaxis,np.newaxis,
+                                        np.newaxis,np.newaxis]
+    da5d = xr.DataArray(d5d, dims=['s','time','z','y','x'],
+                     coords={'s':range(len(s)),'time':range(len(t)),
+                             'z':range(len(z)),'y':range(len(y)),
+                             'x':range(len(x))}
+                       )
+    da = da5d.chunk({'time':1})
+    with pytest.raises(ValueError):
+        func(da.data).compute()
+    with pytest.raises(ValueError):
+        func(da.data, axes=[2,3,4]).compute()
+
 def test_dft_1d():
     """Test the discrete Fourier transform function on one-dimensional data."""
     Nx = 16
@@ -24,7 +116,7 @@ def test_dft_1d():
     da = xr.DataArray(np.random.rand(Nx), coords=[x], dims=['x'])
 
     # defaults with no keyword args
-    ft = xrft.dft(da)
+    ft = xrft.dft(da, detrend='constant')
     # check that the frequency dimension was created properly
     assert ft.dims == ('freq_x',)
     # check that the coords are correct
@@ -40,10 +132,16 @@ def test_dft_1d():
     npt.assert_allclose(ft_data_expected, ft.values, atol=1e-14)
 
     # redo without removing mean
-    ft = xrft.dft(da, remove_mean=False)
+    ft = xrft.dft(da)
     ft_data_expected = np.fft.fftshift(np.fft.fft(da))
     npt.assert_allclose(ft_data_expected, ft.values)
 
+    # redo with detrending linear least-square fit
+    ft = xrft.dft(da, detrend='linear')
+    da_prime = sps.detrend(da.values)
+    ft_data_expected = np.fft.fftshift(np.fft.fft(da_prime))
+    npt.assert_allclose(ft_data_expected, ft.values, atol=1e-14)
+
     # modify data to be non-evenly spaced
     da2 = da.copy()
     da2[-1] = np.nan
@@ -73,61 +171,88 @@ def test_dft_2d():
     da = xr.DataArray(np.random.rand(N,N), dims=['x','y'],
                     coords={'x':range(N),'y':range(N)}
                      )
-    ft = xrft.dft(da, shift=False, remove_mean=False)
+    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)
+    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))
 
+    with pytest.raises(ValueError):
+        xrft.dft(da, shift=False, window=True, detrend='linear')
+
 def test_dft_3d_dask():
     """Test the discrete Fourier transform on 3D dask array data"""
     N=16
-    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}
+    da = xr.DataArray(np.random.rand(N,N,N), dims=['time','x','y'],
+                      coords={'time':range(N),'x':range(N),
+                              'y':range(N)}
                      )
-    daft = xrft.dft(da, dim=['x','y'], shift=False, remove_mean=False)
-    assert hasattr(daft.data, 'dask')
-    npt.assert_almost_equal(daft.values, np.fft.fftn(da.values, axes=[1,2]))
+    daft = xrft.dft(da.chunk({'time': 1}), dim=['x','y'], shift=False)
+    # assert hasattr(daft.data, 'dask')
+    npt.assert_almost_equal(daft.values,
+                        np.fft.fftn(da.chunk({'time': 1}).values, axes=[1,2])
+                           )
 
     with pytest.raises(ValueError):
         xrft.dft(da.chunk({'time': 1, 'x': 1}), dim=['x'])
 
+    daft = xrft.dft(da.chunk({'x': 1}), dim=['time'],
+                    shift=False, detrend='linear')
+    # assert hasattr(daft.data, 'dask')
+    da_prime = sps.detrend(da.chunk({'x': 1}), axis=0)
+    npt.assert_almost_equal(daft.values,
+                        np.fft.fftn(da_prime, axes=[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)
+    ps = xrft.power_spectrum(da, window=True, density=False,
+                            detrend='constant')
     daft = xrft.dft(da,
-                    dim=None, shift=True, remove_mean=True,
+                    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)
-    daft = xrft.dft(da, window=True)
+    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
+                            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, dim=['y','x'], window=True)
+    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.)
 
@@ -143,8 +268,8 @@ def test_power_spectrum_dask():
     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)
-    daft = xrft.dft(da, dim=dim, window=True)
+    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:
@@ -161,12 +286,13 @@ def test_cross_spectrum():
     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)
+    cs = xrft.cross_spectrum(da, da2, window=True, density=False,
+                            detrend='constant')
     daft = xrft.dft(da,
-                    dim=None, shift=True, remove_mean=True,
+                    dim=None, shift=True, detrend='constant',
                     window=True)
     daft2 = xrft.dft(da2,
-                    dim=None, shift=True, remove_mean=True,
+                    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.)
@@ -188,13 +314,21 @@ def test_cross_spectrum_dask():
     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, window=True)
-    daft = xrft.dft(da, dim=dim, window=True)
-    daft2 = xrft.dft(da2, dim=dim, window=True)
+    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/N**4
-    for i in dim:
-        test /= daft['freq_' + i + '_spacing']
+    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.)
 
@@ -212,36 +346,19 @@ def test_parseval():
     for d in dim:
         coord = da[d]
         diff = np.diff(coord)
-        # if pd.core.common.is_timedelta64_dtype(diff):
-        #     # convert to seconds so we get hertz
-        #     diff = diff.astype('timedelta64[s]').astype('f8')
         delta = diff[0]
         delta_x.append(delta)
 
     window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
-    ps = xrft.power_spectrum(da, window=True)
+    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
                             )
 
-    # da = xr.DataArray(np.random.rand(5,2,20,30),
-    #               dims=['time', 'z', 'y', 'x'],
-    #               coords={'time': np.arange(5),'z':range(2),
-    #                     'y': np.arange(20),
-    #                     'x': np.arange(30)}).chunk({'z':1})
-    # dim = ['y','x']
-    # ps = xrft.power_spectrum(da, dim=dim, window=True)
-    # 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)
+    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()
@@ -250,6 +367,17 @@ def test_parseval():
                             ), 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 _synthetic_field(N, dL, amp, s):
     """
     Generate a synthetic series of size N by N
@@ -348,7 +476,7 @@ def test_isotropic_ps_slope(N=512, dL=1., amp=1e1, s=-3.):
     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, remove_mean=True,
+    iso_ps = xrft.isotropic_powerspectrum(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:],
diff --git a/xrft/version.py b/xrft/version.py
@@ -61,4 +61,4 @@
 MICRO = _version_micro
 VERSION = __version__
 PACKAGE_DATA = {'xrft': [pjoin('data', '*')]}
-REQUIRES = ["numpy", "xarray"]
+REQUIRES = ["numpy", "scipy", "xarray", "dask"]
diff --git a/xrft/xrft.py b/xrft/xrft.py
