recursive-zernike2.jl

r.jl

"""
	calcac/b(n, α, β, x)

Two functions calcac and calcb compute the three coefficients in the recurrence
relation for the Jacobi polynomial.  It is broken into two functions to separate
a and c, which do not depend on x and can be calculated once for a vector of x,
from b which does depend on x and must be calculated for each element.

n is the order of the polynomial.

α,β are the two weighting parameters.

x is the argument or point at which the polynomial will be evaluated.
"""
function calcac(n, α, β)
    a = (2n) * (n + α + β) * (2n + α + β - 2)
    c = 2 * (n + α - 1) * (n + β - 1) * (2n + α + β)
    return a, c
end

function calcb(n, α, β, x)
	# the parens emphasize the terms, not PEMDAS
	b1 = (2n + α + β -1)
	b2 = (2n + α + β)
	b2 = b2 * (b2 - 2)
	b = b1 * (b2 * x + α^2 - β^2)
	return b
end

"""
    jacobi(n, α, β, x)

Compute the Jacobi polynomial of order n and weights α,β at point x.

This function uses a recurrence relation and is numerical stable to very high
order.  The computation time is linear w.r.t. the order n.  jacobi.(n, a, b, x)
is also linear w.r.t. the size of argument x.
"""


function jacobi(n, α, β, x)
	if n == 0
		return ones(size(r))
	elseif n == 1
		return (α + 1) .+ (α + β + 2) .* ((x.-1)./2)
	end
	# this uses a loop to avoid recursion
	# we init P of n-2, n-1, and n=2.
	# the latter is the first recursive term.
	# then loop from 3 up to n, updating
	# Pnm2, Pnm1, a, b, c, Pn
	Pnm2 = 1.
	Pnm1 = (α + 1) .+ (α + β + 2) .* ((x.-1)./2)
    a, c = calcac(2, α, β)
	Pn = ((calcb.(2, α, β, x) .* Pnm1) .- (c .* Pnm2)) ./ a
	if n == 2
		return Pn
	end
	for i = 3:n
		Pnm2 = Pnm1
		Pnm1 = Pn
        a, c = calcac(n, α, β)
		Pn = ((calcb.(n, α, β, x) .* Pnm1) .- (c .* Pnm2)) ./ a
	end
	return Pn
end

"""
    kronecker(i,j)

1 if i==j, else 0; mathematical kronecker function
"""
function kronecker(i,j)
    return i==j ? 1 : 0
end

"""
    zernike_norm(n, m)

Norm of Zernike polynomial of radial order n, azimuthal order m.

The norm is the average squared distance to zero.  By multiplying a zernike
value by the norm, the term is given unit stdev or RMS.
"""
function zernike_norm(n, m)
    num = √(2 * (n+1)) / (1 + kronecker(m, 0))
end

"""
    zernike(n, m, ρ, θ[; norm])

Zernike polynomial of radial order n and azimuthal order m, evaluated at the
point (ρ, θ).  No normalization is required of (ρ, θ), though the polynomials
are orthogonal only over the unit disk.

norm is a boolean flag indicating whether the result should be orthonormalized
(scaled to unit RMS) or not.
"""
function zernike(n, m, ρ, θ; norm::Bool=true)
    x = ρ.^2 .- 1
    n_j = (n - m) ÷ 2
    am = abs(m)
    # α=0, β=|m|
    # there is a second syntax where you have x reversed, 1 - ρ^2,
    # in which ase you swap α and β.  It makes no difference
    out = jacobi(n_j, 0, am, x)
    if m != 0
        if m < 0
            out *= (ρ.^am .* sin(m*θ))
        else
            out *= (ρ.^am .* cos(m*θ))
        end

    end
	if norm
		out *= zernike_norm(n,m)
	end
    return out
end