Coverage for src/spectroflat/sensor/smooth_model.py: 100%

1import numpy as np 

2 

3from ..base.sensor_flat_config import SensorFlatConfig 

4from ..utils import Collections 

5 

6 

7class SmoothModel: 

8 """ 

9 The `SmoothModel` is a smoothed version of the actual flat image. 

10 

11 It is created by fitting a polynomial of configured degree (usually 2 to 6) 

12 to each spectral position along the spatial resolution. 

13 As the Image typically suffers from the spectrographic curvature effect and 

14 optical aberrations the fitted version is an ideal representation of the 

15 flat field frame without these. 

16 """ 

17 

18 def __init__(self, img: np.array, config: SensorFlatConfig): 

19 self._orig = img 

20 self._config = config 

21 self._current_poly = None 

22 self.img = None 

23 

24 def create(self): 

25 self._configure() 

26 for x in self._cols: 

27 self._fit_column(x) 

28 self._model_column(x) 

29 self._transpose() 

30 self._remove_spacial_gradient() 

31 self._finalize() 

32 return self 

33 

34 def _configure(self): 

35 self._config.total_rows = self._orig.shape[0] 

36 self._config.total_cols = self._orig.shape[1] 

37 if self._config.roi is None: 

38 self.img = np.empty(self._orig.shape).T 

39 self._rows = range(self._config.total_rows) 

40 self._cols = range(self._config.total_cols) 

41 self._offset = 0 

42 else: 

43 shape = (self._config.roi[0].stop - self._config.roi[0].start, 

44 self._config.roi[1].stop - self._config.roi[1].start) 

45 self.img = np.empty(shape).T 

46 self._offset = self._config.roi[1].start 

47 self._rows = np.arange(self._config.roi[0].start, self._config.roi[0].stop) 

48 self._cols = range(self._config.roi[1].start, self._config.roi[1].stop) 

49 

50 def _fit_column(self, x: int): 

51 b = self._config.fit_border 

52 if self._config.roi is None: 

53 col = self._orig[:, x][b:-b] 

54 else: 

55 col = self._orig[self._config.roi[0], x][b:-b] 

56 col[col < 2] = col.mean() 

57 col = Collections.remove_sigma_outliers(col, self._config.sigma_mask) 

58 self._current_poly = np.poly1d(np.polyfit(self._rows[b:-b], col, self._config.spacial_degree)) 

59 

60 def _model_column(self, x: int): 

61 self.img[x - self._offset] = np.array([self._current_poly(self._rows)]) 

62 

63 def _remove_spacial_gradient(self): 

64 grad = np.average(self.img, axis=1) 

65 self.img = np.array([self.img[i] / grad[i] for i in range(len(grad))]) 

66 

67 def _transpose(self): 

68 self.img = self.img.T 

69 

70 def _finalize(self): 

71 if self._config.roi is None: 

72 return 

73 

74 img = np.ones(self._orig.shape) 

75 img[self._config.roi] = self.img 

76 self.img = img