def test_remove_energy_max_peaks(self): """Test remove_energy with limited number of peaks""" t = np.arange(0, 5, 0.001) p1 = [1.0, 2.0, 0.4] p2 = [3.5, 1.5, 0.3] p3 = [2.0, 1.0, 0.2] peak1 = self.peak_norm(t, p1[0], p1[1], p1[2]) # Biggest peak2 = self.peak_lorentz(t, p2[0], p2[1], p2[2]) # Middle peak3 = self.peak_triangle(t, p3[0], p3[1], p3[2]) # Smallest s = peak1 + peak2 + peak3 max_peaks = 2 preprocessor = Preprocessor() _, p_out = preprocessor.remove_energy(t, s, max_peaks=max_peaks) # Check number of peaks self.assertEqual(len(p_out), 2, "Expected 2 peaks because of flag")
def test_remove_energy_default(self): t = np.arange(0, 5, 0.001) p1 = [1.0, 2.0, 0.4] p2 = [3.5, 1.5, 0.3] p3 = [2.0, 1.0, 0.2] peak1 = self.peak_norm(t, p1[0], p1[1], p1[2]) # Biggest peak2 = self.peak_lorentz(t, p2[0], p2[1], p2[2]) # Middle peak3 = self.peak_triangle(t, p3[0], p3[1], p3[2]) # Smallest s = peak1 + peak2 + peak3 preprocessor = Preprocessor() self.assertEqual(preprocessor.energy_ratio, 0.1, "Default energy ratio is 0.1") _, p_out = preprocessor.remove_energy(t, s) # Check number of peaks self.assertEqual(len(p_out), 3, "Expected 3 peaks by default") # Checking location msg = "Expected: {:.3}, Receive: {:.3}" self.assertTrue( abs(1.00 - p_out[0, 0]) < 0.01, msg.format(1.00, p_out[0, 0])) self.assertTrue( abs(3.50 - p_out[1, 0]) < 0.01, msg.format(3.50, p_out[1, 0])) self.assertTrue( abs(2.00 - p_out[2, 0]) < 0.01, msg.format(2.00, p_out[2, 0])) # Checking amplitude self.assertTrue( abs(1.99 - p_out[0, 1]) < 0.01, msg.format(1.99, p_out[0, 1])) self.assertTrue( abs(1.33 - p_out[1, 1]) < 0.01, msg.format(1.33, p_out[1, 1])) self.assertTrue( abs(0.93 - p_out[2, 1]) < 0.01, msg.format(0.93, p_out[2, 1])) # Checking width self.assertTrue( abs(0.40 - p_out[0, 2]) < 0.01, msg.format(0.40, p_out[0, 2])) self.assertTrue( abs(0.23 - p_out[1, 2]) < 0.01, msg.format(0.23, p_out[1, 2])) self.assertTrue( abs(0.12 - p_out[2, 2]) < 0.01, msg.format(0.12, p_out[2, 2]))
def test_remove_energy_different_energy_ratio(self): """Test remove_energy after changing energy_ratio to different value""" t = np.arange(0, 5, 0.001) p1 = [1.0, 2.0, 0.4] p2 = [3.5, 1.5, 0.3] p3 = [2.0, 1.0, 0.2] peak1 = self.peak_norm(t, p1[0], p1[1], p1[2]) # Biggest peak2 = self.peak_lorentz(t, p2[0], p2[1], p2[2]) # Middle peak3 = self.peak_triangle(t, p3[0], p3[1], p3[2]) # Smallest s = peak1 + peak2 + peak3 energy_ratio = 0.7 preprocessor = Preprocessor() _, p_out = preprocessor.remove_energy(t, s, energy_ratio=energy_ratio) # Check number of peaks self.assertEqual(len(p_out), 1, "One peak expected with large energy ratio")