Coverage for src/susi/reduc/average/framesum.py: 72%

1from tkinter import N 

2import numpy as np 

3from ...io import Fits 

4from ...utils import Collections, MP, progress 

5 

6 

7class FrameSum: 

8 """ 

9 Helper class to sum data arrays of fits files along a given axis 

10 """ 

11 

12 def __init__(self, axis: int = 0, slices=None): 

13 self.header = None 

14 self.result = None 

15 self.axis = axis 

16 self.slices = slices 

17 

18 def add(self, file: str) -> None: 

19 f = Fits(file).read(slices=self.slices) 

20 if len(f.data.shape) < 3: 

21 f.data = np.array([f.data]) 

22 if self.result is None: 

23 self.result = f.data 

24 self.header = f.header 

25 return 

26 self.result = np.sum([self.result, f.data], axis=self.axis) 

27 

28 @staticmethod 

29 def sum_all(args): 

30 """Sums up all the data entries of the fits files in the list.""" 

31 files, axis, slices = args 

32 fs = FrameSum(axis, slices=slices) 

33 for f in files: 

34 fs.add(f) 

35 progress.dot() 

36 return fs 

37 

38 @staticmethod 

39 def sum_in_chunks(files: list, chunk_size=200, slices=None, header_idx=None): 

40 """ 

41 Sums up all data of the fits files in the list in chunks. 

42 param files: list of fits files 

43 param chunk_size: number of files to process in one chunk 

44 param slices: slices to read the data, see `Fits.read` for details 

45 param header_idx: header of this file will be used as the header of the result. 

46 If None, the header of the first chunk finished (random) is be used. 

47 """ 

48 chunks = Collections.chunker(files, chunk_size) 

49 args = [(c, 0, slices) for c in chunks] 

50 frame_sums = MP.simultaneous(FrameSum.sum_all, args) 

51 result = frame_sums[0] 

52 if header_idx is not None: 

53 result.header = Fits(files[header_idx]).read(header_only=True).header 

54 result.result = np.sum([fs.result for fs in frame_sums], axis=0) 

55 progress.dot(flush=True) 

56 return result 

57 

58 @staticmethod 

59 def avg_in_chunks(files: list, chunk_size=200, slices=None, header_idx=None): 

60 """ 

61 Averages the data of the fits files in the list in chunks. 

62 param files: list of fits files 

63 param chunk_size: number of files to process in one chunk 

64 param slices: slices to read the data, see `Fits.read` for details 

65 param header_idx: header of this file will be used as the header of the result. 

66 If None, the header of the first chunk finished (random) is be used. 

67 """ 

68 fs = FrameSum.sum_in_chunks(files, chunk_size, slices, header_idx) 

69 fs.result /= len(files) 

70 return fs