diff --git a/doc/versionHistory.rst b/doc/versionHistory.rst index 4e7bab8..e8ae878 100644 --- a/doc/versionHistory.rst +++ b/doc/versionHistory.rst @@ -4,6 +4,14 @@ Version History ################## +.._lsst.ts.donut.viz-1.2.1 + +------------- +1.2.1 +------------- + +* Add PlotPsfZernTask that creates a per-detector scatter plot of the PSF calculated with the convertZernikesToPsfWidth method. + .. _lsst.ts.donut.viz-1.2.0: ------------- diff --git a/pipelines/scienceDirectDetectPipeline.yaml b/pipelines/scienceDirectDetectPipeline.yaml index bbc6d74..e6f8ded 100644 --- a/pipelines/scienceDirectDetectPipeline.yaml +++ b/pipelines/scienceDirectDetectPipeline.yaml @@ -56,3 +56,7 @@ tasks: class: lsst.donut.viz.PlotDonutTask config: doRubinTVUpload: false + plotPsfZernTask: + class: lsst.donut.viz.PlotPsfZernTask + config: + doRubinTVUpload: false diff --git a/python/lsst/donut/viz/plot_aos_task.py b/python/lsst/donut/viz/plot_aos_task.py index 3cd6214..a436f6f 100644 --- a/python/lsst/donut/viz/plot_aos_task.py +++ b/python/lsst/donut/viz/plot_aos_task.py @@ -8,8 +8,11 @@ import matplotlib.pyplot as plt import numpy as np import yaml +from astropy import units as u +from lsst.ts.wep.utils import convertZernikesToPsfWidth from lsst.utils.timer import timeMethod +from .psf_from_zern import psfPanel from .utilities import ( add_rotated_axis, get_day_obs_seq_num_from_visitid, @@ -30,6 +33,9 @@ "PlotDonutTaskConnections", "PlotDonutTaskConfig", "PlotDonutTask", + "PlotPsfZernTaskConnections", + "PlotPsfZernTaskConfig", + "PlotPsfZernTask", ] @@ -374,3 +380,152 @@ 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.run(zernikes, figsize=(11, 14)) + + 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=psf_zk_panel, + ) + + def run(self, zernikes, **kwargs) -> plt.figure: + """Run the PlotPsfZern AOS task. + + This task create a 3x3 grid of subplots, + each subplot shows the psf value calculated from the Zernike + coefficients for each pair of intra-extra donuts found + for each detector. + + Parameters + ---------- + zernikes: list of tables. + List of tables containing the zernike sets + for each donut pair in each detector. + **kwargs: + Additional keyword arguments passed to + matplotlib.pyplot.figure constructor. + + Returns + ------- + fig: matplotlib.pyplot.figure + The figure. + """ + + 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)) + + psf = [ + [ + np.sqrt(np.sum(convertZernikesToPsfWidth(pair_zset) ** 2)) + for pair_zset in det + ] + for det in zs + ] + + q = qt.meta["extra"]["boresight_par_angle_rad"] + rot = qt.meta["extra"]["boresight_rot_angle_rad"] + rtp = q - rot - np.pi / 2 + + 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)), + } + + fig = plt.figure(**kwargs) + fig.suptitle( + f"PSF from Zernikes\nvisit: {zernikes[-1].meta['extra']['visit']}", + fontsize="xx-large", + fontweight="book", + ) + fig = psfPanel(xs, ys, psf, dname, fig=fig) + + # draw rose + rose(fig, vecs_xy, p0=(0.15, 0.94)) + rose(fig, vecs_NE, p0=(0.85, 0.94)) + + return fig 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..a1c5455 --- /dev/null +++ b/python/lsst/donut/viz/psf_from_zern.py @@ -0,0 +1,104 @@ +import numpy as np +from matplotlib.figure import Figure +from matplotlib.gridspec import GridSpec + + +def psfPanel( + xs, + ys, + psf, + detname, + dettype="LSSTComCam", + fig=None, + figsize=(11, 14), + maxcol=3, + cmap="cool", + **kwargs, +) -> Figure: + """Make a per-detector psf scatter plot + + Subplots shows for each detector the psf retrieved from the zernike value + for each pair of intra-extra focal images. The points are placed using + pixel coordinates. + + Parameters + ---------- + xs, ys: list of list[float], shape (ndet, npair) + Points coordinates in pixel. + psf: list of list[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 'cool'. + maxcol: int, optional + Maximum number of columns to use while creating the subplots grid. + Default 3 + **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) + if len(detname) < maxcol: + det_nrows = 1 + ncols = len(detname) + else: + det_nrows = ( + len(detname) // maxcol + 1 + if len(detname) % maxcol != 0 + else len(detname) // maxcol + ) + ncols = maxcol + + gs = GridSpec( + nrows=det_nrows + 1, # add a final row for the cbar + ncols=ncols, + figure=fig, + width_ratios=[1.0] * ncols, + height_ratios=[1.0] * det_nrows + [0.1], + ) + axs = [] + for i in range(len(detname)): + axs.append(fig.add_subplot(gs[i // ncols, i % ncols])) + 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) + match dettype: + case "LSSTComCam": + det_lim_y = (0.0, 4000.0) + det_lim_x = (0.0, 4072.0) + case "LSSTCam": + det_lim_y = (0.0, 2000.0) + det_lim_x = (0.0, 4072.0) + case _: + raise ValueError("Detector type not known") + + # setting the common colormap limits + pmax = np.nanmax(psf) + pmin = np.nanmin(psf) + + # cycling through the axes. + for i, dn in enumerate(detname): + im = axs[i].scatter(xs[i], ys[i], c=psf[i], cmap=cmap, vmax=pmax, vmin=pmin) + axs[i].set_title(f"{dn}: {np.nanmean(psf[i]):.3f} +/- {np.nanstd(psf[i]):.3f}") + axs[i].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 diff --git a/tests/test_donut_viz_pipeline_science_sensors.py b/tests/test_donut_viz_pipeline_science_sensors.py index 330cf9e..2b1b3a0 100644 --- a/tests/test_donut_viz_pipeline_science_sensors.py +++ b/tests/test_donut_viz_pipeline_science_sensors.py @@ -229,3 +229,12 @@ def testDonutPlotTask(self): self.butler.query_datasets("donutPlotExtra", collections=self.test_run_name) ) self.assertEqual(len(extra_dataset_list), 1) + + def testPlotPsfZernTask(self): + # Test that plots exist in butler + psf_zern_dataset_list = list( + self.butler.query_datasets( + "psfFromZernPanel", collections=self.test_run_name + ) + ) + self.assertEqual(len(psf_zern_dataset_list), 1)