def surface_normal_from_cartesian_derivatives(fx, fy, r, t):
"""Use Cartesian derivatives to compute polar surface normals.
Parameters
----------
fx : numpy.ndarray
derivative of f w.r.t. x
fy : numpy.ndarray
derivative of f w.r.t. y
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
Returns
-------
numpy.ndarray, numpy.ndarray
r, t derivatives; will contain a singularity where r=0,
see fix_zero_singularity
"""
cost = np.cos(t)
sint = np.sin(t)
onebyr = 1 / r
r = fx * cost + fy * sint
t = fx * -sint / onebyr + fy * cost / onebyr
return r, t
def surface_normal_from_cylindrical_derivatives(fp, ft, r, t):
"""Use polar derivatives to compute Cartesian surface normals.
Parameters
----------
fp : numpy.ndarray
derivative of f w.r.t. r
ft : numpy.ndarray
derivative of f w.r.t. t
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
Returns
-------
numpy.ndarray, numpy.ndarray
x, y derivatives; will contain a singularity where r=0,
see fix_zero_singularity
"""
cost = np.cos(t)
sint = np.sin(t)
x = fp * cost - 1 / r * ft * sint
y = fp * sint + 1 / r * ft * cost
return x, y
def off_axis_conic_sigma(c, kappa, r, t, dx, dy=0):
"""Lowercase sigma (direction cosine projection term) for an off-axis conic.
See Eq. (5.2) of oe-20-3-2483.
Parameters
----------
c : float
axial curvature of the conic
kappa : float
conic constant
r : numpy.ndarray
radial coordinate, where r=0 is centered on the off-axis section
t : numpy.ndarray
azimuthal coordinate
dx : float
shift of the surface in x with respect to the base conic vertex,
mutually exclusive to dy (only one may be nonzero)
use dx=0 when dy != 0
dy : float
shift of the surface in y with respect to the base conic vertex
Returns
-------
sigma(r,t)
"""
if dy != 0 and dx != 0:
raise ValueError('only one of dx/dy may be nonzero')
if dx != 0:
s = dx
oblique_term = 2 * s * r * np.cos(t)
else:
s = dy
oblique_term = 2 * s * r * np.sin(t)
aggregate_term = r * r + oblique_term + s * s
csq = c * c
num = np.sqrt(1 - (1 + kappa) * csq * aggregate_term)
den = np.sqrt(1 - kappa * csq * aggregate_term) # flipped sign, 1-kappa
return num / den
def off_axis_conic_sag(c, kappa, r, t, dx, dy=0):
"""Sag of an off-axis conicoid.
Parameters
----------
c : float
axial curvature of the conic
kappa : float
conic constant
r : numpy.ndarray
radial coordinate, where r=0 is centered on the off-axis section
t : numpy.ndarray
azimuthal coordinate
dx : float
shift of the surface in x with respect to the base conic vertex,
mutually exclusive to dy (only one may be nonzero)
use dx=0 when dy != 0
dy : float
shift of the surface in y with respect to the base conic vertex
Returns
-------
numpy.ndarray
surface sag, z(x,y)
"""
if dy != 0 and dx != 0:
raise ValueError('only one of dx/dy may be nonzero')
if dx != 0:
s = dx
oblique_term = 2 * s * r * np.cos(t)
else:
s = dy
oblique_term = 2 * s * r * np.sin(t)
aggregate_term = r * r + oblique_term + s * s
num = c * aggregate_term
csq = c * c
den = 1 + np.sqrt(1 - (1 + kappa) * csq * aggregate_term)
return num / den
def hopkins(a, b, c, r, t, H):
"""Hopkins' aberration expansion.
This function uses the "W020" or "W131" like notation, with Wabc separating
into the a, b, c arguments. To produce a sine term instead of cosine,
make a the negative of the order. In other words, for W222S you would use
hopkins(2, 2, 2, ...) and for W222T you would use
hopkins(-2, 2, 2, ...).
Parameters
----------
a : int
azimuthal order
b : int
radial order
c : int
order in field ("H-order")
r : numpy.ndarray
radial pupil coordinate
t : numpy.ndarray
azimuthal pupil coordinate
H : numpy.ndarray
field coordinate
Returns
-------
numpy.ndarray
polynomial evaluated at this point
"""
# c = "component"
if a < 0:
c1 = np.sin(abs(a)*t)
else:
c1 = np.cos(a*t)
c2 = r ** b
c3 = H ** c
return c1 * c2 * c3
def zernike_nm(n, m, r, t, norm=True):
"""Zernike polynomial of radial order n, azimuthal order m at point r, t.
Parameters
----------
n : int
radial order
m : int
azimuthal order
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
norm : bool, optional
if True, orthonormalize the result (unit RMS)
else leave orthogonal (zero-to-peak = 1)
Returns
-------
numpy.ndarray
zernike mode of order n,m at points r,t
"""
x = 2 * r**2 - 1
am = abs(m)
n_j = (n - am) // 2
out = jacobi(n_j, 0, am, x)
if m != 0:
if m < 0:
out *= (r**am * np.sin(am * t))
else:
out *= (r**am * np.cos(m * t))
if norm:
out *= zernike_norm(n, m)
return out
def Q2d_sequence(nms, r, t):
"""Sequence of 2D-Q polynomials.
Parameters
----------
nms : iterable of tuple
(n,m) for each desired term
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
Returns
-------
generator
yields one term for each element of nms
"""
# see Q2d for general sense of this algorithm.
# the way this one works is to compute the maximum N for each |m|, and then
# compute the recurrence for each of those sequences and storing it. A loop
# is then iterated over the input nms, and selected value with appropriate
# prefixes / other terms yielded.
u = r
x = u**2
def factory():
return 0
# maps |m| => N
m_has_pos = set()
m_has_neg = set()
max_ns = defaultdict(factory)
for n, m in nms:
m_ = abs(m)
if max_ns[m_] < n:
max_ns[m_] = n
if m > 0:
m_has_pos.add(m_)
else:
m_has_neg.add(m_)
# precompute these reusable pieces of data
u_scales = {}
sin_scales = {}
cos_scales = {}
for absm in max_ns.keys():
u_scales[absm] = u**absm
if absm in m_has_neg:
sin_scales[absm] = np.sin(absm * t)
if absm in m_has_pos:
cos_scales[absm] = np.cos(absm * t)
sequences = {}
for m, N in max_ns.items():
if m == 0:
sequences[m] = list(Qbfs_sequence(range(N + 1), r))
else:
sequences[m] = []
P0 = 1 / 2
if m == 1:
P1 = 1 - x / 2
else:
P1 = (m - .5) + (1 - m) * x
f0 = f_q2d(0, m)
Q0 = 1 / (2 * f0)
sequences[m].append(Q0)
if N == 0:
continue
g0 = g_q2d(0, m)
f1 = f_q2d(1, m)
Q1 = (P1 - g0 * Q0) * (1 / f1)
sequences[m].append(Q1)
if N == 1:
continue
# everything above here works, or at least everything in the returns works
if m == 1:
P2 = (3 - x * (12 - 8 * x)) / 6
P3 = (5 - x * (60 - x * (120 - 64 * x))) / 10
g1 = g_q2d(1, m)
f2 = f_q2d(2, m)
Q2 = (P2 - g1 * Q1) * (1 / f2)
g2 = g_q2d(2, m)
f3 = f_q2d(3, m)
Q3 = (P3 - g2 * Q2) * (1 / f3)
sequences[m].append(Q2)
sequences[m].append(Q3)
# Q2, Q3 correct
if N <= 3:
continue
Pnm2, Pnm1 = P2, P3
Qnm1 = Q3
min_n = 4
else:
Pnm2, Pnm1 = P0, P1
Qnm1 = Q1
min_n = 2
for nn in range(min_n, N + 1):
A, B, C = abc_q2d(nn - 1, m)
Pn = (A + B * x) * Pnm1 - C * Pnm2
gnm1 = g_q2d(nn - 1, m)
fn = f_q2d(nn, m)
Qn = (Pn - gnm1 * Qnm1) * (1 / fn)
sequences[m].append(Qn)
Pnm2, Pnm1 = Pnm1, Pn
Qnm1 = Qn
for n, m in nms:
if m != 0:
if m < 0:
# m < 0, double neg = pos
prefix = sin_scales[-m] * u_scales[-m]
else:
prefix = cos_scales[m] * u_scales[m]
yield sequences[abs(m)][n] * prefix
else:
yield sequences[0][n]
def Q2d(n, m, r, t):
"""2D Q polynomial, aka the Forbes polynomials.
Parameters
----------
n : int
radial polynomial order
m : int
azimuthal polynomial order
r : numpy.ndarray
radial coordinate, slope orthogonal in [0,1]
t : numpy.ndarray
azimuthal coordinate, radians
Returns
-------
numpy.ndarray
array containing Q2d_n^m(r,t)
the leading coefficient u^m or u^2 (1 - u^2) and sines/cosines
are included in the return
"""
# Q polynomials have auxiliary polynomials "P"
# which are scaled jacobi polynomials under the change of variables
# x => 2x - 1 with alpha = -3/2, beta = m-3/2
# the scaling prefix may be found in A.4 of oe-20-3-2483
# impl notes:
# Pn is computed using a recurrence over order n. The recurrence is for
# a single value of m, and the 'seed' depends on both m and n.
#
# in general, Q_n^m = [P_n^m(x) - g_n-1^m Q_n-1^m] / f_n^m
# for the sake of consistency, this function takes args of (r,t)
# but the papers define an argument of u (really, u^2...)
# which is what I call rho (or r).
# for the sake of consistency of impl, I alias r=>u
# and compute x = u**2 to match the papers
u = r
x = u**2
if m == 0:
return Qbfs(n, r)
# m == 0 already was short circuited, so we only
# need to consider the m =/= 0 case for azimuthal terms
if sign(m) == -1:
m = abs(m)
prefix = u**m * np.sin(m * t)
else:
prefix = u**m * np.cos(m * t)
m = abs(m)
P0 = 1 / 2
if m == 1:
P1 = 1 - x / 2
else:
P1 = (m - .5) + (1 - m) * x
f0 = f_q2d(0, m)
Q0 = 1 / (2 * f0)
if n == 0:
return Q0 * prefix
g0 = g_q2d(0, m)
f1 = f_q2d(1, m)
Q1 = (P1 - g0 * Q0) * (1 / f1)
if n == 1:
return Q1 * prefix
# everything above here works, or at least everything in the returns works
if m == 1:
P2 = (3 - x * (12 - 8 * x)) / 6
P3 = (5 - x * (60 - x * (120 - 64 * x))) / 10
g1 = g_q2d(1, m)
f2 = f_q2d(2, m)
Q2 = (P2 - g1 * Q1) * (1 / f2)
g2 = g_q2d(2, m)
f3 = f_q2d(3, m)
Q3 = (P3 - g2 * Q2) * (1 / f3)
# Q2, Q3 correct
if n == 2:
return Q2 * prefix
elif n == 3:
return Q3 * prefix
Pnm2, Pnm1 = P2, P3
Qnm1 = Q3
min_n = 4
else:
Pnm2, Pnm1 = P0, P1
Qnm1 = Q1
min_n = 2
for nn in range(min_n, n + 1):
A, B, C = abc_q2d(nn - 1, m)
Pn = (A + B * x) * Pnm1 - C * Pnm2
gnm1 = g_q2d(nn - 1, m)
fn = f_q2d(nn, m)
Qn = (Pn - gnm1 * Qnm1) * (1 / fn)
Pnm2, Pnm1 = Pnm1, Pn
Qnm1 = Qn
# flake8 can't prove that the branches above the loop guarantee that we
# enter the loop and Qn is defined
return Qn * prefix # NOQA
def compute_z_zprime_Q2d(cm0, ams, bms, u, t):
"""Compute the surface sag and first radial and azimuthal derivative of a Q2D surface.
Excludes base sphere.
from Eq. 2.2 and Appendix B of oe-20-3-2483.
Parameters
----------
cm0 : iterable
surface coefficients when m=0 (inside curly brace, top line, Eq. B.1)
span n=0 .. len(cms)-1 and mus tbe fully dense
ams : iterable of iterables
ams[0] are the coefficients for the m=1 cosine terms,
ams[1] for the m=2 cosines, and so on. Same order n rules as cm0
bms : iterable of iterables
same as ams, but for the sine terms
ams and bms must be the same length - that is, if an azimuthal order m
is presnet in ams, it must be present in bms. The azimuthal orders
need not have equal radial expansions.
For example, if ams extends to m=3, then bms must reach m=3
but, if the ams for m=3 span n=0..5, it is OK for the bms to span n=0..3,
or any other value, even just [0].
u : numpy.ndarray
normalized radial coordinates (rho/rho_max)
t : numpy.ndarray
azimuthal coordinate, in the range [0, 2pi]
Returns
-------
numpy.ndarray, numpy.ndarray, numpy.ndarray
surface sag, radial derivative of sag, azimuthal derivative of sag
"""
usq = u * u
z = np.zeros_like(u)
dr = np.zeros_like(u)
dt = np.zeros_like(u)
# this is terrible, need to re-think this
if cm0 is not None and len(cm0) > 0:
zm0, zprimem0 = compute_z_zprime_Qbfs(cm0, u, usq)
z += zm0
dr += zprimem0
# B.1
# cos(mt)[sum a^m Q^m(u^2)] + sin(mt)[sum b^m Q^m(u^2)]
# ~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~
# variables: Sa Sb
# => because of am/bm going into Clenshaw's method, cannot
# simplify, need to do the recurrence twice
# u^m is outside the entire expression, think about that later
m = 0
# initialize to zero and incr at the front of the loop
# to avoid putting an m += 1 at the bottom (too far from init)
for a_coef, b_coef in zip(ams, bms):
m += 1
# TODO: consider zeroing alphas and re-using it to reduce
# alloc pressure inside this func; need care since len of any coef vector
# may be unequal
if len(a_coef) == 0:
continue
# can't use "as" => as keyword
Na = len(a_coef) - 1
Nb = len(b_coef) - 1
alphas_a = clenshaw_q2d_der(a_coef, m, usq)
alphas_b = clenshaw_q2d_der(b_coef, m, usq)
Sa = 0.5 * alphas_a[0][0]
Sb = 0.5 * alphas_b[0][0]
Sprimea = 0.5 * alphas_a[1][0]
Sprimeb = 0.5 * alphas_b[1][0]
if m == 1 and Na > 2:
Sa -= 2 / 5 * alphas_a[0][3]
# derivative is same, but instead of 0 index, index=j==1
Sprimea -= 2 / 5 * alphas_a[1][3]
if m == 1 and Nb > 2:
Sb -= 2 / 5 * alphas_b[0][3]
Sprimeb -= 2 / 5 * alphas_b[1][3]
um = u**m
cost = np.cos(m * t)
sint = np.sin(m * t)
kernel = cost * Sa + sint * Sb
total_sum = um * kernel
z += total_sum
# for the derivatives, we have two cases of the product rule:
# between "cost" and Sa, and between "sint" and "Sb"
# within each of those is a chain rule, just as for Zernike
# then there is a final product rule for the outer term
# differentiating in this way is just like for the classical asphere
# equation; differentiate each power separately
# if F(x) = S(x^2), then
# d/dx(cos(m * t) * Fx) = 2x F'(x^2) cos(mt)
# with u^m in front, taken to its conclusion
# F = Sa, G = Sb
# d/dx(x^m (cos(m y) F(x^2) + sin(m y) G(x^2))) =
# x^(m - 1) (2 x^2 (F'(x^2) cos(m y) + G'(x^2) sin(m y)) + m F(x^2) cos(m y) + m G(x^2) sin(m y))
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# m x "kernel" above
# d/dy(x^m (cos(m y) F(x^2) + sin(m y) G(x^2))) = m x^m (G(x^2) cos(m y) - F(x^2) sin(m y))
umm1 = u**(m - 1)
twousq = 2 * usq
aterm = cost * (twousq * Sprimea + m * Sa)
bterm = sint * (twousq * Sprimeb + m * Sb)
dr += umm1 * (aterm + bterm)
dt += m * um * (-Sa * sint + Sb * cost)
return z, dr, dt
def off_axis_conic_sigma_der(c, kappa, r, t, dx, dy=0):
"""Derivatives of 1/off_axis_conic_sigma.
See Eq. (5.2) of oe-20-3-2483.
Parameters
----------
c : float
axial curvature of the conic
kappa : float
conic constant
r : numpy.ndarray
radial coordinate, where r=0 is centered on the off-axis section
t : numpy.ndarray
azimuthal coordinate
dx : float
shift of the surface in x with respect to the base conic vertex,
mutually exclusive to dy (only one may be nonzero)
use dx=0 when dy != 0
dy : float
shift of the surface in y with respect to the base conic vertex
Returns
-------
numpy.ndarray, numpy.ndarray
d/dr(z), d/dt(z)
"""
if dy != 0 and dx != 0:
raise ValueError('only one of dx/dy may be nonzero')
cost = np.cos(t)
sint = np.sin(t)
if dx != 0:
s = dx
oblique_term = 2 * s * r * cost
ddr_oblique = 2 * r + 2 * s * cost
# I accept the evil in writing this the way I have
# to deduplicate the computation
ddt_oblique_ = r * (-s) * sint
else:
s = dy
oblique_term = 2 * s * r * sint
ddr_oblique = 2 * r + 2 * s * sint
ddt_oblique_ = r * s * cost
aggregate_term = r * r + oblique_term + s * s
csq = c * c
# d/dr first
phi_kernel = (1 + kappa) * csq * aggregate_term
phi = np.sqrt(1 - phi_kernel)
notquitephi_kernel = kappa * csq * aggregate_term
notquitephi = np.sqrt(1 + notquitephi_kernel)
num = csq * (1 + kappa) * ddr_oblique * notquitephi
den = 2 * (1 - phi_kernel)**(3 / 2)
term1 = num / den
num = csq * kappa * ddr_oblique
den = 2 * phi * notquitephi
term2 = num / den
dr = term1 + term2
# d/dt
num = csq * (1 + kappa) * ddt_oblique_ * notquitephi
den = (1 - phi_kernel)**(3 / 2) # phi^3?
term1 = num / den
num = csq * kappa * ddt_oblique_
den = phi * notquitephi
term2 = num / den
dt = term1 + term2 # minus in writing, but sine/cosine
return dr, dt
def off_axis_conic_der(c, kappa, r, t, dx, dy=0):
"""Radial and azimuthal derivatives of an off-axis conic.
Parameters
----------
c : float
axial curvature of the conic
kappa : float
conic constant
r : numpy.ndarray
radial coordinate, where r=0 is centered on the off-axis section
t : numpy.ndarray
azimuthal coordinate
dx : float
shift of the surface in x with respect to the base conic vertex,
mutually exclusive to dy (only one may be nonzero)
use dx=0 when dy != 0
dy : float
shift of the surface in y with respect to the base conic vertex
Returns
-------
numpy.ndarray, numpy.ndarray
d/dr(z), d/dt(z)
"""
if dy != 0 and dx != 0:
raise ValueError('only one of dx/dy may be nonzero')
cost = np.cos(t)
sint = np.sin(t)
if dx != 0:
s = dx
oblique_term = 2 * s * r * cost
ddr_oblique = 2 * r + 2 * s * cost
# I accept the evil in writing this the way I have
# to deduplicate the computation
ddt_oblique_ = r * (-s) * sint
ddt_oblique = 2 * ddt_oblique_
else:
s = dy
oblique_term = 2 * s * r * sint
ddr_oblique = 2 * r + 2 * s * sint
ddt_oblique_ = r * s * cost
ddt_oblique = 2 * ddt_oblique_
aggregate_term = r * r + oblique_term + s * s
csq = c * c
c3 = csq * c
# d/dr first
num = c * ddr_oblique
phi_kernel = (1 + kappa) * csq * aggregate_term
phi = np.sqrt(1 - phi_kernel)
phip1 = 1 + phi
phip1sq = phip1 * phip1
den = phip1
term1 = num / den
num = c3 * (1 + kappa) * ddr_oblique * aggregate_term
den = (2 * phi) * phip1sq
term2 = num / den
dr = term1 + term2
# d/dt
num = c * ddt_oblique
den = phip1
term1 = num / den
num = c3 * (1 + kappa) * ddt_oblique_ * aggregate_term
den = phi * phip1sq
term2 = num / den
dt = term1 + term2
return dr, dt
def zernike_nm_sequence(nms, r, t, norm=True):
"""Zernike polynomial of radial order n, azimuthal order m at point r, t.
Parameters
----------
nms : iterable of tuple of int,
sequence of (n, m); looks like [(1,1), (3,1), ...]
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
norm : bool, optional
if True, orthonormalize the result (unit RMS)
else leave orthogonal (zero-to-peak = 1)
Returns
-------
generator
yields one mode at a time of nms
"""
# this function deduplicates all possible work. It uses a connection
# to the jacobi polynomials to efficiently compute a series of zernike
# polynomials
# it follows this basic algorithm:
# for each (n, m) compute the appropriate Jacobi polynomial order
# collate the unique values of that for each |m|
# compute a set of jacobi polynomials for each |m|
# compute r^|m| , sin(|m|*t), and cos(|m|*t for each |m|
#
# benchmarked at 12.26 ns/element (256x256), 4.6GHz CPU = 56 clocks per element
# ~36% faster than previous impl (12ms => 8.84 ms)
x = 2 * r**2 - 1
ms = [e[1] for e in nms]
am = truenp.abs(ms)
amu = truenp.unique(am)
def factory():
return 0
jacobi_sequences_mjn = defaultdict(factory)
# jacobi_sequences_mjn is a lookup table from |m| to all orders < max(n_j)
# for each |m|, i.e. 0 .. n_j_max
for nm, am_ in zip(nms, am):
n = nm[0]
nj = (n - am_) // 2
if nj > jacobi_sequences_mjn[am_]:
jacobi_sequences_mjn[am_] = nj
for k in jacobi_sequences_mjn:
nj = jacobi_sequences_mjn[k]
jacobi_sequences_mjn[k] = truenp.arange(nj + 1)
jacobi_sequences = {}
jacobi_sequences_mjn = dict(jacobi_sequences_mjn)
for k in jacobi_sequences_mjn:
n_jac = jacobi_sequences_mjn[k]
jacobi_sequences[k] = list(jacobi_sequence(n_jac, 0, k, x))
powers_of_m = {}
sines = {}
cosines = {}
for m in amu:
powers_of_m[m] = r**m
sines[m] = np.sin(m * t)
cosines[m] = np.cos(m * t)
for n, m in nms:
absm = abs(m)
nj = (n - absm) // 2
jac = jacobi_sequences[absm][nj]
if norm:
jac = jac * zernike_norm(n, m)
if m == 0:
# rotationally symmetric Zernikes are jacobi
yield jac
else:
if m < 0:
azpiece = sines[absm]
else:
azpiece = cosines[absm]
radialpiece = powers_of_m[absm]
out = jac * azpiece * radialpiece # jac already contains the norm
yield out
def zernike_nm_der(n, m, r, t, norm=True):
"""Derivatives of Zernike polynomial of radial order n, azimuthal order m, w.r.t r and t.
Parameters
----------
n : int
radial order
m : int
azimuthal order
r : numpy.ndarray
radial coordinates
t : numpy.ndarray
azimuthal coordinates
norm : bool, optional
if True, orthonormalize the result (unit RMS)
else leave orthogonal (zero-to-peak = 1)
Returns
-------
numpy.ndarray, numpy.ndarray
dZ/dr, dZ/dt
"""
# x = 2 * r ** 2 - 1
# R = radial polynomial R_n^m, not dZ/dr
# R = P_(n-m)//2^(0,|m|) (x)
# = modified jacobi polynomial
# dR = 4r R'(x) (chain rule)
# => use jacobi_der
# if m == 0, dZ = dR
# for m != 0, Z = r^|m| * R * cos(mt)
# the cosine term has no impact on the radial derivative,
# for which we need the product rule:
# d/dr(u v) = v(du/dr) + u(dv/dr)
# u = R, which we already have the derivative of
# v = r^|m| = r^k
# dv/dr = k r^(k-1)
# d/dr(Z) = r^k * (4r * R'(x)) + R * k r^(k-1)
# ------------------ -------------
# v du u dv
#
# all of that is multiplied by d/dr( cost ) or sint, which is just a "pass-through"
# since cost does not depend on r
#
# in azimuth it's the other way around: regular old Zernike computation,
# multiplied by d/dt ( cost )
x = 2 * r**2 - 1
am = abs(m)
n_j = (n - am) // 2
# dv from above == d/dr(R(2r^2-1))
dv = (4 * r) * jacobi_der(n_j, 0, am, x)
if norm:
znorm = zernike_norm(n, m)
if m == 0:
dr = dv
dt = np.zeros_like(dv)
else:
v = jacobi(n_j, 0, am, x)
u = r**am
du = am * r**(am - 1)
dr = v * du + u * dv
if m < 0:
dt = am * np.cos(am * t)
dr *= np.sin(am * t)
else:
dt = -m * np.sin(m * t)
dr *= np.cos(m * t)
# dt = dt * (u * v)
# = cost * r^|m| * R
# faster to write it as two in-place ops here
# (no allocations)
dt *= u
dt *= v
# ugly as this is, we skip one multiply
# by doing these extra ifs
if norm:
dt *= znorm
if norm:
dr *= znorm
return dr, dt
def generate_finite_ray_fan(nrays,
na,
P=0,
min_na=None,
azimuth=90,
yangle=0,
xangle=0,
n=1,
distribution='uniform'):
"""Generate a 1D fan of rays.
Parameters
----------
nrays : int
the number of rays in the fan
na : float
object-space numerical aperture
P : numpy.ndarray
length 3 vector containing the position from which the rays emanate
min_na : float, optional
minimum NA for the beam, -na if None
azimuth: float
angle in the XY plane, degrees. 0=X ray fan, 90=Y ray fan
yangle : float
propagation angle of the chief/gut ray with respect to the Y axis, clockwise
xangle : float
propagation angle of the gut ray with respect to the X axis, clockwise
n : float
refractive index at P (1=vacuum)
distribution : str, {'uniform', 'random', 'cheby'}
The distribution to use when placing the rays
a uniform distribution has rays which are equally spaced from minr to maxr,
random has rays randomly distributed in minr and maxr, while cheby has the
Cheby-Gauss-Lobatto roots as its locations from minr to maxr
Returns
-------
numpy.ndarray, numpy.ndarray
"P" and "S" variables, positions and direction cosines of the rays
"""
# TODO: revisit this; tracing a parabola from the focus, the output
# ray spacing is not uniform as it should be. Or is this some manifestation
# of the sine condition?
# more likely it's the square root since it hides unless the na is big
P = _ensure_P_vec(P)
distribution = distribution.lower()
if min_na is None:
min_na = -na
max_t = np.arcsin(na / n)
min_t = np.arcsin(min_na / n)
if distribution == 'uniform':
t = np.linspace(min_t, max_t, nrays, dtype=config.precision)
elif distribution == 'random':
t = np.random.uniform(low=min_t, high=max_t,
size=nrays).astype(config.precision)
# use the even function for the y direction cosine,
# use trig identity to compute the z direction cosine
l = np.sin(t) # NOQA
m = np.sqrt(1 - l * l) # NOQA
k = np.array([0.], dtype=t.dtype)
k = np.broadcast_to(k, (nrays, ))
if azimuth == 0:
k, l = l, k # NOQA swap Y and X axes
S = np.stack([k, l, m], axis=1)
if yangle != 0 and xangle != 0:
R = make_rotation_matrix((0, yangle, -xangle))
# newaxis for batch matmul, squeeze needed for size 1 dim after
S = np.matmul(R, S[..., np.newaxis]).squeeze()
# need to see a copy of P for each ray, -> add empty dim and broadcast
P = P[np.newaxis, :]
P = np.broadcast_to(P, (nrays, 3))
return P, S