def sum_of_xy_modes(modesx, modesy, x, y, weightsx=None, weightsy=None):
"""Weighted sum of separable x and y modes projected over the 2D aperture.
Parameters
----------
modesx : iterable
sequence of x modes
modesy : iterable
sequence of y modes
x : numpy.ndarray
x points
y : numpy.ndarray
y points
weightsx : iterable, optional
weights to apply to modesx. If None, [1]*len(modesx)
weightsy : iterable, optional
weights to apply to modesy. If None, [1]*len(modesy)
Returns
-------
numpy.ndarray
modes summed over the 2D aperture
"""
x, y = optimize_xy_separable(x, y)
if weightsx is None:
weightsx = [1]*len(modesx)
if weightsy is None:
weightsy = [1]*len(modesy)
# apply the weights to the modes
modesx = [m*w for m, w in zip(modesx, weightsx)]
modesy = [m*w for m, w in zip(modesy, weightsy)]
# sum the separable bases in 1D
sum_x = np.zeros_like(x)
sum_y = np.zeros_like(y)
for m in modesx:
sum_x += m
for m in modesy:
sum_y += m
# broadcast to 2D and return
shape = (y.size, x.size)
sum_x = np.broadcast_to(sum_x, shape)
sum_y = np.broadcast_to(sum_y, shape)
return sum_x + sum_y
def jacobi_der(n, alpha, beta, x):
"""First derivative of Pn with respect to x, at points x.
Parameters
----------
n : int
polynomial order
alpha : float
first weight parameter
beta : float
second weight parameter
x : numpy.ndarray
x coordinates to evaluate at
Returns
-------
numpy.ndarray
jacobi polynomial evaluated at the given points
"""
# see https://dlmf.nist.gov/18.9
# dPn = (1/2) (n + a + b + 1)P_{n-1}^{a+1,b+1}
# first two terms are specialized for speed
if n == 0:
return np.zeros_like(x)
if n == 1:
return np.ones_like(x) * (0.5 * (n + alpha + beta + 1))
Pn = jacobi(n - 1, alpha + 1, beta + 1, x)
coef = 0.5 * (n + alpha + beta + 1)
return coef * Pn
def hermite_H_der_sequence(ns, x):
"""First derivative of He_[ns] with respect to x, at points x.
Parameters
----------
ns : iterable of int
rising polynomial orders, assumed to be sorted
x : numpy.ndarray
point(s) to evaluate at. Scalars and arrays both work.
Returns
-------
generator of numpy.ndarray
equivalent to array of shape (len(ns), len(x))
"""
# this function includes all the optimizations in the hermite_He func,
# but excludes the note comments. Read that first if you're looking for
# clarity
# see also: prysm.polynomials.jacobi.jacobi_sequence for the meta machinery
# in use here
ns = list(ns)
min_i = 0
if ns[min_i] == 0:
yield np.zeros_like(x)
min_i += 1
if min_i == len(ns):
return
if ns[min_i] == 1:
yield 2 * np.ones_like(x)
min_i += 1
if min_i == len(ns):
return
x2 = 2 * x
P1 = x2
P2 = 4 * (x * x) - 2
if ns[min_i] == 2:
yield 4 * P1
min_i += 1
if min_i == len(ns):
return
Pnm2, Pnm1 = P1, P2
max_n = ns[-1]
for nn in range(3, max_n + 1):
Pn = x2 * Pnm1 - (2 * (nn - 1)) * Pnm2
if ns[min_i] == nn:
yield 2 * nn * Pnm1
min_i += 1
Pnm2, Pnm1 = Pnm1, Pn
def hermite_H_der(n, x):
"""First derivative of H_n with respect to x, at points x.
Parameters
----------
n : int
polynomial order
x : numpy.ndarray
point(s) to evaluate at. Scalars and arrays both work.
Returns
-------
numpy.ndarray
d/dx[H_n(x)]
"""
if n == 0:
return np.zeros_like(x)
return 2 * n * hermite_H(n - 1, x)
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 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 __init__(self,
ifn,
Nact=50,
sep=10,
shift=(0, 0),
rot=(0, 0, 0),
upsample=1,
spline_order=3,
mask=None):
"""Create a new DM model.
This model is based on convolution of a 'poke lattice' with the influence
function. It has the following idiosyncracies:
1. The poke lattice is always "FFT centered" on the array, i.e.
centered on the sample which would contain the DC frequency bin
after an FFT.
2. The rotation is applied in the same sampling as ifn
3. Shifts and resizing are applied using a Fourier method and not
subject to quantization
Parameters
----------
ifn : numpy.ndarray
influence function; assumes the same for all actuators and must
be the same shape as (x,y). Assumed centered on N//2th sample of x, y.
Assumed to be well-conditioned for use in convolution, i.e.
compact compared to the array holding it
Nact : int or tuple of int, length 2
(X, Y) actuator counts
sep : int or tuple of int, length 2
(X, Y) actuator separation, samples of influence function
shift : tuple of float, length 2
(X, Y) shift of the actuator grid to (x, y), units of x influence
function sampling. E.g., influence function on 0.1 mm grid, shift=1
= 0.1 mm shift. Positive numbers describe (rightward, downward)
shifts in image coordinates (origin lower left).
rot : tuple of int, length <= 3
(Z, Y, X) rotations; see coordinates.make_rotation_matrix
upsample : float
upsampling factor used in determining output resolution, if it is different
to the resolution of ifn.
mask : numpy.ndarray
boolean ndarray of shape Nact used to suppress/delete/exclude
actuators; 1=keep, 0=suppress
"""
if isinstance(Nact, int):
Nact = (Nact, Nact)
if isinstance(sep, int):
sep = (sep, sep)
x, y = make_xy_grid(ifn.shape, dx=1)
# stash inputs and some computed values on self
self.ifn = ifn
self.Ifn = fft.fft2(ifn)
self.Nact = Nact
self.sep = sep
self.shift = shift
self.obliquity = truenp.cos(truenp.radians(truenp.linalg.norm(rot)))
self.rot = rot
self.upsample = upsample
# prepare the poke array and supplimentary integer arrays needed to
# copy it into the working array
out = prepare_actuator_lattice(ifn.shape,
Nact,
sep,
mask,
dtype=x.dtype)
self.mask = out['mask']
self.actuators = out['actuators']
self.actuators_work = np.zeros_like(self.actuators)
self.poke_arr = out['poke_arr']
self.ixx = out['ixx']
self.iyy = out['iyy']
# rotation data
self.rotmat = make_rotation_matrix(rot)
XY = apply_rotation_matrix(self.rotmat, x, y)
XY2 = xyXY_to_pixels((x, y), XY)
self.XY = XY
self.XY2 = XY2
self.needs_rot = True
if np.allclose(rot, [0, 0, 0]):
self.needs_rot = False
# shift data
if shift[0] != 0 or shift[1] != 0:
# caps = Fourier variable (x -> X, y -> Y)
# make 2pi/px phase ramps in 1D (much faster)
# then broadcast them to 2D when they're used as transfer functions
# in a Fourier convolution
Y, X = [forward_ft_unit(1, s, shift=False) for s in x.shape]
Xramp = np.exp(X * (-2j * np.pi * shift[0]))
Yramp = np.exp(Y * (-2j * np.pi * shift[1]))
shpx = x.shape
shpy = tuple(reversed(x.shape))
Xramp = np.broadcast_to(Xramp, shpx)
Yramp = np.broadcast_to(Yramp, shpy).T
self.Xramp = Xramp
self.Yramp = Yramp
self.tf = [self.Ifn * self.Xramp * self.Yramp]
else:
self.tf = [self.Ifn]
def jacobi_der_sequence(ns, alpha, beta, x):
"""First partial derivative of Pn w.r.t. x for order ns, i.e. P_n'.
Parameters
----------
ns : iterable
sorted orders to return, e.g. [1, 2, 3, 10] returns P1', P2', P3', P10'
alpha : float
first weight parameter
beta : float
second weight parameter
x : numpy.ndarray
x coordinates to evaluate at
Returns
-------
generator
equivalent to array of shape (len(ns), len(x))
"""
# the body of this function is very similar to that of jacobi_sequence,
# except note that der is related to jacobi n-1,
# and the actual jacobi polynomial has a different alpha and beta
# special note: P0 is invariant of alpha, beta
# and within this function alphap1 and betap1 are "a+1" and "b+1"
alphap1 = alpha + 1
betap1 = beta + 1
# except when it comes time to yield terms, we yield the modification
# per A&S / the NIST link
# and we modify the arguments to
ns = list(ns)
min_i = 0
if ns[min_i] == 0:
# n=0 is piston, der==0
yield np.zeros_like(x)
min_i += 1
if min_i == len(ns):
return
if ns[min_i] == 1:
yield np.ones_like(x) * (0.5 * (1 + alpha + beta + 1))
min_i += 1
if min_i == len(ns):
return
# min_n is at least two, which means min n-1 is 1
# from here below, Pn is P of order i to keep the reader sane, but Pnm1
# is all that is needed;
# therefor, Pn is computed only after testing if we are done and can return
# to avoid a waste computation at the end of the loop
# note that we can hardcode / unroll the loop up to n=3, one further than
# in jacobi, because we use Pnm1
P1 = alphap1 + 1 + (alphap1 + betap1 + 2) * ((x - 1) / 2)
if ns[min_i] == 2:
yield P1 * (0.5 * (2 + alpha + beta + 1))
min_i += 1
if min_i == len(ns):
return
A, B, C = recurrence_abc(1, alphap1, betap1)
P2 = (A * x + B) * P1 - C # no C * Pnm2 =because Pnm2 = 1
if ns[min_i] == 3:
yield P2 * (0.5 * (3 + alpha + beta + 1))
min_i += 1
if min_i == len(ns):
return
# weird look just above P2, need to prepare for lower loop
# by setting Pnm2 = P1, Pnm1 = P2
Pnm2 = P1
Pnm1 = P1
Pn = P2
# A, B, C = recurrence_abc(2, alpha, beta)
# P3 = (A * x + B) * P2 - C * P1
# Pn = P3
max_n = ns[-1]
for i in range(3, max_n + 1):
Pnm2, Pnm1 = Pnm1, Pn
if ns[min_i] == i:
coef = 0.5 * (i + alpha + beta + 1)
yield Pnm1 * coef
min_i += 1
if min_i == len(ns):
return
A, B, C = recurrence_abc(i - 1, alphap1, betap1)
Pn = (A * x + B) * Pnm1 - C * Pnm2
