fix list input to plot_spectra

fooof-tools · Aug 27, 2023 · d63aae0 · d63aae0
1 parent 6f48eed
commit d63aae0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git 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
@@ -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):