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")
