def plot_rays(phist, lw=1, ls='-', c='r', alpha=1, zorder=4, x='z', y='y', fig=None, ax=None):
"""Plot rays in 2D.
Parameters
----------
phist : list or numpy.ndarray
the first return from spencer_and_murty.raytrace,
iterable of arrays of length 3 (X,Y,Z)
lw : float, optional
linewidth
ls : str, optional
line style
c : color
anything matplotlib interprets as a color, strings, 3-tuples, 4-tuples, ...
alpha : float
opacity of the rays, 1=fully opaque, 0=fully transparent
zorder : int
stack order in the plot, higher z orders are on top of lower z orders
x : str, {'x', 'y', 'z'}
which position to plot on the X axis, defaults to traditional ZY plot
y : str, {'x', 'y', 'z'}
which position to plot on the X axis, defaults to traditional ZY plot
fig : matplotlib.figure.Figure
A figure object
ax : matplotlib.axes.Axis
An axis object
Returns
-------
matplotlib.figure.Figure
A figure object
matplotlib.axes.Axis
An axis object
"""
fig, ax = share_fig_ax(fig, ax)
ph = np.asarray(phist)
xs = ph[..., 0]
ys = ph[..., 1]
zs = ph[..., 2]
sieve = {
'x': xs,
'y': ys,
'z': zs,
}
x = x.lower()
y = y.lower()
x = sieve[x]
y = sieve[y]
ax.plot(x, y, c=c, lw=lw, ls=ls, alpha=alpha, zorder=zorder)
return fig, ax
def plot_transverse_ray_aberration(phist, lw=1, ls='-', c='r', alpha=1, zorder=4, axis='y', fig=None, ax=None):
"""Plot the transverse ray aberration for a single ray fan.
Parameters
----------
phist : list or numpy.ndarray
the first return from spencer_and_murty.raytrace,
iterable of arrays of length 3 (X,Y,Z)
lw : float, optional
linewidth
ls : str, optional
line style
c : color
anything matplotlib interprets as a color, strings, 3-tuples, 4-tuples, ...
alpha : float
opacity of the rays, 1=fully opaque, 0=fully transparent
zorder : int
stack order in the plot, higher z orders are on top of lower z orders
axis : str, {'x', 'y'}
which ray position to plot, x or y
fig : matplotlib.figure.Figure
A figure object
ax : matplotlib.axes.Axis
An axis object
Returns
-------
matplotlib.figure.Figure
A figure object
matplotlib.axes.Axis
An axis object
"""
fig, ax = share_fig_ax(fig, ax)
ph = np.asarray(phist)
sieve = {
'x': 0,
'y': 1,
}
axis = axis.lower()
axis = sieve[axis]
input_rays = ph[0, ..., axis]
output_rays = ph[-1, ..., axis]
ax.plot(input_rays, output_rays, c=c, lw=lw, ls=ls, alpha=alpha, zorder=zorder)
return fig, ax
def plot_spot_diagram(phist, marker='+', c='k', alpha=1, zorder=4, s=None, fig=None, ax=None):
"""Plot a spot diagram from a ray trace.
Parameters
----------
phist : list or numpy.ndarray
the first return from spencer_and_murty.raytrace,
iterable of arrays of length 3 (X,Y,Z)
marker : str, optional
marker style
c : color
anything matplotlib interprets as a color, strings, 3-tuples, 4-tuples, ...
alpha : float
opacity of the rays, 1=fully opaque, 0=fully transparent
zorder : int
stack order in the plot, higher z orders are on top of lower z orders
s : float
marker size or variable used for marker size
axis : str, {'x', 'y'}
which ray position to plot, x or y
fig : matplotlib.figure.Figure
A figure object
ax : matplotlib.axes.Axis
An axis object
Returns
-------
matplotlib.figure.Figure
A figure object
matplotlib.axes.Axis
An axis object
"""
fig, ax = share_fig_ax(fig, ax)
x = phist[-1, ..., 0]
y = phist[-1, ..., 1]
ax.scatter(x, y, c=c, s=s, marker=marker, alpha=alpha, zorder=zorder)
return fig, ax
def test_share_fig_ax_produces_an_axis():
fig = plt.figure()
fig, ax = plotting.share_fig_ax(fig)
assert ax is not None
def test_share_fig_ax_produces_figure_and_axis():
fig, ax = plotting.share_fig_ax()
assert fig
assert ax
def test_share_fig_ax_figure_number_remains_unchanged():
fig, ax = plt.subplots()
fig2, ax2 = plotting.share_fig_ax(fig, ax)
assert fig.number == fig2.number
def barplot_magnitudes(magnitudes,
orientation='h',
sort=False,
buffer=1,
zorder=3,
offset=0,
width=0.8,
fig=None,
ax=None):
"""Create a barplot of magnitudes of coefficient pairs and their names.
e.g., astigmatism will get one bar.
Parameters
----------
magnitudes : dict
keys of names, values of magnitudes. E.g., {'Primary Coma': 1234567}
orientation : str, {'h', 'v', 'horizontal', 'vertical'}
orientation of the plot
sort : bool, optional
whether to sort the zernikes in descending order
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
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.figure.Figure
Figure containing the plot
ax : matplotlib.axes.Axis
Axis containing the plot
Returns
-------
fig : matplotlib.figure.Figure
Figure containing the plot
ax : matplotlib.axes.Axis
Axis containing the plot
"""
from matplotlib import pyplot as plt
mags = magnitudes.values()
names = magnitudes.keys()
idxs = np.arange(len(names))
# idxs = np.asarray(list(range(len(names))))
if sort:
mags, names = sort_xy(mags, names)
mags = list(reversed(mags))
names = list(reversed(names))
lims = (idxs[0] - buffer, idxs[-1] + buffer)
fig, ax = share_fig_ax(fig, ax)
if orientation.lower() in ('h', 'horizontal'):
ax.bar(idxs + offset, mags, zorder=zorder, width=width)
plt.xticks(idxs, names, rotation=90)
ax.set(xlim=lims)
else:
ax.barh(idxs + offset, mags, zorder=zorder, height=width)
plt.yticks(idxs, names)
ax.set(ylim=lims)
return fig, ax
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 plot_optics(prescription, phist, mirror_backing=None, points=100,
lw=1, ls='-', c='k', alpha=1, zorder=3,
x='z', y='y', fig=None, ax=None):
"""Draw the optics of a prescription.
Parameters
----------
prescription : iterable of Surface
a prescription for an optical layout
phist : iterable of numpy.ndarray
the first return of spencer_and_murty.raytrace, the history of positions
through a raytrace
mirror_backing : TODO
TODO
points : int, optional
the number of points used in making the curve for the surface
lw : float, optional
linewidth
ls : str, optional
line style
c : color, optional
anything matplotlib interprets as a color, strings, 3-tuples, 4-tuples, ...
alpha : float, optional
opacity of the rays, 1=fully opaque, 0=fully transparent
zorder : int
stack order in the plot, higher z orders are on top of lower z orders
x : str, {'x', 'y', 'z'}
which position to plot on the X axis, defaults to traditional ZY plot
y : str, {'x', 'y', 'z'}
which position to plot on the X axis, defaults to traditional ZY plot
fig : matplotlib.figure.Figure
A figure object
ax : matplotlib.axes.Axis
An axis object
Returns
-------
matplotlib.figure.Figure
A figure object
matplotlib.axes.Axis
An axis object
"""
x = x.lower()
y = y.lower()
fig, ax = share_fig_ax(fig, ax)
# manual iteration due to how lenses are drawn, start from -1 so the
# increment can be at the top of a large loop
j = -1
jj = len(prescription)
while True:
j += 1
if j == jj:
break
surf = prescription[j]
if surf.typ == STYPE_REFLECT:
z = surf.P[2]
xpt, ypt, mask, ploty = _gather_inputs_for_surface_sag(surf, phist, j, points, y)
sag = surf.F(xpt, ypt)
sag += z
sag[mask] = np.nan
# TODO: mirror backing
ax.plot(sag, ploty, c=c, lw=lw, ls=ls, alpha=alpha, zorder=zorder)
elif surf.typ == STYPE_REFRACT:
if (j + 1) == jj:
raise ValueError('cant draw a prescription that terminates on a refracting surface')
z = surf.P[2]
xpt, ypt, mask, ploty = _gather_inputs_for_surface_sag(surf, phist, j, points, y)
sag = surf.F(xpt, ypt)
sag += z
sag[mask] = np.nan
# now get the points for the second surface of the lens
j += 1
surf = prescription[j]
z = surf.P[2]
xpt2, ypt2, mask2, ploty2 = _gather_inputs_for_surface_sag(surf, phist, j, points, y)
sag2 = surf.F(xpt2, ypt2)
sag2 += z
sag2[mask2] = np.nan
# now bundle the two surfaces together so one line is drawn for the
# whole lens
first_x = sag[0]
first_y = ploty[0]
# the ::-1 are because we need to reverse the order of the second
# surface's points, so that matplotlib doesn't draw an X through the lens
xx = [*sag, *sag2[::-1], first_x]
yy = [*ploty, *ploty2[::-1], first_y]
ax.plot(xx, yy, c=c, lw=lw, ls=ls, alpha=alpha, zorder=zorder)
return fig, ax
