# Thomson's "Multitaper" Estimator

This notebook is a demo & test of new multitaper estimator code.

**TODO**: the jackknife is not working in spawn mode


```python
import multiprocessing as mp
try:
    mp.set_start_method('spawn')
except:
    pass
```


```python
%matplotlib inline
```


```python
# basic stuff
import numpy as np
from scipy.signal import periodogram, detrend
import scipy.stats.distributions as dists

# Ecog stuff
import ecoglib.estimation.multitaper as mt
from ecogdata.devices.data_util import load_experiment_auto
from ecogdata.trigger_fun import extract_epochs
from ecoglib.vis.tile_images import tile_traces_1ax
from ecoglib.vis.ani import animate_frames
from ecoglib.vis.plot_util import filled_interval

# Other plotting stuff
import matplotlib.pyplot as plt
from matplotlib import rcParams
import seaborn as sns
sns.set_context('notebook')
rcParams['figure.figsize'] = 8, 6
```

## Discrete prolate spheroidal sequences and PSD basics

To begin, create a simple test sequence that is a mixture of cosines.


```python
# amplitudes in (2, 5)
amps = np.random.rand(5) * 3 + 2
# phases in (0, 2pi)
phases = np.random.rand(5) * 2 * np.pi
# frequencies in (0.1, 0.4)
mix_freqs = np.random.rand(5) * 0.3 + 0.1
# mixture of cosines
n = np.arange(1024)
x_parts = amps[:, np.newaxis] * np.cos(2 * np.pi * n * mix_freqs[:, np.newaxis] + phases[:, np.newaxis])
x = np.sum(x_parts, axis=0)
plt.figure()
plt.plot(n, x_parts.T + np.arange(2, 7) * 10, lw=0.25)
plt.plot(n, x, lw=0.5, color='k')
plt.yticks([])
sns.despine(left=True)
```

The NW parameter and the resulting eigenvalues relate to the "spectral concentration" property of DPSS. NW is a "time-bandwidth" product. The main lobe of the the DPSS spectrum is about ±NW DFT frequency bins wide. Also, the ratio of energy inside the lobe versus total signal energy is given by the eigenvalue.

Since the DFT bin size is $1/N$, the simple (normalized) bandwidth identity is $W_{(n)}=NW(1/N)$. Note that

1. This is *half* of the full estimator bandwidth
1. To convert to frequency in Hz, multiply by the sampling rate: $W_{(f)}=NW(f_{s}/N)$


```python
NW = 3.5
mte = mt.MultitaperEstimator(len(x), NW, fs=1, low_bias=0.9)
mte.eigs
```


```python
f, axs = plt.subplots(3, 1, figsize=(6, 8))
lns = axs[0].plot(n, mte.dpss.T)
axs[0].set_xlabel('samples')
axs[0].legend(lns[:1], ('DPSS',))

freqs, spectra = periodogram(mte.dpss, fs=1.0, detrend=False, window='boxcar')
spectra = spectra[..., :len(freqs)]
band_limit = float(mte.NW) / len(x)
axs[1].semilogy(freqs, spectra.T)
axs[1].set_xlabel('Normalized frequency')
axs[2].semilogy(freqs, spectra.T, marker='.', ls='-')
axs[2].axvline(band_limit, color='k', linestyle='dashed')
axs[2].set_xlim(0, 15.0 / len(x))
axs[2].set_ylim(1e-8, 2)
axs[2].annotate('Concentration band-limit (NW / N)', (band_limit, 0.01), (band_limit * 1.5, 0.01),
                arrowprops=dict(width=1))
_ = axs[2].set_xlabel('Normalized frequency')
f.tight_layout(pad=0)
```

The DPSS are progressively *less* concentrated as the order increases, as expected from the decreasing eigenvalues. Interestingly, these spectra are computed via periodogram suffer from the broad-band bias (spectral leakage out of the main lobe) that tapering is supposed to address. The next plot shows the reduction of broad-band bias using a single taper (Hamming window) versus the square wave taper.

The advantage of multitaper is to use *multiple* orthogonal tapers to create uncorrelated estimates that can be averaged to reduce variance.


```python
plt.figure()
freqs, spectra = periodogram(mte.dpss, fs=1.0, detrend=False, window='boxcar')
lns_a = plt.semilogy(freqs, spectra.T, lw=1, color='k')
freqs, spectra = periodogram(mte.dpss, fs=1.0, detrend=False, window='hamming')
lns_b = plt.semilogy(freqs, spectra.T, lw=1, color='r')
plt.legend(lns_a[:1] + lns_b[:1], ('Flat taper', 'Hamming taper'))
plt.title('Hamming window: broad-band bias reduction')
```

So the multitaper estimate is essentially an average of $K$ individual "direct spectral estimates" of the series $x(t)$. Let $v_{k}(t)$ be the k'th DPSS taper (technically also uniquely parameterized by $(N, NW)$). A direct estimate is

\begin{equation}
\hat{S}_{k}(\omega)=\left|y_{k}(\omega)\right|^{2}
\tag{1}
\end{equation}

with $y_{k}(\omega)$ being the DFT of the signal-taper product. 

\begin{equation}
y_{k}(\omega)=\sum_{t=0}^{N-1} x(t)v_{k}(t)\exp \{-i\omega t\}
\tag{2}
\end{equation}

It is in fact a tapered periodogram estimate, but all $\hat{S}_{k}(\omega)$ are (basically) uncorrelated. The normalization for a periodogram estimate is usually $1/N$, but this is not needed with the orthonormal taper.

The final multitaper estimator is the average (or weighted average) of these functions

\begin{equation}
\hat{S}(\omega)=\frac{1}{K}\sum_{k=0}^{K-1}\hat{S}_{k}(\omega)
\tag{3}
\end{equation}

Multiplying by a factor of two is conventional to count power from the negative side of the spectrum (which is then discarded). 


```python
freqs, mt_psd = mte.compute_psd(x, adaptive_weights=True)
freqs, pg_psd = periodogram(x, window='boxcar', detrend=False)
freqs, ht_psd = periodogram(x, window='hamming', detrend=False)
f, axs = plt.subplots(2, 1)
axs[0].semilogy(freqs, np.c_[pg_psd, ht_psd, mt_psd])
axs[0].legend(('P-gram (boxcar)', 'P-gram (hamming)', 'Multitaper'))
axs[1].semilogy(freqs, np.c_[pg_psd, ht_psd, mt_psd])
axs[1].set_xlim(mix_freqs[2] - band_limit * 4, mix_freqs[2] + band_limit * 4)
```

The Hamming window improves the broad-band bias compared to a plain periodgram. The multitaper estimator--especially using adaptive weighting for summing low power estimates--has by far the lowest leakage bias. However this comes at the expense of bias within the ±W frequency window. With this test signal, the **true** power spectral density (the quantity we're trying to estimate) is a collection of delta functions at the mixed frequencies. Each type of taper (square, Hamming, MT) blurs that line to an extent: the full power of each component is spread across the estimator bandwidth. This implies that larger NW values would blur the line value across a larger window (more bias).

The full variance is given by $E[(x-\bar{x})^{2}]=\int_{0}^{f_{s}} P_{x}(f)df$


```python
print('Input variance', np.var(x))
print('MT variance:', np.trapz(mt_psd, x=freqs))
print('P-gram(BC) variance:', np.trapz(pg_psd, x=freqs))
print('P-gram(H) variance:', np.trapz(ht_psd, x=freqs))
```

## PSD for real data

Grabbing a random $2^{14}$ point segement from a rat auditory cortex recording.


```python
data = load_experiment_auto('viventi/2017-11-28_acute', 'test_003', mapped=True)
```


```python
n0 = int(100 * data.Fs)
n_pts = 2 ** 14
rando_data = detrend(data.data[:, n0:n0 + n_pts], axis=1)
fig = tile_traces_1ax(rando_data, p=data.chan_map, twin=(0, 1e3 * n_pts / data.Fs), calib_unit='uV', linewidths=0.25)
```

Despite the larger in-window bias, the multitaper PSD has much less estimator variance than either the boxcar or Hamming-window periodogram. Note the amount of deviation from the $1/f$ LFP + transistor noise spectrum in the different estimates.


```python
x = rando_data[data.chan_map.lookup(3, 0)]
# Using the "psd" class method
freqs, mt_psd_1 = mt.MultitaperEstimator.psd(x, NW=3.5, fs=data.Fs, adaptive_weights=True)
freqs, mt_psd_2 = mt.MultitaperEstimator.psd(x, NW=7, fs=data.Fs, adaptive_weights=True)
freqs, pg_psd = periodogram(x, window='boxcar', detrend=False, fs=data.Fs)
freqs, ht_psd = periodogram(x, window='hamming', detrend=False, fs=data.Fs)
f, axs = plt.subplots(1, 1)
axs.loglog(freqs, np.c_[pg_psd, ht_psd, mt_psd_1, mt_psd_2], alpha=0.6)
axs.set_ylim(bottom=1e-3)
axs.set_xlim(left=0.5)
axs.legend(('P-gram (boxcar)', 'P-gram (hamming)', 'Multitaper (NW=3.5)', 'Multitaper (NW=7)'))
```

The multitaper variance relationship is closer than a simple tapered periodogram. The naked periodogram is essentially the same as the sample variance because of Parseval's theorem of the DFT--in short the sum of squares is the same under $x(t)$ and $X(\omega_k)$.


```python
print('Input variance', np.var(x))
print('MT(NW=3.5) variance:', np.trapz(mt_psd_1, x=freqs))
print('MT(NW=7) variance:', np.trapz(mt_psd_2, x=freqs))
print('P-gram(BC) variance:', np.trapz(pg_psd, x=freqs))
print('P-gram(H) variance:', np.trapz(ht_psd, x=freqs))
```

## PSD estimator variance & confidence intervals
We can explicitly find the variance of the spectral estimator, and from that variance form confidence intervals. Confidence intervals can be estimated two ways.

* intervals based on a $\chi{}^{2}_{2K}$ distribution
* intervals based on the Jackknife standard error estimate and Student's $t_{2K-1}$ distribution

### $\chi^{2}$ argument
The $\chi{}^{2}$ interval is derived from normality assumptions about the timeseries (serial independence, stationarity). Both these assumptions are particularly bad for LFP, btw. For periodograms in general, the estimator can be considered the sum of squares:

\begin{equation}
\hat{S}_{\omega}=\left(A_{\omega}^{2}+B_{\omega}^{2}\right)
\tag{4}
\end{equation}

and $A$ and $B$ are the cosine and sine components of the DFT of taper-signal product $v(t)x(t)$. 

If the series $x(t)$ is iid zero mean Gaussian with variance $\sigma^{2}_{x}$, then $A$ and $B$ are both sums of Gaussians (thus also Gaussian). The expected value is $\sum E\{x(t)\}\cos\omega t=0$ and variance is proportional to $\sigma_{x}^{2}$ (in the case of eq 2 above, $Var(A_{\omega})=\sigma_{x}^{2}/2$). So the appropriately normalized sum of squares in eq 4 is distributed as $\chi^{2}$ with 2 degrees of freedom (DOF)

\begin{equation}
\frac{A_{\omega}^{2}+B_{\omega}^{2}}{\sigma_{x}^{2}/2}=\frac{2\hat{S}_{\omega}}{\sigma_{x}^{2}}\sim \chi^{2}_{2}
\end{equation}

Asymptotically (as $N\rightarrow \infty$) $\sigma_{x}^{2}$ is replaced by the *true* power spectral density (the quantity we're trying to estimate) for this fundamental result

\begin{equation}
\frac{2\hat{S}_{\omega}}{S_{\omega}}\sim \chi^{2}_{2}
\end{equation}

For the multitaper estimator, $K\hat{S}^{(mt)}(\omega)=\sum_{k=0}^{K-1}\hat{S}_{k}(\omega)$ is a sum or more uncorrelated squared-Gaussians, which just increases the $\chi^{2}$ DOF

\begin{equation}
\frac{2K\hat{S}^{(mt)}_{\omega}}{S_{\omega}}\sim \chi_{2K}^{2}
\end{equation}

So the ratio of the PSD estimator to the actual PSD is $chi^{2}$, which means we can calculate a confidence interval for the lowest and highest points that the ratio should reach. That CI can then be manipulated to show the CI for the PSD itself. The 95% CI is

\begin{equation}
\frac{2K\hat{S}^{(mt)}_{\omega}}{\chi^{2}_{0.975,2K}}\lt S_{\omega} \lt \frac{2K\hat{S}^{(mt)}_{\omega}}{\chi^{2}_{0.025,2K}}
\end{equation}

where $\chi^{2}_{\alpha,2K}$ denotes the inverse cumulative distribution function at point $\alpha$. If the direct spectral estimates are weighted with the adaptive weighting scheme, then the effective DOF can be slightly different at each frequency, as a function of weights. In this case

\begin{equation}
\hat{S}^{(mt)}=\frac{\sum_{k=0}^{K-1}\left|d_{k}(\omega)\right|^{2}\hat{S}_{k}(\omega)}{\sum_{k=0}^{K-1}\left|d_{k}(\omega)\right|^{2}}
\end{equation}

and the effective DOF is $\nu(\omega)=2\sum_{k=0}^{K-1}\left|d_{k}(\omega)\right|^{2}$.

For the standard case, the bounds get clearly tighter with additional estimates averaged in


```python
tk = np.arange(1, 8) * 2
print('upper bound scaling:', tk / dists.chi2.ppf(0.025, tk))
print('lower bound scaling:', tk / dists.chi2.ppf(0.975, tk))
plt.figure()
plt.semilogy(tk, tk / dists.chi2.ppf(0.025, tk), label='upper scaling', marker='*', ms=10)
plt.semilogy(tk, tk / dists.chi2.ppf(0.975, tk), label='lower scaling', marker='*', ms=10)
plt.axhline(1, color='k')
plt.legend()
plt.xlabel('DOF')
```

### Jackknife variance

The "Jackknife" estimate of variance follows from a leave-one-out resampling technique called the [jackknife](https://en.wikipedia.org/wiki/Jackknife_resampling). For the purpose of the multitaper estimator of a PSD, the jackknife calculates $K$ versions of eq 3 leaving out one direct spectral estimate at a time. Call these $\hat{S}_{-k}(\omega)$, and their average is $\bar{S}(\omega)$. The jackknife variance of an estimator is just the scaled sample variance of these $K$ sub-estimators:

\begin{equation}
\operatorname{var}\left[S^{(mt)}(\omega)\right]=\frac{n-1}{n}\sum_{k=0}^{K-1}\left(\hat{S}_{-k}(\omega)-\bar{S}(\omega)\right)^{2}
\end{equation}

and the standard error (SE) is the square root of this variance. Using the SE, the normal CI is based on Student's t distribution with $K-1$ degrees of freedom.

\begin{equation}
\hat{S}^{(mt)}(\omega)-(SE)t_{0.975,K-1}\lt S(\omega)\lt\hat{S}^{(mt)}(\omega)+(SE)t_{0.975,K-1}
\end{equation}


```python
x = rando_data[data.chan_map.lookup(3, 0)]
# Using the "psd" class method
mte_1 = mt.MultitaperEstimator(len(x), NW=3.5, fs=data.Fs, low_bias=True)
dof_1 = 2 * len(mte_1.dpss)
freqs, mt_psd_1, ci_1 = mte_1.compute_psd(x, adaptive_weights=True, ci=True)
mte_2 = mt.MultitaperEstimator(len(x), NW=7, fs=data.Fs, low_bias=True)
dof_2 = 2 * len(mte_2.dpss)
freqs, mt_psd_2, ci_2 = mte_2.compute_psd(x, adaptive_weights=True, ci=True)
freqs, mt_psd_3, ci_3 = mt.MultitaperEstimator.psd(x, NW=7, fs=data.Fs, adaptive_weights=True, 
                                                   ci=True, jackknife=True, jn_jobs=4)

f, axs = plt.subplots(2, 1, sharex=True, sharey=True)
axs[0].fill_between(freqs, ci_1[0], ci_1[1], color=(0.2, 0.2, 0.2), label='Chi2 {} dof'.format(dof_1))
axs[0].fill_between(freqs, ci_2[0], ci_2[1], color=(0.6, 0.6, 0.6), label='Chi2 {} dof'.format(dof_2))
axs[1].fill_between(freqs, ci_3[0], ci_3[1], color=(0.6, 0.2, 0.2), label='Jackknife SE')

axs[1].set_yscale('log')
axs[1].set_xscale('log')
# axs.loglog(freqs, mt_psd_3)
f_cutoff = freqs.searchsorted(freqs.max() * 0.9)
axs[1].set_ylim(bottom=0.25 * ci_1[0][f_cutoff])
axs[1].set_xlim(left=0.5)
axs[0].legend()
axs[1].legend()
_ = axs[0].set_title('Relative chi-squared confidence intervals')
```

## Cross spectral density (CSD)

For this section, we'll focus on an evoked tone response from this dataset.


```python
fft_pts = 256
all_tone_responses = extract_epochs(data.data, data.exp, pre=64, post=fft_pts - 64)
all_tone_responses = detrend(all_tone_responses, type='linear', axis=-1)
row_tone_responses = np.array([all_tone_responses[data.chan_map.lookup(3, i)] for i in range(8)])
tone_ss = np.sum(np.sum(row_tone_responses ** 2, axis=0), axis=1)
big_tone = np.argsort(tone_ss)[int(0.9 * len(tone_ss))]
```


```python
fig = tile_traces_1ax(all_tone_responses[:, big_tone], 
                      p=data.chan_map, 
                      twin=(-64 / data.Fs, (fft_pts - 64) / data.Fs),
                      calib_unit='uV', 
                      linewidths=0.5)
```


```python
frames = data.chan_map.embed(all_tone_responses[:, big_tone].T, axis=1)
time = (np.arange(frames.shape[0]) - 64) * 1e3 / data.Fs
animate_frames(frames, notebook=True, time=time, axis_toggle='off', colorbar=True)
```

This response has, in general, a right-to-left sweep. Earliest onset on the right-side electrodes and then sequential activation moving left. Cross spectral density (CSD) is a tool to look at magnitude and phase of covarying parts of the signal power per frequency. We should expect to see a structured phase lag between electrodes along the row transect.

The `MultitaperEstimator.csd` class method returns a "matrix" of cross spectral densities for the signals in `x`. Each row i and column j in the matrix has the CSD $C_{ij}(f)$. Like a covariance matrix, this is somewhat redundant with $C_{ij}(f)=C_{ji}^{*}(f)$ (* is complex conjugate). The conjugation just reflects that the phase relationship between two signals is reversed from the perspective of one signal versus the other. The *diagonal* of the matrix is actually equivalent to the normal PSDs of the signals.


```python
# x = np.row_stack([rando_data[data.chan_map.lookup(3, i)] for i in range(8)])
x = row_tone_responses[:, big_tone]
freqs, csd_matrix = mt.MultitaperEstimator.csd(x, NW=2.5, fs=data.Fs, adaptive_weights=True)
```


```python
with sns.husl_palette(n_colors=8):
    f, axs = plt.subplots(2, 1, sharex=True, figsize=(8, 8))
    axs[0].loglog(freqs, np.abs(csd_matrix[7, :7]).T, alpha=0.5)
    axs[0].loglog(freqs, csd_matrix[7, 7].real, color='k', alpha=0.5, ls='--')
    # axs[0].set_xlim(left=2, right=15)
    labels = ['(3,7) to (3,{})'.format(i) for i in range(7)]
    labels.append('(3,7) PSD')
    axs[0].legend(labels, ncol=2)
    axs[0].set_title('Cross-spectral density (power)')
    axs[1].semilogx(freqs, np.unwrap(np.angle(csd_matrix[7, :7])).T, alpha=0.5)
    # axs[1].semilogx(freqs, np.angle(csd_matrix[7, :7]).T, alpha=0.5)
    axs[1].set_title('Cross-spectral density (phase)')
    axs[1].set_ylim(-1 * np.pi, 3 * np.pi)
```

The CSD reflects two things about the response.

1. The covarying power between ~10-50 Hz decreases from right to left: i.e. as a function of distance
1. The covarying part of the response has a phase lag that increases from right to left.

The phase tells us the response happens later on channel (3, 0) than on (3, 7) because the phase difference is positive. This is most clear up to ~40 Hz. After 40 Hz the phase estimates become pretty ragged, partly due to phase wrapping but mostly due to truly random noise. Phase is unwrapped here, but after a certain point the unwrapping is a random walk.

### Coherence: normalized CSD
If CSD is analogous to the covariance, we can calculate an analogous correlation coefficient called coherence. Coherence is the CSD normalized by the square-root product of the two signals' PSDs. Its magnitude is directly analogous to the correlation coefficient. Its phase is equal to the CSD phase (since the normalization introduces no new phase). For this reason, magnitude squared coherence (MSC) is usually calculated.

MSC is a *very* noisy estimator, and it typically needs averaging over several trials.


```python
all_8k = np.where(data.exp.tones == 8000)[0][:10]
```


```python
coh_specs = []
row_sites = [data.chan_map.lookup(3, i) for i in range(8)]
for trial in all_8k:
    x = np.array([all_tone_responses[chan, trial] for chan in row_sites])
    coh_spec = mt.coherence(x, 2.5, msc=True)
    coh_specs.append(coh_spec)
coh_spec_avg = np.mean(coh_specs, axis=0)
```


```python
with sns.husl_palette(n_colors=8):
    f, axs = plt.subplots(1, 1, sharex=True)
    axs.semilogx(freqs, np.abs(coh_spec_avg[7, 0:7]).T ** 2, alpha=0.5)
    labels = ['(3,7) to (3,{})'.format(i) for i in range(7)]
    axs.legend(labels, ncol=2)
    axs.set_ylabel('Mag. squared coherence')
```

The jackknife can be used to find the estimator variance and compute a confidence interval. This gives some insight into the reliability of MSC on short windows. The confidence interval improves with higher NW, but the bandwidth resolution goes down. Note that when using a jackknife, you really want more than ~3 samples. For this reason, we'll tune down the eigenvalue threshold from a default 0.99 to 0.9 for DPSS with NW=2.5.


```python
coh_spec_1, ci_1 = mt.coherence(x, 2.5, msc=True, ci=True, low_bias=.9, jn_jobs=4)
bw_1 = mt.nw2bw(2.5, x.shape[-1], data.Fs)
coh_spec_2, ci_2 = mt.coherence(x, 5, msc=True, ci=True, low_bias=True, jn_jobs=4)
bw_2 = mt.nw2bw(6.5, x.shape[-1], data.Fs)
```


```python
f, ax = plt.subplots(1, 1)
l1 = 'BW={:.1f}'.format(bw_1)
l2 = 'BW={:.1f}'.format(bw_2)
filled_interval(ax.plot, freqs, coh_spec_1[7, 0], ci_1[:, 7, 0], alpha=0.5, ax=ax, label=l1)
filled_interval(ax.plot, freqs, coh_spec_2[7, 0], ci_2[:, 7, 0], alpha=0.5, ax=ax, label=l2)
ax.legend()
# ax.set_xscale('log')
```

The higher resolution estimator shows a lot more peaks than the lower resolution estimator. But the CIs are pretty big, so hard to say if it's legit.

## Todo: higher order spectra, spectrograms

0a1a87e2fffab9ced868788530479ed468705425
ipynb
Jupyter Notebook
docs/source/usage_demos/multitaper_estimator.ipynb
miketrumpis/ecoglib
2ecc5bc64920d96e01297cce5472d4b4797c3a7d
BSD-3-Clause
2021-11-06T21:39:01.000Z
2021-11-06T21:39:01.000Z
docs/source/usage_demos/multitaper_estimator.ipynb
miketrumpis/ecoglib
2ecc5bc64920d96e01297cce5472d4b4797c3a7d
BSD-3-Clause
docs/source/usage_demos/multitaper_estimator.ipynb
miketrumpis/ecoglib
2ecc5bc64920d96e01297cce5472d4b4797c3a7d
BSD-3-Clause
2022-01-10T20:40:18.000Z
2022-01-10T20:40:18.000Z
Qwen/Qwen-72B
1. YES
2. YES