Coverage for python/lsst/sims/utils/ZernikeModule.py : 12%

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
4__all__ = ["_FactorialGenerator", "ZernikePolynomialGenerator"]
7class _FactorialGenerator(object):
8 """
9 A class that generates factorials
10 and stores them in a dict to be referenced
11 as needed.
12 """
14 def __init__(self):
15 self._values = {0:1, 1:1}
16 self._max_i = 1
18 def evaluate(self, num):
19 """
20 Return the factorial of num
21 """
22 if num<0:
23 raise RuntimeError("Cannot handle negative factorial")
25 i_num = int(np.round(num));
26 if i_num in self._values:
27 return self._values[i_num]
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
34 self._max_i = num
35 return self._values[num]
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 """
46 def __init__(self):
47 self._factorial = _FactorialGenerator()
48 self._coeffs = {}
49 self._powers = {}
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.
56 n is the radial order
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')
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')
72 n = int(n)
73 m = int(m)
75 return (n, m)
77 def _make_polynomial(self, n, m):
78 """
79 Make the radial part of the n, m Zernike
80 polynomial.
82 n is the radial order
84 m is the angular order
86 Returns 2 numpy arrays: coeffs and powers.
88 The radial part of the Zernike polynomial is
90 sum([coeffs[ii]*power(r, powers[ii])
91 for ii in range(len(coeffs))])
92 """
94 n, m = self._validate_nm(n, m)
96 # coefficients taken from
97 # https://en.wikipedia.org/wiki/Zernike_polynomials
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
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)
113 local_coeffs[k] = sgn*num_fac/(k_fac*d1_fac*d2_fac)
114 local_powers[k] = n-2*k
116 self._coeffs[(n,m)] = local_coeffs
117 self._powers[(n,m)] = local_powers
119 def _evaluate_radial_number(self, r, nm_tuple):
120 """
121 Evaluate the radial part of a Zernike polynomial.
123 r is a scalar value
125 nm_tuple is a tuple of the form (radial order, angular order)
126 denoting the polynomial to evaluate
128 Return the value of the radial part of the polynomial at r
129 """
130 if r > 1.0:
131 return np.NaN
133 r_term = np.power(r, self._powers[nm_tuple])
134 return (self._coeffs[nm_tuple]*r_term).sum()
136 def _evaluate_radial_array(self, r, nm_tuple):
137 """
138 Evaluate the radial part of a Zernike polynomial.
140 r is a numpy array of radial values
142 nm_tuple is a tuple of the form (radial order, angular order)
143 denoting the polynomial to evaluate
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)
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]))
160 results = np.dot(r_power, self._coeffs[nm_tuple])
161 return np.where(r<1.0, results, np.NaN)
163 def _evaluate_radial(self, r, n, m):
164 """
165 Evaluate the radial part of a Zernike polynomial
167 r is a radial value or an array of radial values
169 n is the radial order of the polynomial
171 m is the angular order of the polynomial
173 Return the value(s) of the radial part of the polynomial at r
174 (returns np.NaN if r>1.0)
175 """
177 is_array = False
178 if not isinstance(r, numbers.Number):
179 is_array = True
181 nm_tuple = self._validate_nm(n,m)
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
188 if nm_tuple not in self._coeffs:
189 self._make_polynomial(nm_tuple[0], nm_tuple[1])
191 if is_array:
192 return self._evaluate_radial_array(r, nm_tuple)
194 return self._evaluate_radial_number(r, nm_tuple)
196 def evaluate(self, r, phi, n, m):
197 """
198 Evaluate a Zernike polynomial in polar coordinates
200 r is the radial coordinate (a scalar or an array)
202 phi is the angular coordinate in radians (a scalar or an array)
204 n is the radial order of the polynomial
206 m is the angular order of the polynomial
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)
216 def norm(self, n, m):
217 """
218 Return the normalization of the n, m Zernike
219 polynomial
221 n is the radial order
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)
232 def evaluate_xy(self, x, y, n, m):
233 """
234 Evaluate a Zernike polynomial at a point in
235 Cartesian space.
237 x and y are the Cartesian coordinaes (either scalars
238 or arrays)
240 n is the radial order of the polynomial
242 m is the angular order of the polynomial
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)