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

1import numpy as np 

2from ...io import Fits 

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

4 

5 

6class FrameSum: 

7 """ 

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

9 """ 

10 

11 def __init__(self, axis: int = 0): 

12 self.header = None 

13 self.result = None 

14 self.axis = axis 

15 

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

17 f = Fits(file).read() 

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

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

20 if self.result is None: 

21 self.result = f.data 

22 self.header = f.header 

23 return 

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

25 

26 @staticmethod 

27 def sum_all(files: list, axis: int = 0): 

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

29 fs = FrameSum(axis) 

30 for f in files: 

31 fs.add(f) 

32 progress.dot() 

33 return fs 

34 

35 @staticmethod 

36 def sum_in_chunks(files: list, chunk_size=200): 

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

38 frame_sums = MP.simultaneous(FrameSum.sum_all, Collections.chunker(files, chunk_size)) 

39 result = frame_sums[0] 

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

41 progress.dot(flush=True) 

42 return result