Coverage for src/susi/reduc/fields/dark_field_correction.py: 86%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4module for dark correction provides DFCorrector 

5 

6@author: hoelken 

7""" 

8 

9# package imports 

10import numpy as np 

11 

12from ...io import Fits 

13from ...base import Logging, DataMissMatchException 

14from ...base.header_keys import * 

15 

16logger = Logging.get_logger() 

17 

18 

19class DFCorrector: 

20 """ 

21 Provides algorithm for dark image correction 

22 """ 

23 

24 def __init__(self, dark: Fits, img: Fits): 

25 self.dark = dark 

26 self.img = img 

27 

28 def run(self, header_check=True): 

29 """ 

30 Applies the dark img correction to the average image. 

31 

32 ### Params 

33 - header_check: [Bool, optional] set to `False` to bypass the header check. 

34 

35 ### Returns 

36 [Array] of img data with the dark correction applied (img - dark) 

37 """ 

38 self.__check_header(header_check, CAMERA_ID, 'Dark Image is not from the same camera!') 

39 self.__check_header(header_check, INTEGRATION_TIME, 'Integration time of dark image does not match!') 

40 self.__check_shapes() 

41 return self.__correct_img() 

42 

43 def __check_header(self, fail_on_error, key, message): 

44 if self.dark.value_of(key) != self.img.value_of(key): 

45 if fail_on_error: 

46 logger.error('%s != %s', self.dark.value_of(key), self.img.value_of(key)) 

47 raise DataMissMatchException(message) 

48 else: 

49 logger.warning(message) 

50 

51 def __check_shapes(self): 

52 if len(self.img.data.shape) == 2: 

53 self.img.data = np.array([self.img.data]) 

54 for shape, array_type in zip([self.dark.data.shape, self.img.data.shape], ['dark', 'img']): 

55 if len(shape) != 3: 

56 raise ValueError(f"{array_type} does not have 3 dimensions but {len(self.dark.data.shape)}") 

57 if self.dark.data.shape[0] != 1: 

58 raise ValueError("Dark array must have shape (1, x, y) but has %s", self.dark.data.shape) 

59 if self.dark.data.shape[1:] != self.img.data.shape[1:]: 

60 if self.img.data.shape[1:] == (2047, 2048): 

61 # HACK: Cope with wrongly cropped legacy files. 

62 # TODO: Delete me, once no longer needed 

63 self.dark.data = self.dark.data[:, :-1, :] 

64 logger.warning('Adjusting dark image shape to legacy cropping') 

65 return 

66 msg = f"Shapes do not match: dark: {self.dark.data.shape[1:]} img: {self.img.data.shape[1:]}" 

67 raise DataMissMatchException(msg) 

68 

69 def __correct_img(self): 

70 result = np.empty(self.img.data.shape) 

71 for i in range(self.img.data.shape[0]): 

72 if (i + 1) % 100 == 0: 

73 logger.info('\t> %s/%s', i + 1, self.img.data.shape[0]) 

74 result[i] = np.subtract(self.img.data[i], self.dark.data[0], dtype='float32') 

75 return result