-
Notifications
You must be signed in to change notification settings - Fork 112
Expand file tree
/
Copy pathspectra.py
More file actions
171 lines (136 loc) · 6.87 KB
/
spectra.py
File metadata and controls
171 lines (136 loc) · 6.87 KB
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
"""Power spectrum plotting functions.
Notes
-----
This file contains functions for plotting power spectra, that take in data directly.
"""
from itertools import repeat
import numpy as np
from fooof.core.modutils import safe_import, check_dependency
from fooof.plts.settings import PLT_FIGSIZES
from fooof.plts.style import check_n_style, style_spectrum_plot
from fooof.plts.utils import check_ax, add_shades, check_plot_kwargs
plt = safe_import('.pyplot', 'matplotlib')
###################################################################################################
###################################################################################################
@check_dependency(plt, 'matplotlib')
def plot_spectrum(freqs, power_spectrum, log_freqs=False, log_powers=False,
ax=None, plot_style=style_spectrum_plot, **plot_kwargs):
"""Plot a power spectrum.
Parameters
----------
freqs : 1d array
Frequency values, to be plotted on the x-axis.
power_spectrum : 1d array
Power values, to be plotted on the y-axis.
log_freqs : bool, optional, default: False
Whether to plot the frequency axis in log spacing.
log_powers : bool, optional, default: False
Whether to plot the power axis in log spacing.
ax : matplotlib.Axes, optional
Figure axis upon which to plot.
plot_style : callable, optional, default: style_spectrum_plot
A function to call to apply styling & aesthetics to the plot.
**plot_kwargs
Keyword arguments to be passed to the plot call.
"""
ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
# Set plot data & labels, logging if requested
plt_freqs = np.log10(freqs) if log_freqs else freqs
plt_powers = np.log10(power_spectrum) if log_powers else power_spectrum
# Create the plot
plot_kwargs = check_plot_kwargs(plot_kwargs, {'linewidth' : 2.0})
ax.plot(plt_freqs, plt_powers, **plot_kwargs)
check_n_style(plot_style, ax, log_freqs, log_powers)
@check_dependency(plt, 'matplotlib')
def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, labels=None,
ax=None, plot_style=style_spectrum_plot, **plot_kwargs):
"""Plot multiple power spectra on the same plot.
Parameters
----------
freqs : 2d array or 1d array or list of 1d array
Frequency values, to be plotted on the x-axis.
power_spectra : 2d array or list of 1d array
Power values, to be plotted on the y-axis.
log_freqs : bool, optional, default: False
Whether to plot the frequency axis in log spacing.
log_powers : bool, optional, default: False
Whether to plot the power axis in log spacing.
labels : list of str, optional
Legend labels, for each power spectrum.
ax : matplotlib.Axes, optional
Figure axes upon which to plot.
plot_style : callable, optional, default: style_spectrum_plot
A function to call to apply styling & aesthetics to the plot.
**plot_kwargs
Keyword arguments to be passed to the plot call.
"""
ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
# Make inputs iterable if need to be passed multiple times to plot each spectrum
freqs = repeat(freqs) if isinstance(freqs, np.ndarray) and freqs.ndim == 1 else freqs
labels = repeat(labels) if not isinstance(labels, list) else labels
for freq, power_spectrum, label in zip(freqs, power_spectra, labels):
plot_spectrum(freq, power_spectrum, log_freqs, log_powers, label=label,
plot_style=None, ax=ax, **plot_kwargs)
check_n_style(plot_style, ax, log_freqs, log_powers)
@check_dependency(plt, 'matplotlib')
def plot_spectrum_shading(freqs, power_spectrum, shades, shade_colors='r', add_center=False,
ax=None, plot_style=style_spectrum_plot, **plot_kwargs):
"""Plot a power spectrum with a shaded frequency region (or regions).
Parameters
----------
freqs : 1d array
Frequency values, to be plotted on the x-axis.
power_spectrum : 1d array
Power values, to be plotted on the y-axis.
shades : list of [float, float] or list of list of [float, float]
Shaded region(s) to add to plot, defined as [lower_bound, upper_bound].
shade_colors : str or list of string
Color(s) to plot shades.
add_center : bool, optional, default: False
Whether to add a line at the center point of the shaded regions.
ax : matplotlib.Axes, optional
Figure axes upon which to plot.
plot_style : callable, optional, default: style_spectrum_plot
A function to call to apply styling & aesthetics to the plot.
**plot_kwargs
Keyword arguments to be passed to the plot call.
"""
ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
plot_spectrum(freqs, power_spectrum, plot_style=None, ax=ax, **plot_kwargs)
add_shades(ax, shades, shade_colors, add_center, plot_kwargs.get('log_freqs', False))
check_n_style(plot_style, ax,
plot_kwargs.get('log_freqs', False),
plot_kwargs.get('log_powers', False))
@check_dependency(plt, 'matplotlib')
def plot_spectra_shading(freqs, power_spectra, shades, shade_colors='r', add_center=False,
ax=None, plot_style=style_spectrum_plot, **plot_kwargs):
"""Plot a group of power spectra with a shaded frequency region (or regions).
Parameters
----------
freqs : 2d array or 1d array or list of 1d array
Frequency values, to be plotted on the x-axis.
power_spectra : 2d array or list of 1d array
Power values, to be plotted on the y-axis.
shades : list of [float, float] or list of list of [float, float]
Shaded region(s) to add to plot, defined as [lower_bound, upper_bound].
shade_colors : str or list of string
Color(s) to plot shades.
add_center : bool, optional, default: False
Whether to add a line at the center point of the shaded regions.
ax : matplotlib.Axes, optional
Figure axes upon which to plot.
plot_style : callable, optional, default: style_spectrum_plot
A function to call to apply styling & aesthetics to the plot.
**plot_kwargs
Keyword arguments to be passed to `plot_spectra` or to the plot call.
Notes
-----
Parameters for `plot_spectra` can also be passed into this function as keyword arguments.
This includes `log_freqs`, `log_powers` & `labels`. See `plot_spectra` for usage details.
"""
ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
plot_spectra(freqs, power_spectra, ax=ax, plot_style=None, **plot_kwargs)
add_shades(ax, shades, shade_colors, add_center, plot_kwargs.get('log_freqs', False))
check_n_style(plot_style, ax,
plot_kwargs.get('log_freqs', False),
plot_kwargs.get('log_powers', False))