Coverage for src/spectroflat/utils/line_detection.py: 92%

1import numpy as np 

2from scipy.signal import find_peaks 

3 

4from ..base.config import SmileConfig 

5from ..fitting.line_fit import LineFit 

6 

7 

8def find_line_cores(row: np.array, conf: SmileConfig, peaks: list = None) -> tuple: 

9 allow_shift = False if peaks is None else True 

10 if not conf.emission_spectrum: 

11 row = row * -1 

12 row, peaks = _find_peaks(row, conf, peaks) 

13 return peaks, _find_minima(peaks, row, conf.line_distance // 2, allow_shift, conf.error_threshold) 

14 

15 

16def _find_peaks(row: np.array, conf: SmileConfig, peaks: list = None) -> np.array: 

17 row = row / np.abs(row.mean()) 

18 row = row - row.min() 

19 h = row.mean() + conf.height_sigma * row.std() 

20 if peaks is None: 

21 peaks, _ = find_peaks(row, height=h, prominence=conf.line_prominence, distance=conf.line_distance) 

22 return row, peaks 

23 

24 

25def _find_minima(peaks: np.array, row: np.array, win: int, allow_shift: bool = False, max_error: float = 0.5) -> list: 

26 return [_fit_peak(p, row, win, allow_shift, max_error) for p in peaks] 

27 

28 

29def _fit_peak(peak: int, row: np.array, win: int, allow_shift: bool, max_error: float) -> float: 

30 thresholds = [0.25, 0.5, 0.75, 1, 1.5, 2, max_error] if allow_shift else [0.3, max_error] 

31 thresholds = sorted(thresholds) 

32 if allow_shift: 

33 a, b = _get_borders(peak, row, win) 

34 alt = np.argmax(row[a:b]) 

35 peaks = [peak, a + alt] 

36 else: 

37 peaks = [peak] 

38 for t in thresholds: 

39 for p in peaks: 

40 a, b = _get_borders(p, row, win) 

41 lf = LineFit(np.arange(a, b), row[a:b], error_threshold=t) 

42 try: 

43 return lf.run().max_location 

44 except RuntimeError: 

45 pass 

46 

47 return peaks[1] if allow_shift else None 

48 

49 

50def _get_borders(peak: int, row: np.array, win: int): 

51 return int(max(0, peak - win)), int(min(len(row), peak + win))