diff --git a/pcpostprocess/hergQC.py b/pcpostprocess/hergQC.py index ceb04e3..2342151 100644 --- a/pcpostprocess/hergQC.py +++ b/pcpostprocess/hergQC.py @@ -461,7 +461,8 @@ def qc6(self, recording1, win=None, label=''): val = np.mean(recording1[i:f]) if self.plot_dir and self._debug: - plt.axvspan(win[0], win[1], color='grey', alpha=.1) + if win is not None: + plt.axvspan(win[0], win[1], color='grey', alpha=.1) plt.plot(recording1, label='recording1') plt.savefig(os.path.join(self.plot_dir, f"qc6_{label}")) plt.clf() diff --git a/tests/test_herg_qc.py b/tests/test_herg_qc.py index af924ec..4561257 100644 --- a/tests/test_herg_qc.py +++ b/tests/test_herg_qc.py @@ -104,21 +104,17 @@ def test_qc2(self): voltage=self.voltage) # qc2 checks that raw and subtracted SNR are above a minimum threshold - recording = np.asarray([0, 1] * 500 + [0, 10] * 500) # snr = 70.75 - result = hergqc.qc2(recording) - self.assertTrue(result[0], f"({result[1]})") - - recording = np.asarray([0, 1] * 500 + [0, 6.03125] * 500) # snr = 25.02 - result = hergqc.qc2(recording) - self.assertTrue(result[0], f"({result[1]})") - - recording = np.asarray([0, 1] * 500 + [0, 6.015625] * 500) # snr = 24.88 - result = hergqc.qc2(recording) - self.assertFalse(result[0], f"({result[1]})") + test_matrix = [ + (10, True), # snr = 70.75 + (6.03125, True), # snr = 25.02 + (6.015625, False), # snr = 24.88 + (1, False), # snr = 1.0 + ] - recording = np.asarray([0, 1] * 1000) # snr = 1.0 - result = hergqc.qc2(recording) - self.assertFalse(result[0], f"({result[1]})") + for i, expected in test_matrix: + recording = np.asarray([0, 1] * 500 + [0, i] * 500) # snr = 70.75 + result = hergqc.qc2(recording) + self.assertEqual(result[0], expected, f"({i}: {result[1]})") # TODO: Test on select data @@ -200,6 +196,7 @@ def test_qc5(self): (0.5, False), (0.75, False), (1.0, False), + (2.0, False), ] for i, expected in test_matrix: