From 2c9fd8fb02d192040b3b6e22b97862d546c6c210 Mon Sep 17 00:00:00 2001 From: LR-inaf Date: Fri, 15 Nov 2024 18:38:23 +0000 Subject: [PATCH] Feature branch that add PSF-realted plots. --- python/lsst/donut/viz/plot_aos_task.py | 138 +++++++++++++++++++++++++ python/lsst/donut/viz/psf_from_zern.py | 71 +++++++++++++ 2 files changed, 209 insertions(+) create mode 100644 python/lsst/donut/viz/psf_from_zern.py diff --git a/python/lsst/donut/viz/plot_aos_task.py b/python/lsst/donut/viz/plot_aos_task.py index 5d1853c..142b512 100644 --- a/python/lsst/donut/viz/plot_aos_task.py +++ b/python/lsst/donut/viz/plot_aos_task.py @@ -7,8 +7,10 @@ import lsst.pipe.base.connectionTypes as ct import matplotlib.pyplot as plt import numpy as np +from astropy import units as u import yaml from lsst.utils.timer import timeMethod +from lsst.ts.wep.utils import convertZernikesToPsfWidth from .utilities import ( add_rotated_axis, @@ -17,6 +19,7 @@ rose, ) from .zernike_pyramid import zernikePyramid +from .psf_from_zern import psfPanel try: from lsst.rubintv.production.uploaders import MultiUploader @@ -30,6 +33,9 @@ "PlotDonutTaskConnections", "PlotDonutTaskConfig", "PlotDonutTask", + "PlotPsfZernTaskConnections", + "PlotPsfZernTaskConfig", + "PlotPsfZernTask", ] @@ -380,3 +386,135 @@ def runQuantum( seqNum=seq_num, filename=donut_gallery_fn, ) + +class PlotPsfZernTaskConnections( + pipeBase.PipelineTaskConnections, + dimensions=("visit", "instrument"), +): + zernikes = ct.Input( + doc="Zernikes catalog", + dimensions=("visit", "instrument", "detector"), + storageClass="AstropyTable", + multiple=True, + name="zernikes", + ) + psfFromZernPanel = ct.Output( + doc="PSF value retrieved from zernikes", + dimensions=("visit", "instrument"), + storageClass="Plot", + name="psfFromZernPanel", + ) + + +class PlotPsfZernTaskConfig( + pipeBase.PipelineTaskConfig, + pipelineConnections=PlotPsfZernTaskConnections, +): + doRubinTVUpload = pexConfig.Field( + dtype=bool, + doc="Upload to RubinTV", + default=False, + ) + + +class PlotPsfZernTask(pipeBase.PipelineTask): + ConfigClass = PlotPsfZernTaskConfig + _DefaultName = "plotPsfZernTask" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if self.config.doRubinTVUpload: + if not MultiUploader: + raise RuntimeError("MultiUploader is not available") + self.uploader = MultiUploader() + + @timeMethod + def runQuantum( + self, + butlerQC: pipeBase.QuantumContext, + inputRefs: pipeBase.InputQuantizedConnection, + outputRefs: pipeBase.OutputQuantizedConnection, + ) -> None: + zernikes = butlerQC.get(inputRefs.zernikes) + + zkPanel = self.plotPsfFromZern(zernikes) + zkPanel.suptitle( + f"PSF from Zernikes\nvisit: {inputRefs.zernikes[-1].dataId['visit']}", + fontsize="xx-large", + fontweight="book" + ) + + butlerQC.put(zkPanel, outputRefs.psfFromZernPanel) + + if self.config.doRubinTVUpload: + instrument = inputRefs.zernikes.dataId["instrument"] + visit = inputRefs.zernikes.dataId["visit"] + day_obs, seq_num = get_day_obs_seq_num_from_visitid(visit) + with tempfile.TemporaryDirectory() as tmpdir: + psf_zk_panel = Path(tmpdir) / "psf_zk_panel.png" + zkPanel.savefig(psf_zk_panel) + + self.uploader.uploadPerSeqNumPlot( + instrument=get_instrument_channel_name(instrument), + plotName="psf_zk_panel", + dayObs=day_obs, + seqNum=seq_num, + filename=zk_meas_fn, + ) + + def get_psf_degr(self, zset): + return np.sqrt(np.sum(convertZernikesToPsfWidth(zset)**2)) + + def get_rtp_q(self, qtable): + q = qtable.meta["extra"]["boresight_par_angle_rad"] + rot = qtable.meta["extra"]["boresight_rot_angle_rad"] + rtp = q - rot - np.pi / 2 + return rtp, q + + def get_rose_vecs(self, rtp, q): + vecs_xy = { + r"$x_\mathrm{Opt}$": (1, 0), + r"$y_\mathrm{Opt}$": (0, -1), + r"$x_\mathrm{Cam}$": (np.cos(rtp), -np.sin(rtp)), + r"$y_\mathrm{Cam}$": (-np.sin(rtp), -np.cos(rtp)), + } + + vecs_NE = { + "az": (1, 0), + "alt": (0, +1), + "N": (np.sin(q), np.cos(q)), + "E": (np.sin(q - np.pi / 2), np.cos(q - np.pi / 2)), + } + + return vecs_xy, vecs_NE + + def plotPsfFromZern(self, zernikes): + xs = [] + ys = [] + zs = [] + dname = [] + for i, qt in enumerate(zernikes): + dname.append(qt.meta["extra"]["det_name"]) + xs.append(qt["extra_centroid"]["x"][1:].value) + ys.append(qt["extra_centroid"]["y"][1:].value) + z = [] + for row in qt[[col for col in qt.colnames if "Z" in col]][1:].iterrows(): + z.append([el.to(u.micron).value for el in row]) + zs.append(np.array(z)) + + xs = np.array(xs) + ys = np.array(ys) + zs = np.array(zs) + psf = np.array([[self.get_psf_degr(pair) for pair in det] for det in zs]) + + fig = psfPanel(xs, ys, psf, dname) + + # draw rose + rtp, q = self.get_rtp_q(zernikes[-1]) + vecs_xy, vecs_NE = self.get_rose_vecs(rtp, q) + rose(fig, vecs_xy, p0=(0.15, 0.94)) + rose(fig, vecs_NE, p0=(0.85, 0.94)) + + return fig + \ No newline at end of file diff --git a/python/lsst/donut/viz/psf_from_zern.py b/python/lsst/donut/viz/psf_from_zern.py new file mode 100644 index 0000000..21f4cae --- /dev/null +++ b/python/lsst/donut/viz/psf_from_zern.py @@ -0,0 +1,71 @@ +import numpy as np +from matplotlib.figure import Figure +from matplotlib.gridspec import GridSpec + +def psfPanel( + xs, + ys, + psf, + detname, + fig=None, + figsize=(11,14), + cmap="cool", + **kwargs +) -> Figure: + """Make a per-detector psf scatter plot + + Subplots shows for each detector the psf retrieve from the zernike value + for each pair of intra-extra focal images. The points are placed using + pixel coordinates. + + Parameters + ---------- + xs, ys: array of float, shape (ndet, npair) + Points coordinates in pixel. + psf: array of float, shape (ndet, npair) + PSF value for each point. + detname: list of strings, shape (ndet,) + Detector labels. + fig: matplotlib Figure, optional + If provided, use this figure. Default None. + figsize: tuple of float, optional + Figure size in inches. Default (11, 12). + cmap: str, optional + Colormap name. Default 'seismic'. + **kwargs: + Additional keyword arguments passed to matplotlib Figure constructor. + + Returns + ------- + fig: matplotlib Figure + The figure. + """ + + # generating figure if None + if fig is None: + fig = Figure(figsize=figsize, **kwargs) + + # creating the gridspec grid (3x3 equal axes and the bottom cbar ax) + gs = GridSpec(nrows=4, ncols=3, figure=fig, width_ratios=[1, 1, 1], height_ratios=[1, 1, 1, 0.1]) + axs = [fig.add_subplot(gs[i, j]) for i in range(3) for j in range(3)] + ax_cbar = fig.add_subplot(gs[-1, :]) + + # setting the detector size (maybe there is a more wise way to retrieve it from the data metadata) + det_lim_y = (0., 4000.) + det_lim_x = (0., 4072.) + + # setting the common colormap limits + pmax = np.nanmax(psf) + pmin = np.nanmin(psf) + + # cycling through the axes. + for i, ax in enumerate(axs): + im = ax.scatter(xs[i], ys[i], c=psf[i], cmap=cmap, vmax=pmax, vmin=pmin) + ax.set_title(f"{detname[i]}: {np.nanmean(psf[i]):.3f} +/- {np.nanstd(psf[i]):.3f}") + ax.set(xlim=det_lim_x, ylim=det_lim_y, xticks=[], yticks=[], aspect="equal") + + # setting the colorbar + cb = fig.colorbar(im, cax=ax_cbar, location="bottom") + cb.set_label(label="PSF width, arcsecond", fontsize="large") + + return fig \ No newline at end of file