Coverage for src/spectroflat/smile/offset_map.py: 58%

1#!/usr/bin/env python3 

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

3""" 

4provides OffsetMap 

5 

6@author: hoelken 

7""" 

8from __future__ import annotations 

9 

10import os 

11 

12import numpy as np 

13from astropy.io import fits 

14 

15from ..base import Logging 

16 

17log = Logging.get_logger() 

18 

19 

20class OffsetMap: 

21 """ 

22 ## Offset Map 

23 This class holds the map of offsets by the smile per line. 

24 """ 

25 

26 @staticmethod 

27 def from_file(path) -> OffsetMap: 

28 with fits.open(path) as hdul: 

29 hdul[0].verify('silentfix') 

30 om = OffsetMap() 

31 om.header = hdul[0].header 

32 om.map = hdul[0].data 

33 om.error = hdul[1].data 

34 om.path = path 

35 return om 

36 

37 def __init__(self, header: dict = None): 

38 #: The map holds a numpy array with the shifts per mod state. 

39 #: If the map was squashed the map has shape y,x, else state,y,x 

40 self.map = np.zeros((0, 0, 0)) 

41 #: The error is the per pixel chi² error from the central line. 

42 #: If the map was squashed the map has shape y,x, else state,y,x 

43 self.error = np.zeros((0, 0, 0)) 

44 self.header = None 

45 self.path = None 

46 header = header if header is not None else {} 

47 header['squashed'] = False 

48 self.set_header(header) 

49 

50 def set_header(self, header: dict) -> None: 

51 self.header = fits.Header() 

52 for key, value in header.items(): 

53 self.header[key] = str(value) 

54 

55 def get_offsets(self, row: int, state: int = 0) -> np.array: 

56 img = self.map[state] if self.is_squashed else self.map 

57 return img[int(row)] 

58 

59 def get_map(self, state: int = 0) -> np.array: 

60 return self.map[state] if self.is_squashed() else self.map 

61 

62 def get_mean_errors(self): 

63 if self.is_squashed(): 

64 return [np.mean(self.error)] 

65 return [np.mean(self.error[s]) for s in range(self.map.shape[0])] 

66 

67 def get_max_errors(self): 

68 if self.is_squashed(): 

69 return [np.max(self.error)] 

70 return [np.max(self.error[s]) for s in range(self.map.shape[0])] 

71 

72 def dump(self, out_file, overwrite: bool = True): 

73 """ 

74 Dump a SmileMap to a `fits` file. 

75 HDU 0: shifts + header 

76 HDU 1: chi2 error map 

77 

78 :param out_file: The file to write to 

79 :param overwrite: Overwrite existing file with same name? 

80 """ 

81 os.makedirs(os.path.dirname(out_file), exist_ok=True) 

82 hdul = fits.HDUList() 

83 phdu = fits.PrimaryHDU(self.map, header=self.header) 

84 hdul.append(phdu) 

85 hdul.append(fits.ImageHDU(data=self.error)) 

86 log.info('Write OffsetMap to "%s"...', out_file) 

87 hdul.writeto(out_file, overwrite=overwrite) 

88 self.path = out_file 

89 

90 def is_squashed(self) -> bool: 

91 return self.header['squashed'] in ['True', True] 

92 

93 def squash(self) -> None: 

94 if self.is_squashed(): 

95 return 

96 

97 self.header['squashed'] = True 

98 s = self.map.shape[0] 

99 self.map = np.repeat(self.map.mean(axis=0, keepdims=True), s, axis=0) 

100 self.error = np.repeat(self.error.mean(axis=0, keepdims=True), s, axis=0) 

101 

102 def enforce_same_offsets_on_all_states(self): 

103 states = self.map.shape[0] 

104 self.map = np.repeat(np.mean(self.map, axis=0, keepdims=True, dtype='float32'), states, axis=0) 

105 self.error = np.repeat(np.mean(self.error, axis=0, keepdims=True), states, axis=0) 

106 

107 def __repr__(self) -> str: 

108 text = f'{self.__class__.__name__}\n' 

109 for k, v in self.header.items(): 

110 text += f'{k}:\t {v}\n' 

111 text += f'Shape:\t\t {self.map.shape}\n' 

112 text += f'Mean Error:\t{np.mean(self.error)}' 

113 return text