python-meegkit/examples/example_dss_line.py at 5f0ea01047055295ef81094d14a055e50abe8e2a · eliasbenyahia/python-meegkit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
"""
Remove line noise with ZapLine
==============================

Find a spatial filter to get rid of line noise [1]_.

Uses meegkit.dss_line().

References
----------
.. [1] de Cheveigné, A. (2019). ZapLine: A simple and effective method to remove
   power line artifacts. NeuroImage, 116356.
   https://doi.org/10.1016/j.neuroimage.2019.116356

"""
# Authors: Maciej Szul <maciej.szul@isc.cnrs.fr>
#          Nicolas Barascud <nicolas.barascud@gmail.com>
import os

import matplotlib.pyplot as plt
import numpy as np
from scipy import signal

from meegkit import dss
from meegkit.utils import create_line_data, unfold

###############################################################################
# Line noise removal
# =============================================================================

###############################################################################
# Remove line noise with dss_line()
# -----------------------------------------------------------------------------
# We first generate some noisy data to work with
sfreq = 250
fline = 50
nsamples = 10000
nchans = 10
data = create_line_data(n_samples=3 * nsamples, n_chans=nchans,
                        n_trials=1, fline=fline / sfreq, SNR=2)[0]
data = data[..., 0]  # only take first trial

# Apply dss_line (ZapLine)
out, _ = dss.dss_line(data, fline, sfreq, nkeep=1)

###############################################################################
# Plot before/after
f, ax = plt.subplots(1, 2, sharey=True)
f, Pxx = signal.welch(data, sfreq, nperseg=500, axis=0, return_onesided=True)
ax[0].semilogy(f, Pxx)
f, Pxx = signal.welch(out, sfreq, nperseg=500, axis=0, return_onesided=True)
ax[1].semilogy(f, Pxx)
ax[0].set_xlabel("frequency [Hz]")
ax[1].set_xlabel("frequency [Hz]")
ax[0].set_ylabel("PSD [V**2/Hz]")
ax[0].set_title("before")
ax[1].set_title("after")
plt.show()


###############################################################################
# Remove line noise with dss_line_iter()
# -----------------------------------------------------------------------------
# We first load some noisy data to work with
data = np.load(os.path.join("..", "tests", "data", "dss_line_data.npy"))
fline = 50
sfreq = 200
print(data.shape)  # n_samples, n_chans, n_trials

# Apply dss_line(), removing only one component
out1, _ = dss.dss_line(data, fline, sfreq, nfft=400, nremove=1)

###############################################################################
# Now try dss_line_iter(). This applies dss_line() repeatedly until the
# artifact is gone
out2, iterations = dss.dss_line_iter(data, fline, sfreq, nfft=400, show=True)
print(f"Removed {iterations} components")

###############################################################################
# Plot results with dss_line() vs. dss_line_iter()
f, ax = plt.subplots(1, 2, sharey=True)
f, Pxx = signal.welch(unfold(out1), sfreq, nperseg=200, axis=0,
                      return_onesided=True)
ax[0].semilogy(f, Pxx, lw=.5)
f, Pxx = signal.welch(unfold(out2), sfreq, nperseg=200, axis=0,
                      return_onesided=True)
ax[1].semilogy(f, Pxx, lw=.5)
ax[0].set_xlabel("frequency [Hz]")
ax[1].set_xlabel("frequency [Hz]")
ax[0].set_ylabel("PSD [V**2/Hz]")
ax[0].set_title("dss_line")
ax[1].set_title("dss_line_iter")
plt.tight_layout()
plt.show()