Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np 

2import numbers 

3 

4__all__ = ["_FactorialGenerator", "ZernikePolynomialGenerator"] 

5 

6 

7class _FactorialGenerator(object): 

8 """ 

9 A class that generates factorials 

10 and stores them in a dict to be referenced 

11 as needed. 

12 """ 

13 

14 def __init__(self): 

15 self._values = {0:1, 1:1} 

16 self._max_i = 1 

17 

18 def evaluate(self, num): 

19 """ 

20 Return the factorial of num 

21 """ 

22 if num<0: 

23 raise RuntimeError("Cannot handle negative factorial") 

24 

25 i_num = int(np.round(num)); 

26 if i_num in self._values: 

27 return self._values[i_num] 

28 

29 val = self._values[self._max_i] 

30 for ii in range(self._max_i, num): 

31 val *= (ii+1) 

32 self._values[ii+1] = val 

33 

34 self._max_i = num 

35 return self._values[num] 

36 

37 

38class ZernikePolynomialGenerator(object): 

39 """ 

40 A class to generate and evaluate the Zernike 

41 polynomials. Definitions of Zernike polynomials 

42 are taken from 

43 https://en.wikipedia.org/wiki/Zernike_polynomials 

44 """ 

45 

46 def __init__(self): 

47 self._factorial = _FactorialGenerator() 

48 self._coeffs = {} 

49 self._powers = {} 

50 

51 def _validate_nm(self, n, m): 

52 """ 

53 Make sure that n, m are a valid pair of indices for 

54 a Zernike polynomial. 

55 

56 n is the radial order 

57 

58 m is the angular order 

59 """ 

60 if not isinstance(n, int) and not isinstance(n, np.int64): 

61 raise RuntimeError('Zernike polynomial n must be int') 

62 if not isinstance(m,int) and not isinstance(m, np.int64): 

63 raise RuntimeError('Zernike polynomial m must be int') 

64 

65 if n<0: 

66 raise RuntimeError('Radial Zernike n cannot be negative') 

67 if m<0: 

68 raise RuntimeError('Radial Zernike m cannot be negative') 

69 if n<m: 

70 raise RuntimeError('Radial Zerniki n must be >= m') 

71 

72 n = int(n) 

73 m = int(m) 

74 

75 return (n, m) 

76 

77 def _make_polynomial(self, n, m): 

78 """ 

79 Make the radial part of the n, m Zernike 

80 polynomial. 

81 

82 n is the radial order 

83 

84 m is the angular order 

85 

86 Returns 2 numpy arrays: coeffs and powers. 

87 

88 The radial part of the Zernike polynomial is 

89 

90 sum([coeffs[ii]*power(r, powers[ii]) 

91 for ii in range(len(coeffs))]) 

92 """ 

93 

94 n, m = self._validate_nm(n, m) 

95 

96 # coefficients taken from 

97 # https://en.wikipedia.org/wiki/Zernike_polynomials 

98 

99 n_coeffs = 1+(n-m)//2 

100 local_coeffs = np.zeros(n_coeffs, dtype=float) 

101 local_powers = np.zeros(n_coeffs, dtype=float) 

102 for k in range(0, n_coeffs): 

103 if k%2 == 0: 

104 sgn = 1.0 

105 else: 

106 sgn = -1.0 

107 

108 num_fac = self._factorial.evaluate(n-k) 

109 k_fac = self._factorial.evaluate(k) 

110 d1_fac = self._factorial.evaluate(((n+m)//2)-k) 

111 d2_fac = self._factorial.evaluate(((n-m)//2)-k) 

112 

113 local_coeffs[k] = sgn*num_fac/(k_fac*d1_fac*d2_fac) 

114 local_powers[k] = n-2*k 

115 

116 self._coeffs[(n,m)] = local_coeffs 

117 self._powers[(n,m)] = local_powers 

118 

119 def _evaluate_radial_number(self, r, nm_tuple): 

120 """ 

121 Evaluate the radial part of a Zernike polynomial. 

122 

123 r is a scalar value 

124 

125 nm_tuple is a tuple of the form (radial order, angular order) 

126 denoting the polynomial to evaluate 

127 

128 Return the value of the radial part of the polynomial at r 

129 """ 

130 if r > 1.0: 

131 return np.NaN 

132 

133 r_term = np.power(r, self._powers[nm_tuple]) 

134 return (self._coeffs[nm_tuple]*r_term).sum() 

135 

136 def _evaluate_radial_array(self, r, nm_tuple): 

137 """ 

138 Evaluate the radial part of a Zernike polynomial. 

139 

140 r is a numpy array of radial values 

141 

142 nm_tuple is a tuple of the form (radial order, angular order) 

143 denoting the polynomial to evaluate 

144 

145 Return the values of the radial part of the polynomial at r 

146 (returns np.NaN if r>1.0) 

147 """ 

148 if len(r) == 0: 

149 return np.array([],dtype=float) 

150 

151 # since we use np.where to handle cases of 

152 # r==0, use np.errstate to temporarily 

153 # turn off the divide by zero and 

154 # invalid double scalar RuntimeWarnings 

155 with np.errstate(divide='ignore', invalid='ignore'): 

156 log_r = np.log(r) 

157 log_r = np.where(np.isfinite(log_r), log_r, -1.0e10) 

158 r_power = np.exp(np.outer(log_r, self._powers[nm_tuple])) 

159 

160 results = np.dot(r_power, self._coeffs[nm_tuple]) 

161 return np.where(r<1.0, results, np.NaN) 

162 

163 def _evaluate_radial(self, r, n, m): 

164 """ 

165 Evaluate the radial part of a Zernike polynomial 

166 

167 r is a radial value or an array of radial values 

168 

169 n is the radial order of the polynomial 

170 

171 m is the angular order of the polynomial 

172 

173 Return the value(s) of the radial part of the polynomial at r 

174 (returns np.NaN if r>1.0) 

175 """ 

176 

177 is_array = False 

178 if not isinstance(r, numbers.Number): 

179 is_array = True 

180 

181 nm_tuple = self._validate_nm(n,m) 

182 

183 if (nm_tuple[0]-nm_tuple[1]) % 2 == 1: 

184 if is_array: 

185 return np.zeros(len(r), dtype=float) 

186 return 0.0 

187 

188 if nm_tuple not in self._coeffs: 

189 self._make_polynomial(nm_tuple[0], nm_tuple[1]) 

190 

191 if is_array: 

192 return self._evaluate_radial_array(r, nm_tuple) 

193 

194 return self._evaluate_radial_number(r, nm_tuple) 

195 

196 def evaluate(self, r, phi, n, m): 

197 """ 

198 Evaluate a Zernike polynomial in polar coordinates 

199 

200 r is the radial coordinate (a scalar or an array) 

201 

202 phi is the angular coordinate in radians (a scalar or an array) 

203 

204 n is the radial order of the polynomial 

205 

206 m is the angular order of the polynomial 

207 

208 Return the value(s) of the polynomial at r, phi 

209 (returns np.NaN if r>1.0) 

210 """ 

211 radial_part = self._evaluate_radial(r, n, np.abs(m)) 

212 if m>=0: 

213 return radial_part*np.cos(m*phi) 

214 return radial_part*np.sin(m*phi) 

215 

216 def norm(self, n, m): 

217 """ 

218 Return the normalization of the n, m Zernike 

219 polynomial 

220 

221 n is the radial order 

222 

223 m is the angular order 

224 """ 

225 nm_tuple = self._validate_nm(n, np.abs(m)) 

226 if nm_tuple[1] == 0: 

227 eps = 2.0 

228 else: 

229 eps = 1.0 

230 return eps*np.pi/(nm_tuple[0]*2+2) 

231 

232 def evaluate_xy(self, x, y, n, m): 

233 """ 

234 Evaluate a Zernike polynomial at a point in 

235 Cartesian space. 

236 

237 x and y are the Cartesian coordinaes (either scalars 

238 or arrays) 

239 

240 n is the radial order of the polynomial 

241 

242 m is the angular order of the polynomial 

243 

244 Return the value(s) of the polynomial at x, y 

245 (returns np.NaN if sqrt(x**2+y**2)>1.0) 

246 """ 

247 # since we use np.where to handle r==0 cases, 

248 # use np.errstate to temporarily turn off the 

249 # divide by zero and invalid double scalar 

250 # RuntimeWarnings 

251 with np.errstate(divide='ignore', invalid='ignore'): 

252 r = np.sqrt(x**2+y**2) 

253 cos_phi = np.where(r>0.0, x/r, 0.0) 

254 arccos_phi = np.arccos(cos_phi) 

255 phi = np.where(y>=0.0, arccos_phi, 0.0-arccos_phi) 

256 return self.evaluate(r, phi, n, m)