Skip to content

Commit

Permalink
Merge pull request #70 from xgcm/azi
Browse files Browse the repository at this point in the history 
for loop for azimuthal averaging in isotropic functions
  • Loading branch information
Takaya Uchida authored Apr 10, 2019
2 parents 8af5fc4 + 0a5c152 commit afdf053
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 55 deletions.
1 change: 1 addition & 0 deletions doc/conf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
15 changes: 13 additions & 2 deletions doc/whats-new.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>`__
66 changes: 34 additions & 32 deletions xrft/tests/test_xrft.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand All @@ -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.)

Expand All @@ -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
Expand Down
69 changes: 48 additions & 21 deletions xrft/xrft.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand Down Expand Up @@ -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)
Expand All @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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}

Expand All @@ -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
Expand Down Expand Up @@ -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}

Expand Down

0 comments on commit afdf053

Please sign in to comment.