From 6aea5009c2e80e7b56cb1d48aaa2b95b9b0cbc7c Mon Sep 17 00:00:00 2001 From: John Franklin Crenshaw Date: Thu, 7 Dec 2023 17:21:53 -0800 Subject: [PATCH] Added plotting utils for pupil masks and the ImageMapper round trip. --- python/lsst/ts/wep/imageMapper.py | 13 +- python/lsst/ts/wep/utils/plotUtils.py | 235 +++++++++++++++++++++++++- 2 files changed, 238 insertions(+), 10 deletions(-) diff --git a/python/lsst/ts/wep/imageMapper.py b/python/lsst/ts/wep/imageMapper.py index 4dcb8d808..54b6d3569 100644 --- a/python/lsst/ts/wep/imageMapper.py +++ b/python/lsst/ts/wep/imageMapper.py @@ -25,15 +25,10 @@ import numpy as np from lsst.ts.wep.image import Image from lsst.ts.wep.instrument import Instrument -from lsst.ts.wep.utils import ( - DefocalType, - PlaneType, - centerWithTemplate, - configClass, - mergeConfigWithFile, - polygonContains, - zernikeGradEval, -) +from lsst.ts.wep.utils.enumUtils import DefocalType, PlaneType +from lsst.ts.wep.utils.ioUtils import configClass, mergeConfigWithFile +from lsst.ts.wep.utils.miscUtils import centerWithTemplate, polygonContains +from lsst.ts.wep.utils.zernikeUtils import zernikeGradEval from scipy.interpolate import interpn diff --git a/python/lsst/ts/wep/utils/plotUtils.py b/python/lsst/ts/wep/utils/plotUtils.py index 73e3bed99..44d08fb67 100644 --- a/python/lsst/ts/wep/utils/plotUtils.py +++ b/python/lsst/ts/wep/utils/plotUtils.py @@ -19,9 +19,16 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -__all__ = ["plotZernike"] +__all__ = ["plotZernike", "plotPupilMaskElements", "plotRoundTrip"] + +from typing import Optional, Union import matplotlib.pyplot as plt +import numpy as np +from lsst.ts.wep.image import Image +from lsst.ts.wep.imageMapper import ImageMapper +from lsst.ts.wep.instrument import Instrument +from lsst.ts.wep.utils.enumUtils import BandLabel, DefocalType def plotZernike(zkIdx, zk, unit, saveFilePath=None): @@ -50,3 +57,229 @@ def plotZernike(zkIdx, zk, unit, saveFilePath=None): else: plt.savefig(saveFilePath, bbox_inches="tight") plt.close() + + +def plotPupilMaskElements( + instrument: Instrument, + fieldAngle: tuple, + legend: bool = True, + minimal: bool = True, + ax: Optional[plt.Axes] = None, +) -> None: + """Plot the mask elements as circles. + + Outer and inner elements have the same color, with the inner elements dashed. + The pupil is highlighted in yellow. + + Parameters + ---------- + instrument : Instrument + The ts_wep instrument that contains the mask parameters. + fieldAngle : tuple + Tuple of x and y field angles in degrees. + legend : bool, optional + Whether to draw the legend. (the default is True) + minimal : bool, optional + Whether to only draw the elements that presently determine the mask. + (the default is True) + ax : plt.Axes, optional + A matplotlib axis to plot on. If None passed, plt.gca() is used. + (the default is None) + """ + # Get angle radius + rTheta = np.sqrt(np.sum(np.square(fieldAngle))) + + # Generate angles around the circle + theta = np.linspace(0, 2 * np.pi, 10_000) + + # Make lists to save the inner and outer radius at each angle + innerR = [] + outerR = [] + + # Create an empty dictionary for the colors + colors = dict() + + # Make lists for plotting + keys = [] + colorKeys = [] + xs = [] + ys = [] + rs = [] + + # Loop over every mask element + for key, val in instrument.maskParams.items(): + # Determine the color + colorKey = key.removesuffix("Outer").removesuffix("Inner") + if colorKey not in colors: + colors[colorKey] = f"C{len(colors)}" + + # Skip elements for which we are not far enough out + if rTheta < val["thetaMin"]: + continue + + # Calculate the radius and center of the mask in meters + radius = np.polyval(val["radius"], rTheta) + rCenter = np.polyval(val["center"], rTheta) + + # Calculate x and y coordinates of the center + xCenter = 0 if rTheta == 0 else rCenter * fieldAngle[0] / rTheta + yCenter = 0 if rTheta == 0 else rCenter * fieldAngle[1] / rTheta + + # Calculate x and y of circle + # Using the polar equation of a circle so that points are gridded on theta + A = xCenter * np.cos(theta) + yCenter * np.sin(theta) + B = radius**2 - (xCenter**2 + yCenter**2) + A**2 + with np.errstate(invalid="ignore"): + r = np.sqrt(B) + A + x = r * np.cos(theta) + y = r * np.sin(theta) + + # Save the radii to either the inner or outer lists + if "Inner" in key: + innerR.append(r) + else: + outerR.append(r) + + # Save the values for plotting + keys.append(key) + colorKeys.append(colorKey) + xs.append(x) + ys.append(y) + rs.append(r) + + # Fill dummy values if inner or outer radii are missing + innerR = [np.full_like(theta, 0)] if len(innerR) == 0 else innerR + outerR = [np.full_like(theta, np.inf)] if len(outerR) == 0 else outerR + + # Determine the minimum and maximum radii at each angle + innerR = np.nanmax(innerR, axis=0) + outerR = np.nanmin(outerR, axis=0) + + # Select the axis + ax = plt.gca() if ax is None else ax + + # Fill the pupil in yellow + xIn = innerR * np.cos(theta) + yIn = innerR * np.sin(theta) + xOut = outerR * np.cos(theta) + yOut = outerR * np.sin(theta) + ax.fill(xOut, yOut, "gold", alpha=0.2) + ax.fill(xIn, yIn, "w") + + # Loop over elements for plotting + for key, colorKey, x, y, r in zip(keys, colorKeys, xs, ys, rs): + # Determine the color + c = colors[colorKey] + + # Determine the line style + ls = "--" if "Inner" in key else "-" + + if not minimal: + ax.plot(x, y, ls=ls, c=c, label=key) + continue + + # Determine which points to plot + idx = np.where((r >= innerR) & (r <= outerR))[0] + if len(idx) == 0: + continue + + # Split the points into consecutive segments + cutPoints = np.where(np.diff(idx) != 1)[0] + 1 + splits = np.split(idx, cutPoints) + for split in splits: + ax.plot(x[split], y[split], ls=ls, c=c) + + # Add legend for this element + ax.plot([], [], ls=ls, c=c, label=key) + + # Set the limits, axis labels, and title + rmax = np.abs(outerR).max() + rmax = rmax if np.isfinite(rmax) else np.abs(innerR).max() + lim = 1.15 * rmax + lim = lim if np.isfinite(lim) else None + ax.set( + xlim=(-lim, +lim), + ylim=(-lim, +lim), + xlabel="meters", + ylabel="meters", + aspect="equal", + title=f"$\\theta\,=\,${rTheta:.2f}$\!^\circ$", + ) + + # Draw the legend? + if legend: + ax.legend(bbox_to_anchor=(1.04, 1), borderaxespad=0, loc="upper left") + + +def plotRoundTrip( + fieldAngle: tuple, + defocalType: Union[DefocalType, str], + band: Union[BandLabel, str] = BandLabel.REF, + opticalModel: str = "offAxis", + zk: np.ndarray = np.zeros(19), + nPixels: int = 180, +): + """Plot a roundtrip of the ImageMapper + + Parameters + ---------- + fieldAngle : tuple + Tuple of x and y field angles in degrees. + defocalType : Union[DefocalType, str] + The DefocalType (or corresponding string) + band : Union[BandLabel, str], optional + The BandLabel (or corresponding string) + (the default is BandLabel.REF) + opticalModel : str, optional + A string specifying the optical model + (the default is "offAxis") + zk : np.ndarray, optional + Array of Zernike coefficients for wavefront aberrations (in meters) + (the default is an array of zeros) + nPixels : int, optional + The number of pixels on a side for the images. + (the default is 180) + + Returns + ------- + fig + The matplotlib Figure + axes + The matplotlib Axes + """ + # Create the image mapper + mapper = ImageMapper(opticalModel=opticalModel) + + # Forward model an image + image = Image( + np.zeros((nPixels, nPixels)), + fieldAngle, + defocalType, + band, + ) + image = mapper.mapPupilToImage(image, zk) + + # Then map back to the pupil + pupilRecon = mapper.mapImageToPupil(image, zk) + + # Create the pupil mask + pupil = mapper.createPupilMask(image) + + # Plot everything! + fig, axes = plt.subplots(1, 4, figsize=(10, 2), dpi=150) + + settings = {"origin": "lower", "vmin": 0, "vmax": 1} + + axes[0].imshow(pupil, **settings) + axes[0].set(title="Original") + + axes[1].imshow(image.image, **settings) + axes[1].set(title="Mapped to image") + + axes[2].imshow(pupilRecon.image, **settings) + axes[2].set(title="Back to pupil") + + axes[3].imshow(np.abs(pupilRecon.image - pupil), **settings) + axes[3].set(title="Abs Pupil difference") + + return fig, axes