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

1#!/usr/bin/env python3 

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

3""" 

4Demodulation module provides Demodulator 

5 

6@author: hoelken 

7""" 

8import numpy as np 

9from spectroflat.sensor.flat import Flat 

10 

11from ...base import Logging, Globals 

12from ...io import Fits 

13 

14logger = Logging.get_logger() 

15 

16 

17class Demodulator: 

18 """ 

19 ## Demodulator 

20 

21 This class is performing the demodulation of the frames from one mod cycle to the four Stokes images. 

22 """ 

23 

24 def __init__(self, demod_matrix: np.array, img_stack: np.array, states: list, fname: str = ''): 

25 #: The config data item 

26 self.demod_matrix = demod_matrix 

27 #: the state map for the files 

28 self.states = states 

29 #: The data stack created from the input files. 

30 self.img_stack = np.array(img_stack) 

31 #: contains the result in a 4 x (img shape) np.array 

32 self.stokes = Fits(fname) 

33 

34 def run(self) -> Fits: 

35 """ 

36 Executes the processing 

37 

38 :return: the stokes images as a 4 x (img shape) np.array 

39 """ 

40 self._adjust_demod_matrix() 

41 self._compute_stokes() 

42 self._normalize() 

43 return self.stokes 

44 

45 def _adjust_demod_matrix(self) -> None: 

46 if len(self.states) == Globals.MOD_CYCLE_FRAMES: 

47 return 

48 

49 logger.debug('Adjusting demodulation matrix for states %s', self.states) 

50 self.demod_matrix = np.transpose(np.transpose(self.demod_matrix)[self.states]) 

51 

52 def _compute_stokes(self) -> None: 

53 data = np.transpose(self.img_stack, axes=(1, 2, 0))[:, :, :, None] 

54 stokes = np.matmul(self.demod_matrix, data, dtype=np.float64) 

55 self.stokes.set_data(np.transpose(stokes[:, :, :, 0], axes=(2, 0, 1))) 

56 

57 def _normalize(self) -> None: 

58 for s in range(1, 4): 

59 self.stokes.data[s] = Flat.save_divide(self.stokes.data[s], self.stokes.data[0])