Merge pull request #16 from lsst-ts/u/LRosignoli/psf_addon

Feature branch that add PSF-related plots.

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:
 
 -------------

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

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

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

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)