Coverage for src/susi/reduc/demodulation/demod_assembler.py: 100%

1""" 

2Module provides a data assembler for the demodulator 

3 

4@author hoelken 

5""" 

6 

7import numpy as np 

8 

9from ...db import FileDB 

10from ...io import FitsBatch 

11from ...base.header_keys import * 

12 

13 

14class DemodAssembler: 

15 """ 

16 The DemodAssembler provides suitable input for the Demodulator depending on a fits batch 

17 with at least one modulation cycle included. 

18 """ 

19 

20 def __init__(self, db: FileDB, batch: FitsBatch, with_data: bool = True): 

21 self.state_data = [] 

22 self.state_map = [] 

23 self.state_files = [] 

24 self.batch = batch 

25 self.with_data = with_data 

26 self.db = db 

27 self.blocks = [] 

28 

29 def run(self) -> list: 

30 """ 

31 Returns a list with a dictionary per modulation cycle providing the 

32 keys 

33 - 'data': [np.array] the data to demodulate with shape (state, x, y) 

34 - 'states': [list] the sorted mod states included in the data 

35 - 'name': [str] suggested file name 

36 - 'start': [str] file name of the first fits file taken into account. 

37 """ 

38 self.batch.sort_by(self.db.config.base.timestamp_field) 

39 for entry in self.batch.batch: 

40 state = DemodAssembler._state(entry) 

41 if state == 0: 

42 self._start_new_block() 

43 if self.with_data: 

44 self.state_data.append(entry['data'][0]) 

45 self.state_map.append(state) 

46 self.state_files.append(entry['file']) 

47 self._start_new_block() 

48 return self.blocks 

49 

50 def _start_new_block(self): 

51 if self.state_map: 

52 # if data was already binned in time the name is already an average 

53 avg_flag = (self.db.config.cam.temporal_binning != 1) & self.batch.is_applied('B') 

54 name = self.db.avrg_fname(self.state_files[0], self.state_files[-1], avg=avg_flag) 

55 self.blocks.append( 

56 { 

57 'states': self.state_map, 

58 'data': np.array(self.state_data), 

59 'name': name, 

60 'start': self.state_files[0], 

61 } 

62 ) 

63 self.state_files = [] 

64 self.state_data = [] 

65 self.state_map = [] 

66 

67 @staticmethod 

68 def _state(entry: dict) -> int: 

69 return int(entry['header'][MOD_STATE])