def ray_aim(P, S, prescription, j, wvl, target=(0, 0, np.nan), debug=False):
"""Aim a ray such that it encounters the jth surface at target.
Parameters
----------
P : numpy.ndarray
shape (3,), a single ray's initial positions
S : numpy.ndarray
shape (3,) a single ray's initial direction cosines
prescription : iterable
sequence of surfaces in the prescription
j : int
the surface index in prescription at which the ray should hit (target)
wvl : float
wavelength of light to use in ray aiming, microns
target : iterable of length 3
the position at which the ray should intersect the target surface
NaNs indicate to ignore that position in aiming
debug : bool, optional
if True, returns the (ray-aiming) optimization result as well as the
adjustment P
Returns
-------
numpy.ndarray
deltas to P which result in ray intersection
"""
P = np.asarray(P).astype(config.precision).copy()
S = np.asarray(S).astype(config.precision).copy()
target = np.asarray(target)
trace_path = prescription[:j + 1]
def optfcn(x):
P[:2] = x
phist, _ = spencer_and_murty.raytrace(trace_path, P, S, wvl)
final_position = phist[-1]
euclidean_dist = (final_position - target)**2
euclidean_dist = np.nansum(
euclidean_dist) / 3 # /3 = div by number of axes
return euclidean_dist
res = optimize.minimize(optfcn, np.zeros(2), method='L-BFGS-B')
P[:] = 0
P[:2] = res.x
if debug:
return P, res
else:
return P
def _ensure_P_vec(P):
if not hasattr(P, '__iter__'):
P = np.array([0, 0, P], dtype=config.precision)
else:
# iterable
P2 = np.zeros(3, dtype=config.precision)
P2[-len(P):] = P
P = P2
return np.asarray(P).astype(config.precision)
def fix_zero_singularity(arr, x, y, fill='xypoly', order=2):
"""Fix a singularity at the origin of arr by polynomial interpolation.
Parameters
----------
arr : numpy.ndarray
array of dimension 2 to modify at the origin (x==y==0)
x : numpy.ndarray
array of dimension 2 of X coordinates
y : numpy.ndarray
array of dimension 2 of Y coordinates
fill : str, optional, {'xypoly'}
how to fill. Not used/hard-coded to X/Y polynomials, but made an arg
today in case it may be added future for backwards compatibility
order : int
polynomial order to fit
Returns
-------
numpy.ndarray
arr (modified in-place)
"""
zloc = find_zero_indices_2d(x, y)
min_y = zloc[0] - order
max_y = zloc[0] + order + 1
min_x = zloc[1] - order
max_x = zloc[1] + order + 1
# newaxis schenanigans to get broadcasting right without
# meshgrid
ypts = np.arange(min_y, max_y)[:, np.newaxis]
xpts = np.arange(min_x, max_x)[np.newaxis, :]
window = arr[ypts, xpts].copy()
c = [s // 2 for s in window.shape]
window[c] = np.nan
# no longer need xpts, ypts
# really don't care about fp64 vs fp32 (very small arrays)
xpts = xpts.astype(float)
ypts = ypts.astype(float)
# use Hermite polynomials as
# XY polynomial-like basis orthogonal
# over the infinite plane
# H0 = 0
# H1 = x
# H2 = x^2 - 1, and so on
ns = np.arange(order + 1)
xbasis = hermite_He_sequence(ns, xpts)
ybasis = hermite_He_sequence(ns, ypts)
xbasis = [mode_1d_to_2d(mode, xpts, ypts, 'x') for mode in xbasis]
ybasis = [mode_1d_to_2d(mode, xpts, ypts, 'y') for mode in ybasis]
basis_set = np.asarray([*xbasis, *ybasis])
coefs = lstsq(basis_set, window)
projected = np.dot(basis_set[:, c[0], c[1]], coefs)
arr[zloc] = projected
return arr
def sum_of_2d_modes(modes, weights):
"""Compute a sum of 2D modes.
Parameters
----------
modes : iterable
sequence of ndarray of shape (k, m, n);
a list of length k with elements of shape (m,n) works
weights : numpy.ndarray
weight of each mode
Returns
-------
numpy.ndarry
ndarray of shape (m, n) that is the sum of modes as given
"""
modes = np.asarray(modes)
weights = np.asarray(weights).astype(modes.dtype)
# dot product of the 0th dim of modes and weights => weighted sum
return np.tensordot(modes, weights, axes=(0, 0))
def lstsq(modes, data):
"""Least-Squares fit of modes to data.
Parameters
----------
modes : iterable
modes to fit; sequence of ndarray of shape (m, n)
data : numpy.ndarray
data to fit, of shape (m, n)
place NaN values in data for points to ignore
Returns
-------
numpy.ndarray
fit coefficients
"""
mask = np.isfinite(data)
data = data[mask]
modes = np.asarray(modes)
modes = modes.reshape((modes.shape[0], -1)) # flatten second dim
modes = modes[:, mask.ravel()].T # transpose moves modes to columns, as needed for least squares fit
c, *_ = np.linalg.lstsq(modes, data, rcond=None)
return c
def barplot(coefs,
names=None,
orientation='h',
buffer=1,
zorder=3,
number=True,
offset=0,
width=0.8,
fig=None,
ax=None):
"""Create a barplot of coefficients and their names.
Parameters
----------
coefs : dict
with keys of Zn, values of numbers
names : dict
with keys of Zn, values of names (e.g. Primary Coma X)
orientation : str, {'h', 'v', 'horizontal', 'vertical'}
orientation of the plot
buffer : float, optional
buffer to use around the left and right (or top and bottom) bars
zorder : int, optional
zorder of the bars. Use zorder > 3 to put bars in front of gridlines
number : bool, optional
if True, plot numbers along the y=0 line showing indices
offset : float, optional
offset to apply to bars, useful for before/after Zernike breakdowns
width : float, optional
width of bars, useful for before/after Zernike breakdowns
fig : matplotlib.figurnp.Figure
Figure containing the plot
ax : matplotlib.axes.Axis
Axis containing the plot
Returns
-------
fig : matplotlib.figurnp.Figure
Figure containing the plot
ax : matplotlib.axes.Axis
Axis containing the plot
"""
from matplotlib import pyplot as plt
fig, ax = share_fig_ax(fig, ax)
coefs2 = np.asarray(list(coefs.values()))
idxs = np.asarray(list(coefs.keys()))
coefs = coefs2
lims = (idxs[0] - buffer, idxs[-1] + buffer)
if orientation.lower() in ('h', 'horizontal'):
vmin, vmax = coefs.min(), coefs.max()
drange = vmax - vmin
offsetY = drange * 0.01
ax.bar(idxs + offset, coefs, zorder=zorder, width=width)
plt.xticks(idxs, names, rotation=90)
if number:
for i in idxs:
ax.text(i, offsetY, str(i), ha='center')
else:
ax.barh(idxs + offset, coefs, zorder=zorder, height=width)
plt.yticks(idxs, names)
if number:
for i in idxs:
ax.text(0, i, str(i), ha='center')
ax.set(xlim=lims)
return fig, ax
def raytrace(surfaces, P, S, wvl, n_ambient=1):
"""Perform a raytrace through a sequence of surfaces.
Notes
-----
When P and S are single dimensional, a single ray is traced.
When they have two dimensions, the first dimension is the "batch" and the
second contains [X,Y,Z] and [k,l,m] for each ray in the batch.
There is no internal ray aiming or other adjustment to P and S.
In a batch raytrace, there is no reason all rows of P and S must belong to
the same ray bundle.
wvl does not matter and is not used in raytraces with only reflective
surfaces
A ray originating "at infinity" would have
P = [Px, Py, -1e99]
S = [0, 0, 1] # propagating in the +z direction
though the value of P is not so important, since S defines the ray as moving in the +z direction only
Parameters
----------
surfaces : iterable
the surfaces to trace through;
a surface is defined by the interface:
surf.F(x,y) -> z sag
surf.Fp(x,y) -> (Fx, Fy, 1) derivatives (with S&M convention for 1 in z)
surf.typ in {STYPE}
surf.P, surface global coordinates, [X,Y,Z]
surf.R, surface rotation matrix (may be None)
surf.n(wvl) -> refractive index (wvl in um)
P : numpy.ndarray
shape (3,) or (N,3), any float dtype
position (X0,Y0,Z0) at the outset of the raytrace
S : numpy.ndarray
shape (3,) or (N,3), any float dtype
(k,l,m) starting direction cosines
wvl : float
wavelength of light, um
n_ambient : float
ambient index of refraction (1=vacuum)
Returns
-------
P_hist, S_hist
position history and direction cosine history
Implementation Notes
--------------------
See Spencer & Murty, General Ray-Tracing Procedure JOSA 1961
Steps (I, II, III, IV) utilize the functions:
I -> transform_to_local_coords
II -> newton_raphson_solve_s
III -> reflect or refract
IV -> transform_to_global_coords
"""
P = np.asarray(P)
S = np.asarray(S)
jj = len(surfaces)
P_hist = np.empty((jj+1, *P.shape), dtype=P.dtype)
S_hist = np.empty((jj+1, *S.shape), dtype=P.dtype)
Pj = P
Sj = S
P_hist[0] = P
S_hist[0] = S
nj = n_ambient
for j, surf in enumerate(surfaces):
# I - transform from global to local coordinates
P0, Sj = transform_to_local_coords(Pj, surf.P, Sj, surf.R)
# II - find ray intersection
Pj, r = intersect(P0, Sj, surf.sag_normal)
# III - reflection or refraction
if surf.typ == STYPE_REFLECT:
Sjp1 = reflect(Sj, r)
elif surf.typ == STYPE_REFRACT:
nprime = surf.n(wvl)
Sjp1 = refract(nj, nprime, Sj, r)
nj = nprime
else:
# other surface types do not bend rays
Sjp1 = Sj
Pjp1 = Pj
# IV - back to world coordinates
if surf.R is None:
Rt = None
else:
# transformation matrix has inverse which is its transpose
Rt = surf.R.T
Pjp1, Sjp1 = transform_to_global_coords(Pj, surf.P, Sjp1, Rt)
P_hist[j+1] = Pjp1
S_hist[j+1] = Sjp1
Pj, Sj = Pjp1, Sjp1
return P_hist, S_hist
def generate_collimated_ray_fan(nrays,
maxr,
z=0,
minr=None,
azimuth=90,
yangle=0,
xangle=0,
distribution='uniform',
aim_at=None):
"""Generate a 1D fan of rays.
Colloquially, an extended field in Y for an object at inf is represented by a ray fan with yangle != 0.
Parameters
----------
nrays : int
the number of rays in the fan
maxr : float
maximum radial value of the fan
z : float
z position for the ray fan
minr : float, optional
minimum radial value of the fan, -maxr if None
azimuth: float
angle in the XY plane, degrees. 0=X ray fan, 90=Y ray fan
yangle : float
propagation angle of the rays with respect to the Y axis, clockwise
xangle : float
propagation angle of the rays with respect to the X axis, clockwise
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
aim_at : numpy.ndarray or float
position [X,Y,Z] aim the rays such that the gut ray of the fan,
when propagated in an open medium, hits (aim_at)
if a float, interpreted as if x=0,y=0, z=aim_at
This argument mimics ray aiming in commercial optical design software, and
is used in conjunction with xangle/yangle to form off-axis ray bundles which
properly go through the center of the stop
Returns
-------
numpy.ndarray, numpy.ndarray
"P" and "S" variables, positions and direction cosines of the rays
"""
dtype = config.precision
distribution = distribution.lower()
if minr is None:
minr = -maxr
S = np.array([0, 0, 1], dtype=dtype)
R = make_rotation_matrix((0, yangle, -xangle))
S = np.matmul(R, S)
# need to see a copy of S for each ray, -> add empty dim and broadcast
S = S[np.newaxis, :]
S = np.broadcast_to(S, (nrays, 3))
# now generate the radial part of P
if distribution == 'uniform':
r = np.linspace(minr, maxr, nrays, dtype=dtype)
elif distribution == 'random':
r = np.random.uniform(low=minr, high=maxr, size=nrays).astype(dtype)
t = np.asarray(np.radians(azimuth), dtype=dtype)
t = np.broadcast_to(t, r.shape)
x, y = polar_to_cart(r, t)
z = np.array(z, dtype=x.dtype)
z = np.broadcast_to(z, x.shape)
xyz = np.stack([x, y, z], axis=1)
return xyz, S