"""Math functions for computing Zernike polynomial information.
TODO: write tests and document these functions properly
"""
import numpy as np
[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 *np.math.factorial(int(n-ik))
den = np.math.factorial(int(ik))*np.math.factorial(int((n+m)/2-ik))*np.math.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