def test_as_list(self):
obj = 1
self.assertEqual(utils.as_list(obj), [obj])
list_obj = [1, 2, 3]
self.assertEqual(utils.as_list(list_obj), list_obj)
def get(self,
pupil,
rho=None,
theta=None,
n=None,
m=None,
coefficient=None,
**kwargs):
m_list = as_list(m)
n_list = as_list(n)
coefficients = as_list(coefficient)
assert len(m_list) == len(
n_list), "The number of indices need to match"
assert len(m_list) == len(
coefficients
), "The number of indices need to match the number of coefficients"
pupil_bool = pupil != 0
rho = rho[pupil_bool]
theta = theta[pupil_bool]
Z = 0
for n, m, coefficient in zip(n_list, m_list, coefficients):
if (n - m) % 2 or coefficient == 0:
continue
R = 0
for k in range((n - np.abs(m)) // 2 + 1):
R += ((-1)**k * np.math.factorial(n - k) /
(np.math.factorial(k) *
np.math.factorial((n - m) // 2 - k) *
np.math.factorial((n + m) // 2 - k)) * rho**(n - 2 * k))
if m > 0:
R = R * np.cos(m * theta) * (np.sqrt(2 * n + 2) * coefficient)
elif m < 0:
R = R * np.sin(-m * theta) * (np.sqrt(2 * n + 2) * coefficient)
else:
R = R * (np.sqrt(n + 1) * coefficient)
Z += R
phase = np.exp(1j * Z)
pupil[pupil_bool] *= phase
return pupil