Coverage for src/spectroflat/sensor/flat.py: 59%

1import numpy as np 

2 

3from .smooth_model import SmoothModel 

4from ..base.sensor_flat_config import SensorFlatConfig 

5 

6 

7class Flat: 

8 """ 

9 The `SensorFlat` class is a pre-flat field that contains most of the hard flat field features. 

10 Most prominently the following is corrected: 

11 - Sensor features (e.g. dust on the sensor itself) 

12 - Slit features (e.g. dust on the slit resulting in line features in the spectral direction) 

13 - (Some of) polarimetric fringes. 

14 

15 Note 1: The `SensorFlat` is not a full flat and does not correct for gradients nor all fringes. 

16 

17 Note 2: This class expects 2D data. It is a good idea to have such a flat per modulation state. 

18 """ 

19 

20 @staticmethod 

21 def from_frame(flat_frame: np.array, config: SensorFlatConfig): 

22 """ 

23 Create a `SensorFlat` from an average flat frame. 

24 Please note that this method expects a 2d image. 

25 """ 

26 model = SmoothModel(flat_frame, config).create() 

27 if config.roi is None: 

28 flat = Flat.save_divide(flat_frame, model.img) 

29 flat = flat / np.mean(flat) 

30 else: 

31 temp = Flat.save_divide(flat_frame[config.roi], model.img[config.roi]) 

32 temp = temp / np.mean(temp) 

33 flat = np.ones(flat_frame.shape) 

34 flat[config.roi] = temp 

35 # import matplotlib.pyplot as plt 

36 # fig, ax = plt.subplots(nrows=3) 

37 # ax[0].imshow(model.img) 

38 # ax[1].imshow(flat_frame) 

39 # ax[2].imshow(flat) 

40 # plt.show() 

41 return Flat(flat) 

42 

43 def __init__(self, flat: np.array): 

44 self.flat = flat 

45 

46 def correct(self, img: np.array): 

47 return Flat.save_divide(img, self.flat) 

48 

49 @staticmethod 

50 def save_divide(a: np.array, b: np.array) -> np.array: 

51 return np.true_divide(a.astype('float32'), b.astype('float32'), out=a.astype('float32'), 

52 where=b != 0, dtype='float64')