From 228ae50bc309fda8aa79654e987fd17f6d861452 Mon Sep 17 00:00:00 2001 From: Matteo Bachetti Date: Thu, 30 Nov 2023 17:08:36 +0100 Subject: [PATCH] Force correct format of excluded frequency intervals --- stingray/fourier.py | 9 +++++++-- stingray/tests/test_fourier.py | 5 +++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/stingray/fourier.py b/stingray/fourier.py index 47b5a8e08..31b9a71d8 100644 --- a/stingray/fourier.py +++ b/stingray/fourier.py @@ -157,9 +157,14 @@ def power_color( power_err = np.asarray(power_err) if frequencies_to_exclude is not None: - if len(np.shape(frequencies_to_exclude)): + if len(np.shape(frequencies_to_exclude)) == 1: frequencies_to_exclude = [frequencies_to_exclude] - if np.shape(frequencies_to_exclude)[1] != 2: + + if ( + not isinstance(frequencies_to_exclude, Iterable) + or len(np.shape(frequencies_to_exclude)) != 2 + or np.shape(frequencies_to_exclude)[1] != 2 + ): raise ValueError("frequencies_to_exclude must be of format [[f0, f1], [f2, f3], ...]") for f0, f1 in frequencies_to_exclude: frequency_mask = (input_frequency_low_edges > f0) & (input_frequency_high_edges < f1) diff --git a/stingray/tests/test_fourier.py b/stingray/tests/test_fourier.py index ebf90ccdf..18f851518 100644 --- a/stingray/tests/test_fourier.py +++ b/stingray/tests/test_fourier.py @@ -726,6 +726,11 @@ def test_bad_edges(self): with pytest.raises(ValueError, match="The maximum frequency is lower "): power_color(self.freq[good], self.power[good]) + def test_bad_excluded_interval(self): + for fte in ([1, 1.1, 3.0], [4], [[1, 1.1, 3.0]], 0, [[[1, 3]]]): + with pytest.raises(ValueError, match="frequencies_to_exclude must be of "): + power_color(self.freq, self.power, frequencies_to_exclude=fte) + def test_excluded_frequencies(self): pc0, _, pc1, _ = power_color(self.freq, self.power, frequencies_to_exclude=[1, 1.1]) # The colors calculated with these frequency edges on a 1/f spectrum should be 1