Jashcraf / poke

Poke (pronounced poh-keh) is a Polarization Ray Tracing and Gaussian Beamlet module for Python
BSD 3-Clause "New" or "Revised" License
32 stars 6 forks source link

Zernike polynomial generator #65

Closed emoryljenkins closed 10 months ago

emoryljenkins commented 1 year ago

This little script generates the values of the Zernike polynomials for a set of rho and phi coordinate pairs, up to a given ANSI index of the Zernike set.

import numpy as np
# import matplotlib.pyplot as plt

x = np.linspace(-1,1,50)
x,y = np.meshgrid(x,x)
rho = np.sqrt(x**2 + y**2)
phi = np.arctan2(y,x)

# rho=np.linspace(0, 1, 50)
# phi=np.linspace(0, 2*np.pi, 50)
# rho, phi = np.meshgrid(rho, phi)

rho = rho.ravel()
phi = phi.ravel()

def zernike(rho, phi, J):
    N=int(np.ceil(np.sqrt(2*J + 0.25)-0.5)) # N = number of rows on the zernike pyramid
    values = np.zeros([rho.size, J+1])
    j=0 # ANSI index of zernike
    for n in range(0,N):
        m=-n
        while m<=n:
            R = 0
            for k in range(0,1+(n-abs(m))//2):
                c_k = ((-1)**k * np.math.factorial(n-k))/(np.math.factorial(k) * np.math.factorial((n+abs(m))//2 - k) * np.math.factorial((n-abs(m))//2 - k))
                R += c_k * rho**(n-2*k)
            if m>0:
                Z = R*np.cos(m*phi)
            elif m<0:
                Z = R*np.sin(abs(m)*phi)
            else:
                Z = R
            values[:,j] = Z
            j=j+1
            if j>J:
                break
            m=m+2
    return values

# vals = zernike(rho, phi, 0)
Jashcraf commented 10 months ago

Epic and merged a while ago, also put Emory on the SPIE proceedings