Spectral (FFT) analysis

xr-scipy wraps some of scipy spectral analysis functions such as scipy.signal.spectrogram(), scipy.signal.csd() etc. For convenience, the xrscipy.signal namespace will be imported under the alias dsp

In [1]: import xrscipy.signal as dsp

To demonstrate the basic functionality, let’s create two simple example DataArray at a similar frequency but one with a frequency drift and some noise:

In [2]: time_ax = np.arange(0,100,0.01)

In [3]: sig_1 = xr.DataArray(np.sin(100 * time_ax) + np.random.rand(len(time_ax))*3,
   ...:                      coords=[("time", time_ax)], name='sig_1')
   ...: 

In [4]: sig_2 = xr.DataArray((np.cos(100 * time_ax) + np.random.rand(len(time_ax))*3 +
   ...:                       3*np.sin(30 * time_ax**1.3)),
   ...:                      coords=[("time", time_ax)], name='sig_2')
   ...: 

Power spectra

The spectrogram() function can be used directly on an xarray DataArray object. The returned object is again an xarray.DataArray object.

In [5]: spec_1 = dsp.spectrogram(sig_1)

In [6]: spec_2 = dsp.spectrogram(sig_2)

In [7]: spec_2
Out[7]: 
<xarray.DataArray 'spectrogram_sig_2' (frequency: 129, time: 77)>
array([[2.44424801e-05, 5.22947169e-03, 2.89224670e-04, ...,
        1.49873467e-03, 7.65608700e-04, 9.59582636e-05],
       [1.75153585e-02, 1.89154990e-02, 7.80560915e-03, ...,
        1.44778153e-03, 1.09230319e-02, 9.86477734e-03],
       [2.67981444e-02, 1.72320528e-03, 1.77923215e-02, ...,
        3.47197400e-03, 7.94489056e-05, 1.24722011e-02],
       ...,
       [7.31455285e-03, 1.98518741e-02, 5.07390687e-03, ...,
        6.44165409e-03, 5.43732551e-03, 7.00189408e-03],
       [3.07786421e-03, 1.69232295e-03, 1.04356585e-03, ...,
        2.04653030e-02, 5.20056837e-04, 1.61887611e-02],
       [2.69532456e-03, 1.81679619e-05, 2.76618829e-03, ...,
        5.93410480e-03, 6.43149358e-03, 1.45961774e-02]])
Coordinates:
  * time       (time) float64 1.28 2.56 3.84 5.12 6.4 ... 94.72 96.0 97.28 98.56
  * frequency  (frequency) float64 0.0 0.3906 0.7812 1.172 ... 49.22 49.61 50.0

The frequency dimension coords are based on the transformed dimension (time in this case) coords sampling (i.e. inverse units). When the signal is 1D, the dimension does not have to be provided.

In [8]: norm = plt.matplotlib.colors.LogNorm()

In [9]: plt.subplot(211)
Out[9]: <matplotlib.axes._subplots.AxesSubplot at 0x7fdc21db2d10>

In [10]: spec_1.plot(norm=norm)
Out[10]: <matplotlib.collections.QuadMesh at 0x7fdc21a20250>

In [11]: plt.subplot(212)
Out[11]: <matplotlib.axes._subplots.AxesSubplot at 0x7fdc21f5c190>

In [12]: spec_2.plot(norm=norm)
Out[12]: <matplotlib.collections.QuadMesh at 0x7fdc21d1b550>

In [13]: plt.show()
_images/spectrograms.png

These routines calculate the FFT on segments of the signal of a length controlled by nperseg and nfft parameters. The routines here offer a convenience parameter seglen which makes it possible to specify the segment length in the units of the transformed dimension’s coords. If seglen is specified, nperseg is then calculated from it and nfft is set using scipy.fftpack.next_fast_len (or to closest higher power of 2). A desired frequency resolution spacing df can be achieved by specifying seglen=1/df.

Another convenience parameter is overlap_ratio which calculates the noverlap parameter (by how many points the segments overlap) as noverlap = np.rint(overlap_ratio * nperseg)

For example, these parameters calculate the spectrogram with a higher frequency resolution and try to make for the longer segments by overlapping them by 75%.

In [14]: dsp.spectrogram(sig_1, seglen=1, overlap_ratio=0.75)
Out[14]: 
<xarray.DataArray 'spectrogram_sig_1' (frequency: 51, time: 397)>
array([[3.99319244e-04, 5.72385049e-06, 4.63213189e-04, ...,
        1.30700319e-03, 8.42935876e-04, 2.13630973e-03],
       [2.36519351e-03, 2.63675999e-03, 2.70759396e-03, ...,
        1.35248503e-02, 1.52467384e-02, 1.52477065e-02],
       [1.14508286e-03, 6.57120911e-03, 1.59951326e-02, ...,
        2.59556948e-02, 2.86589719e-02, 2.80344654e-02],
       ...,
       [3.22485304e-02, 5.98876777e-03, 4.75008719e-03, ...,
        1.27105913e-02, 1.10435949e-02, 2.60625766e-03],
       [6.48727646e-03, 4.24654407e-03, 5.71875070e-03, ...,
        3.54197551e-03, 4.97810155e-02, 5.84757962e-02],
       [7.82088430e-04, 3.02992866e-04, 8.29188527e-03, ...,
        1.46543043e-02, 4.78417261e-02, 3.80860346e-02]])
Coordinates:
  * time       (time) float64 0.5 0.75 1.0 1.25 1.5 ... 98.75 99.0 99.25 99.5
  * frequency  (frequency) float64 0.0 1.0 2.0 3.0 4.0 ... 47.0 48.0 49.0 50.0

All the functions can be calculated on N-dimensional signals if the dimension is provided. Here the power spectral density (PSD) \(P_{xx}\) is calculated using Welch’s method (i.e. time average of the spectrogram) is shown

In [15]: sig_2D = xr.concat([sig_1,sig_2], dim="sigs")

In [16]: psd_2D = dsp.psd(sig_2D, dim="time")
In [17]: psd_2D.plot.line(x='frequency')
Out[17]: 
[<matplotlib.lines.Line2D at 0x7fdc207d7790>,
 <matplotlib.lines.Line2D at 0x7fdc207d7b50>]

In [18]: plt.loglog()
Out[18]: []

In [19]: plt.grid(which='both')

In [20]: plt.show()
_images/psd.png

Cross-coherence and correlation

The same windowed FFT approach is also used to calculate the cross-spectral density \(P_{xy}\) (using xrscipy.signal.csd()) and from it coherency \(\gamma\) as

\[\gamma = \frac{\langle P_{xy}\rangle}{\sqrt{\langle P_{xx} \rangle \langle P_{yy} \rangle}}\]

where \(\langle \dots \rangle\) is an average over the FFT windows, i.e. ergodicity is assumed.

In [21]: coher_12 = dsp.coherence(sig_1, sig_2)

In [22]: coher_12[:10]
Out[22]: 
<xarray.DataArray 'coherence_sig_1_sig_2' (frequency: 10)>
array([ 0.15445836+0.j        ,  0.09626349+0.14572021j,
        0.01668929+0.1171455j , -0.01160877+0.0149668j ,
       -0.01835304-0.09690598j,  0.05856481-0.23405385j,
       -0.10611867-0.03007782j, -0.08271153-0.10426975j,
        0.01914803-0.07631421j,  0.0691541 -0.03660832j])
Coordinates:
  * frequency  (frequency) float64 0.0 0.3906 0.7812 1.172 ... 2.734 3.125 3.516

The returned \(\gamma\) DataArray is complex (because so is \(P_{xy}\)) and the modulus is what is more commonly called coherence and the angle is the phase shift.

In [23]: coh = np.abs(coher_12)

In [24]: xphase = xr.apply_ufunc(np.angle, coher_12) / np.pi

In [25]: fig, axs = plt.subplots(2, 1, sharex=True)

In [26]: coh.plot(ax=axs[0])
Out[26]: [<matplotlib.lines.Line2D at 0x7fdc207b8c10>]

In [27]: xphase.where(coh > 0.6).plot.line('o--', ax=axs[1])
Out[27]: [<matplotlib.lines.Line2D at 0x7fdc207c1910>]

In [28]: axs[1].set(yticks=[-1, -0.5, 0, 0.5, 1]);

In [29]: plt.show()
_images/coher.png

In the future more convenient wrappers returning the coherence magnitude and cross-phase might be developed.

The cross-correlation is calculated similarly as \(\gamma\), but with \(\mathcal{F}^{-1} [\langle P_*\rangle ]\), i.e. in the inverse-FFT domain. The lag coordinates are the inverse of the frequency coordinates.

In [30]: xcorr_12 = dsp.xcorrelation(sig_1, sig_2)

In [31]: xcorr_12.loc[-0.1:0.1].plot()
Out[31]: [<matplotlib.lines.Line2D at 0x7fdc20811cd0>]

In [32]: plt.grid()

In [33]: plt.show()
_images/xcorr.png

A partially averaged counterpart to coherence() is coherogram() which uses a running average over nrolling FFT windows.