diff --git a/fooof/plts/spectra.py b/fooof/plts/spectra.py index dbb8cdf2..81ec7eb5 100644 --- a/fooof/plts/spectra.py +++ b/fooof/plts/spectra.py @@ -60,8 +60,8 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, freq_r freq_range = np.log10(freq_range) if log_freqs else freq_range # Make inputs iterable if need to be passed multiple times to plot each spectrum - plt_powers = np.reshape(power_spectra, (1, -1)) if np.ndim(power_spectra) == 1 else \ - power_spectra + plt_powers = np.reshape(power_spectra, (1, -1)) if isinstance(freqs, np.ndarray) and \ + np.ndim(power_spectra) == 1 else power_spectra plt_freqs = repeat(freqs) if isinstance(freqs, np.ndarray) and freqs.ndim == 1 else freqs # Set labels diff --git a/fooof/tests/plts/test_spectra.py b/fooof/tests/plts/test_spectra.py index 5ca7dbd9..6a181340 100644 --- a/fooof/tests/plts/test_spectra.py +++ b/fooof/tests/plts/test_spectra.py @@ -15,27 +15,31 @@ @plot_test def test_plot_spectra(tfm, tfg, skip_if_no_mpl): - # Test with 1d inputs - 1d freq array and list of 1d power spectra + # Test with 1d inputs - 1d freq array & list of 1d power spectra plot_spectra(tfm.freqs, tfm.power_spectrum, save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_plot_spectra_1d.png') - # Test with 1d inputs - 1d freq array and list of 1d power spectra + # Test with 1d inputs - 1d freq array & list of 1d power spectra plot_spectra(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]], - save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_plot_spectra_list_1d.png') + save_fig=True, file_path=TEST_PLOTS_PATH, + file_name='test_plot_spectra_list_1d.png') - # Test with multiple freq inputs - list of 1d freq array and list of 1d power spectra - plot_spectra([tfg.freqs, tfg.freqs], [tfg.power_spectra[0, :], tfg.power_spectra[1, :]], + # Test with multiple lists - list of 1d freqs & list of 1d power spectra (different f ranges) + plot_spectra([tfg.freqs, tfg.freqs[:-5]], + [tfg.power_spectra[0, :], tfg.power_spectra[1, :-5]], save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_plot_spectra_lists_1d.png') # Test with 2d array inputs plot_spectra(np.vstack([tfg.freqs, tfg.freqs]), np.vstack([tfg.power_spectra[0, :], tfg.power_spectra[1, :]]), - save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_plot_spectra_2d.png') + save_fig=True, file_path=TEST_PLOTS_PATH, + file_name='test_plot_spectra_2d.png') # Test with labels plot_spectra(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]], labels=['A', 'B'], - save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_plot_spectra_labels.png') + save_fig=True, file_path=TEST_PLOTS_PATH, + file_name='test_plot_spectra_labels.png') @plot_test def test_plot_spectra_shading(tfm, tfg, skip_if_no_mpl):