Coverage for src/susi/reduc/pipeline/processing_data.py: 97%

1#!/usr/bin/env python3 

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

3""" 

4Provides ProcessingData and ProcessorArgs data items 

5 

6@author: hoelken 

7""" 

8from dataclasses import dataclass 

9from typing import Union 

10 

11import numpy as np 

12 

13from ...base import Config 

14from ...io import Fits 

15 

16 

17@dataclass 

18class ProcessingData: 

19 """ 

20 DataItem to store and pass demodulation data between the different collaborating classes 

21 """ 

22 

23 #: The usual configuration 

24 config: Config 

25 #: The loaded modulation matrix FITS 

26 mod_matrix: np.array = None 

27 #: The demodulation matrix data 

28 demod_matrix: np.array = None 

29 #: The loaded dark image FITS 

30 dark_image: Fits = None 

31 #: The loaded slit flat 

32 slit_flat: Fits = None 

33 #: The loaded sensor flat 

34 sensor_flat: Fits = None 

35 #: The loaded soft flat (optionaly containing prefilter calibration) 

36 soft_flat: Fits = None 

37 #: The loaded prefilter calibration map 

38 prefilter_map: Fits = None 

39 #: The loaded wavelength calibration axis 

40 wl_cal_axis: Fits = None 

41 #: The loaded offset (smile) map 

42 offset_map: Fits = None 

43 #: Shear distortion factor 

44 shear_factor: float = None 

45 #: Global rigid rotation 

46 rot_angle: float = None 

47 #: Slit shifts to apply (processed version of slit mask position minus reference) 

48 #: {'files:': files, 'shifts': shifts} 

49 slit_shifts: np.array = None 

50 #: Reference position for shifts in slit dim and common roi of shifted batch 

51 #: {'offset': float, 'slope': float, 'file': str, 'common_roi': tuple} 

52 slit_flat_shift_ref: dict = None 

53 

54 def has(self, field: Union[str, list]) -> bool: 

55 """Check if the requested field is set""" 

56 if isinstance(field, str): 

57 if field == "dark_image" and self.config.base.no_dark_correction: 

58 return True 

59 return self.__dict__[field] is not None 

60 

61 return all(self.has(key) for key in field) 

62 

63 def release(self, field: Union[str, list]) -> None: 

64 if isinstance(field, list): 

65 for f in field: 

66 self.__dict__[f] = None 

67 else: 

68 self.__dict__[field] = None 

69 

70 

71@dataclass 

72class ProcessorArgs: 

73 """ 

74 DataItem to shift data for processing between threads and processes 

75 """ 

76 

77 #: The index of the batch. 

78 idx: int 

79 #: Frame maps are a list of lists, where each list contains a number of (file_name, mod_state) tuples 

80 #: The processor wil average all frames of the given mod state prior to demodulation. 

81 frame_maps: list 

82 #: The processing data configuration 

83 proc_data: ProcessingData