"""Math functions for computing Zernike polynomial information.
TODO: write tests and document these functions properly
"""
import numpy as np
import scipy
[docs]def Rnm_p(n, m):
"""Generate radial polynomial for radial Zernikee."""
pc = []
m = abs(m)
for ik in range(int((n - m) / 2) + 1):
num = (-1) ** ik * scipy.special.factorial(int(n - ik))
den = (
scipy.special.factorial(int(ik))
* scipy.special.factorial(int((n + m) / 2 - ik))
* scipy.special.factorial(int((n - m) / 2 - ik))
)
pc.append(num / den)
return np.array(pc)
[docs]def Rnm_eval(_r, _phi, n, m, a0):
"""Function to evaluate radial components."""
# Obtain the polynomial coeffs:
pn = Rnm_p(n, m)
Rnm = np.zeros(_r.shape)
for idx, ip in enumerate(pn):
Rnm += ip * (_r / a0) ** (n - 2 * idx)
# Noll normalissation:
Nnm = np.sqrt(2 * n + 2)
return Nnm * Rnm
[docs]def ZPhi_eval(_phi, m):
"""Fuction to generate azimuthal component."""
if m < 0:
ZPhi = 1 / np.sqrt(np.pi) * np.sin(m * _phi)
elif m > 0:
ZPhi = 1 / np.sqrt(np.pi) * np.cos(m * _phi)
else:
ZPhi = 1 / np.sqrt(2 * np.pi) * np.ones(_phi.shape)
return ZPhi
[docs]def Znm_eval(_r, _phi, n, m, a0):
_Rnm = Rnm_eval(_r, _phi, n, m, a0)
_Pnm = ZPhi_eval(_phi, m)
return _Rnm * _Pnm
[docs]def Gen_nm(n):
"""Generate n and m vectors containing n and m indices up to n, excluding zeroth
mode."""
vlen = np.sum(np.arange(2, n + 2))
_n = np.zeros(vlen)
_m = np.zeros(vlen)
for iN in range(1, n + 1):
iStart = np.sum(np.arange(2, iN - 1 + 2))
iStop = np.sum(np.arange(2, iN + 2))
_n[iStart:iStop] = iN
for im in range(iN + 1):
_m[iStart + im] = -iN + im * 2
return _n, _m