pre-commit run, added test unit, run() method implemented in plot_aos…

…_task

python/lsst/donut/viz/plot_aos_task.py

@@ -7,19 +7,19 @@
 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 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,
     get_instrument_channel_name,
     rose,
 )
 from .zernike_pyramid import zernikePyramid
-from .psf_from_zern import psfPanel
 
 try:
     from lsst.rubintv.production.uploaders import MultiUploader
@@ -437,15 +437,15 @@ def runQuantum(
         inputRefs: pipeBase.InputQuantizedConnection,
         outputRefs: pipeBase.OutputQuantizedConnection,
     ) -> None:
+
         zernikes = butlerQC.get(inputRefs.zernikes)
 
-        zkPanel = self.plotPsfFromZern(zernikes, figsize=(11, 14))
+        zkPanel = self.run(zernikes, figsize=(11, 14))
         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:
@@ -464,33 +464,28 @@ def runQuantum(
                     filename=psf_zk_panel,
                 )
 
-    def get_psf_degr(self, zset):
-        return np.sqrt(np.sum(convertZernikesToPsfWidth(zset) ** 2))
+    def run(self, zernikes, **kwargs) -> plt.figure:
+        """Run the PlotPsfZern AOS task.
 
-    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
+        This task create a 3x3 grid of subplots,
+        each axe contains the psf value for each pair of
+        intra-extra donuts found for each detector.
 
-    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)),
-        }
+        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.
 
-        return vecs_xy, vecs_NE
+        Returns
+        -------
+        fig: matplotlib.pyplot.figure
+            The figure.
+        """
 
-    def plotPsfFromZern(self, zernikes, **kwargs):
         xs = []
         ys = []
         zs = []
@@ -504,17 +499,36 @@ def plotPsfFromZern(self, zernikes, **kwargs):
                 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])
+        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 = psfPanel(xs, ys, psf, dname, fig=fig)
 
         # 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))
 

python/lsst/donut/viz/psf_from_zern.py

@@ -14,9 +14,9 @@ def psfPanel(
 
     Parameters
     ----------
-    xs, ys: array of float, shape (ndet, npair)
+    xs, ys: list of list[float], shape (ndet, npair)
         Points coordinates in pixel.
-    psf: array of float, shape (ndet, npair)
+    psf: list of list[float], shape (ndet, npair)
         PSF value for each point.
     detname: list of strings, shape (ndet,)
         Detector labels.

tests/test_donut_viz_pipeline_science_sensors.py

@@ -227,3 +227,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)