Coverage for src/susi/reduc/poincare_rotation/poincare.py: 47%

1#!/usr/bin/env python3 

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

3""" 

4Provides Poincare rotation for Stokes images 

5 

6@author: iglesias 

7""" 

8 

9import numpy as np 

10import matplotlib as mpl 

11import matplotlib.pyplot as plt 

12import logging 

13 

14log = logging.getLogger('SUSI') 

15 

16 

17class Poincare: 

18 """ 

19 Rotates the input stokes images in the Poincare sphere. 

20 

21 """ 

22 

23 def __init__(self, array, rm_offset=False, off_deg=1): 

24 """ 

25 :param array: numpy array [[I0,q0,u0,v0,I1,q1,u1,v1,...],x,y] (note that Stokes q,u and v are normalized) 

26 :param rm_offset: If True (default False) the mean value of q, u and v is forced to 0. 

27 :param off_deg: 0 means the image average is substracted per Stokes. 

28 1 means a linear fit to dim 0 avg is fit per stokes 

29 """ 

30 self.data = array 

31 self.result = None 

32 self.rm_offset = rm_offset 

33 self.off_deg = off_deg 

34 

35 def run(self, ang): 

36 """ 

37 Applies the poincare rotation in q and u direction. 

38 :param ang: Tuple with the rotation angles [deg] in Stokes q and u direction (q_ang, u_ang) 

39 :return: Numpy array with rotated images 

40 """ 

41 self.result = np.array(self.data) 

42 if ang != (0, 0): 

43 sz = np.shape(self.result) 

44 nstok = sz[0] // 4 

45 self.result = np.reshape(self.result, (nstok, 4, sz[1], sz[2])) 

46 self.result = np.transpose(self.result, (0, 2, 3, 1)) 

47 crossq1 = np.deg2rad(ang[0]) 

48 crossu1 = np.deg2rad(ang[1]) 

49 cq = np.cos(crossq1) 

50 sq = np.sin(crossq1) 

51 cu = np.cos(crossu1) 

52 su = np.sin(crossu1) 

53 rot_m = np.array([[1, 0, 0, 0], [0, cu, sq * su, cq * su], [0, 0, cq, -sq], [0, -su, cu * sq, cq * cu]]) 

54 self.result = rot_m[None, None, None, :, :] @ self.result[:, :, :, :, None] 

55 self.result = self.result[:, :, :, :, 0] 

56 self.result = np.transpose(self.result, (0, 3, 1, 2)) 

57 self.result = np.reshape(self.result, (sz[0], sz[1], sz[2])) 

58 if self.rm_offset: 

59 for i in range(nstok): 

60 self.result[1 + i * 4, :, :] = self.remove_off(self.result[1 + i * 4, :, :]) 

61 self.result[2 + i * 4, :, :] = self.remove_off(self.result[2 + i * 4, :, :]) 

62 self.result[3 + i * 4, :, :] = self.remove_off(self.result[3 + i * 4, :, :]) 

63 return self.result 

64 

65 def remove_off(self, im): 

66 """ 

67 Removes the offset of the image 

68 :param img: image to remove the offset 

69 :return: image with offset removed 

70 """ 

71 if self.off_deg == 0: 

72 return im - np.mean(im) 

73 elif self.off_deg == 1: 

74 im_y_mean = np.mean(im, axis=0) 

75 fit_x = np.arange(len(im_y_mean)) 

76 fit_m, fit_b = np.polyfit(fit_x, im_y_mean, 1) 

77 im_x_fit = fit_m * fit_x + fit_b 

78 im -= np.repeat(im_x_fit[None, :], np.shape(im)[0], axis=0) 

79 return im 

80 else: 

81 log.error('Unknown offset degree %d, must be 0 or 1', self.off_deg) 

82 raise ValueError('Unknown offset degree %d, must be 0 or 1' % self.off_deg) 

83 

84 def interactive(self, rng=(-3, 3), roi=None): 

85 """ 

86 Applies a poincare rotation in q and u direction interactively by ploting the rotated images. 

87 You must have a working X server! 

88 

89 param: rng: (min,max) color scale range in fractions of the image std dev 

90 param: roi: (slice(ymin,ymax),slice(xmin,xmax)); ROI for the plots, default all image 

91 return: array, angles Rotated array and rotationa angles 

92 """ 

93 bend = mpl.get_backend() 

94 mpl.use('TkAgg') 

95 numfmt = '{:8.4e}' # format string to print sci notation numbers 

96 stks_names = ['I', 'q', 'u', 'v'] 

97 

98 loop = True 

99 self.run((0, 0)) 

100 while loop: 

101 if roi is not None: 

102 toplot = self.result[:, roi[0], roi[1]] 

103 else: 

104 toplot = self.result 

105 fig = plt.figure(figsize=(18, 8)) 

106 # imgs 

107 for i in range(4): 

108 fig.add_subplot(2, 4, i + 1) 

109 m = np.nanmean(toplot[i, :, :]) 

110 sd = np.nanstd(toplot[i, :, :]) 

111 tit = 'Stokes ' + stks_names[i] + ' \n mean=' + numfmt.format(m) + '; std=' + numfmt.format(sd) 

112 im = plt.imshow(toplot[i, :, :], cmap='Greys_r', vmin=m + rng[0] * sd, vmax=m + rng[1] * sd) 

113 plt.title(tit) 

114 # cuts 

115 for i in range(4): 

116 fig.add_subplot(2, 4, i + 5) 

117 mean = np.mean(toplot[i, :, :], axis=0) 

118 m = np.nanmean(mean) 

119 sd = np.nanstd(mean) 

120 tit = 'Stokes ' + stks_names[i] + ' \n mean=' + numfmt.format(m) + '; std=' + numfmt.format(sd) 

121 im = plt.plot(mean) # , vmin=m + rng[0] * sd, vmax=m + rng[1] * sd) 

122 plt.title(tit) 

123 

124 plt.show(block=False) 

125 ang = input('Enter a new pair (q, u) of rotation angles (e.g., -2.56, 10.0) use 360,360 to exit:') 

126 ang = tuple(map(float, ang.split(','))) 

127 plt.close() 

128 print('Rotating with angles...%s', ang) 

129 if ang == (360, 360): 

130 loop = False 

131 else: 

132 oang = ang 

133 self.run(ang) 

134 

135 mpl.use(bend) 

136 return self.result, oang