Coverage for src/spectroflat/smile/smile_detection.py: 75%

1#!/usr/bin/env python3 

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

3""" 

4provides SmileMapGenerator, SmileDetector, LineSmileFitter 

5 

6@author: hoelken 

7""" 

8from datetime import datetime 

9 

10import numpy as np 

11 

12from .offset_map import OffsetMap 

13from .smile_fit import SmileFit 

14from ..base import Logging 

15from ..base.smile_config import SmileConfig 

16 

17logger = Logging.get_logger() 

18 

19 

20class SmileMapGenerator: 

21 """ 

22 ##SmileMapGenerator 

23 This class generates a `OffsetMap` from a set of input files. 

24 For each state: In every row lines are detected and line cores fitted. 

25 Then the cores are shifted wrt to the central row until the chi² error 

26 between current and central row is minimal 

27 

28 The `OffsetMap` can be saved to a fits file. 

29 Metadata in the smile map with keywords similar to the FITS images will help to test 

30 the applicability of the map for an image to correct. 

31 """ 

32 

33 def __init__(self, smile_config: SmileConfig, cube: np.array): 

34 #: or for emission lines (thus the image is inverted) 

35 self._config = smile_config 

36 #: The image cube to process 

37 self._cube = cube 

38 #: Result: The generated OffsetMap 

39 self.omap = OffsetMap() 

40 

41 def run(self, out_file=None): 

42 self._collect_metadata() 

43 self._configure_builder() 

44 self._detect_smile() 

45 self._write_to_disk(out_file) 

46 return self 

47 

48 def _collect_metadata(self): 

49 self.omap.header['Type'] = 'OffsetMap' 

50 self.omap.header['GenTime'] = datetime.now().strftime("%Y-%m-%dT%H:%M:%S") 

51 self.omap.header['DIMS'] = repr(['state', 'y', 'lambda']) 

52 self.omap.header['Shape'] = repr(self._cube.shape) 

53 self.omap.header['Rotation'] = repr(self._config.rotation_correction) 

54 if self._config.roi is not None: 

55 self.omap.header['ROI'] = repr(self._config.roi) 

56 

57 def _configure_builder(self): 

58 if self._config.roi is None: 

59 self._config.total_rows = self._cube.shape[-2] 

60 self._config.total_cols = self._cube.shape[-1] 

61 else: 

62 self._config.total_rows = self._config.roi[0].stop - self._config.roi[0].start 

63 self._config.total_cols = self._config.roi[1].stop - self._config.roi[1].stop 

64 

65 def _detect_smile(self): 

66 if self._config.roi is None: 

67 sf = SmileFit(self._cube, self._config) 

68 else: 

69 roi = (slice(0, self._cube.shape[0]), self._config.roi[0], self._config.roi[1]) 

70 sf = SmileFit(self._cube[roi], self._config) 

71 sf.run() 

72 if self._config.roi is None: 

73 self.omap.map = sf.shift_map 

74 self.omap.error = sf.chi2_map 

75 else: 

76 y0 = roi[1].start 

77 y1 = self._cube.shape[1] - roi[1].stop 

78 x0 = roi[2].start 

79 x1 = self._cube.shape[2] - roi[2].stop 

80 self.omap.map = np.pad(sf.shift_map, ((0, 0), (y0, y1), (x0, x1)), 

81 'constant', constant_values=(0, 0)) 

82 self.omap.error = np.pad(sf.chi2_map, ((0, 0), (y0, y1), (x0, x1)), 

83 'constant', constant_values=(0, 0)) 

84 

85 def _write_to_disk(self, out_file): 

86 if out_file is None: 

87 return 

88 

89 logger.info('Writing OffsetMap to %s', out_file) 

90 self.omap.dump(out_file)