Coverage for src/spectroflat/sensor/artificial_flat.py: 90%

1from __future__ import annotations 

2 

3import copy 

4 

5import numpy as np 

6 

7from ..shift.img_rotation import RotationCorrection 

8from ..smile import OffsetMap, SmileInterpolator 

9from ..utils.processing import MP 

10 

11 

12class ArtificialFlat: 

13 

14 def __init__(self, img: np.array, roi: tuple = None): 

15 self._roi = roi if roi is not None else tuple([slice(0, s) for s in img.shape]) 

16 self.artificial = np.ones(img.shape) 

17 self._input = img 

18 

19 def create(self) -> ArtificialFlat: 

20 self._expand_vertical_average() 

21 return self 

22 

23 def _expand_vertical_average(self): 

24 average = np.mean(self._input[self._roi], axis=1) 

25 states, rows, _ = self.artificial.shape 

26 for s in range(states): 

27 for r in range(rows): 

28 self.artificial[s, r, self._roi[2]] = average[s] 

29 

30 def resmile(self, offsets: OffsetMap) -> ArtificialFlat: 

31 b = int(np.max(np.abs(offsets.map))) + 1 

32 rot = - float(offsets.header['Rotation']) 

33 omap = copy.deepcopy(offsets) 

34 omap.header['Rotation'] = 0 

35 args = [] 

36 for s in range(self.artificial.shape[0]): 

37 omap.map[s] *= -1 

38 args.append((omap, self.artificial[s], s)) 

39 result = dict(MP.simultaneous(SmileInterpolator.desmile_state, args)) 

40 img = [RotationCorrection(result[s], rot).bicubic() for s in range(self.artificial.shape[0])] 

41 self.artificial = np.array(img) 

42 self.artificial[self._roi][:, :, 0:b] = self._input[self._roi][:, :, 0:b] 

43 self.artificial[self._roi][:, :, -b:] = self._input[self._roi][:, :, -b:] 

44 return self 

45 

46 def pad(self, shape: tuple): 

47 y = (shape[1] - self.artificial.shape[1]) // 2 

48 x = (shape[2] - self.artificial.shape[2]) // 2 

49 self.artificial = np.pad(self.artificial, ((0, 0), (y, y), (x, x)), 'constant', constant_values=1) 

50 return self 

51 

52 def remove(self, img: np.array): 

53 return np.true_divide(img.astype('float32'), self.artificial.astype('float32'), 

54 out=img.astype('float32'), where=self.artificial != 0, dtype='float64')