Enable sparse zernikes to be used from ts_wep.

doc/versionHistory.rst

@@ -4,6 +4,14 @@
 Version History
 ##################
 
+.. _lsst.ts.donut.viz-1.2.0:
+
+-------------
+1.2.0
+-------------
+
+* Enable sparse zernikes to be used from ts_wep.
+
 .. _lsst.ts.donut.viz-1.1.2:
 
 -------------

pipelines/scienceDirectDetectPipeline.yaml

@@ -37,7 +37,7 @@ tasks:
   calcZernikesTask:
     class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask
     config:
-      estimateZernikes.maxNollIndex: 28
+      estimateZernikes.nollIndices: [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 27, 28]
       estimateZernikes.saveHistory: False
       estimateZernikes.maskKwargs: {'doMaskBlends': False}
   aggregateZernikeTablesTask:

python/lsst/donut/viz/aggregate_visit.py

@@ -87,11 +87,14 @@ def runQuantum(
             zernike_table = butlerQC.get(zernikesRef)
             raw_table = Table()
             zernikes_merged = []
+            noll_indices = []
             for col_name in zernike_table.colnames:
                 # Grab zernike output columns
                 if col_name.startswith("Z"):
                     zernikes_merged.append(zernike_table[col_name].to(u.um).value)
+                    noll_indices.append(int(col_name[1:]))
             zernikes_merged = np.array(zernikes_merged).T
+            noll_indices = np.array(noll_indices)
             raw_table["zk_CCS"] = np.atleast_2d(zernikes_merged[1:])
             raw_table["detector"] = camera[zernikesRef.dataId["detector"]].getName()
             raw_tables.append(raw_table)
@@ -123,17 +126,23 @@ def runQuantum(
         meta["alt"] = table_meta["extra"]["boresight_alt_rad"]
         # Average mjds
         meta["mjd"] = 0.5 * (table_meta["extra"]["mjd"] + table_meta["intra"]["mjd"])
+        meta["nollIndices"] = noll_indices
 
         q = meta["parallacticAngle"]
         rtp = meta["rotTelPos"]
 
-        jmax = out_raw["zk_CCS"].shape[1] + 3
+        jmin = np.min(noll_indices)
+        jmax = np.max(noll_indices)
         rot_OCS = galsim.zernike.zernikeRotMatrix(jmax, -rtp)[4:, 4:]
         rot_NW = galsim.zernike.zernikeRotMatrix(jmax, -q)[4:, 4:]
         for cat in (out_raw, out_avg):
             cat.meta = meta
-            cat["zk_OCS"] = cat["zk_CCS"] @ rot_OCS
-            cat["zk_NW"] = cat["zk_CCS"] @ rot_NW
+            full_zk_ccs = np.zeros((len(cat), jmax-jmin+1))
+            full_zk_ccs[:, noll_indices-4] = cat["zk_CCS"]
+            cat["zk_OCS"] = full_zk_ccs @ rot_OCS
+            cat["zk_NW"] = full_zk_ccs @ rot_NW
+            cat["zk_OCS"] = cat["zk_OCS"][:, noll_indices-4]
+            cat["zk_NW"] = cat["zk_NW"][:, noll_indices-4]
 
         # Find the right output references
         butlerQC.put(out_raw, outputRefs.aggregateZernikesRaw)

python/lsst/donut/viz/plot_aos_task.py

@@ -140,8 +140,9 @@ def doPyramid(
         zk,
         rtp,
         q,
+        nollIndices
     ):
-        fig = zernikePyramid(x, y, zk, cmap="seismic", s=10)
+        fig = zernikePyramid(x, y, zk, nollIndices, cmap="seismic", s=10)
         vecs_xy = {
             r"$x_\mathrm{Opt}$": (1, 0),
             r"$y_\mathrm{Opt}$": (0, -1),
@@ -175,8 +176,9 @@ def plotZernikePyramids(
         zk = aos_raw["zk_OCS"].T
         rtp = aos_raw.meta["rotTelPos"]
         q = aos_raw.meta["parallacticAngle"]
+        nollIndices = aos_raw.meta["nollIndices"]
 
-        zkPyramid = self.doPyramid(x, y, zk, rtp, q)
+        zkPyramid = self.doPyramid(x, y, zk, rtp, q, nollIndices)
 
         # We want residuals from the intrinsic design too.
         path = Path(__file__).parent.parent.parent.parent.parent / "data"
@@ -193,10 +195,10 @@ def plotZernikePyramids(
             [z.coef for z in dzs(aos_raw["thx_OCS"], aos_raw["thy_OCS"])]
         ).T[4:29]
         intrinsic = intrinsic[: len(zk)]
-        intrinsicPyramid = self.doPyramid(x, y, intrinsic, rtp, q)
+        intrinsicPyramid = self.doPyramid(x, y, intrinsic, rtp, q, nollIndices)
 
         resid = zk - intrinsic
-        residPyramid = self.doPyramid(x, y, resid, rtp, q)
+        residPyramid = self.doPyramid(x, y, resid, rtp, q, nollIndices)
 
         return zkPyramid, residPyramid, intrinsicPyramid
 

python/lsst/donut/viz/plot_zernike_pyramid.py

@@ -17,11 +17,13 @@ def plotZernikePyramid(
         instrument=instrument,
     )
     zs = cat["zs_OCS"]
+    noll_indices = cat.meta["nollIndices"]
 
     fig = zernikePyramid(
         cat["thx_OCS"],
         cat["thy_OCS"],
         zs.T,
+        noll_indices,
         cmap="seismic",
         s=2,
     )
@@ -61,6 +63,7 @@ def plotZernikePyramid(
         cat["thx_OCS"],
         cat["thy_OCS"],
         resid.T,
+        noll_indices,
         cmap="seismic",
         s=2,
     )
@@ -73,7 +76,8 @@ def plotZernikePyramid(
     fig3 = zernikePyramid(
         cat["thx_OCS"],
         cat["thy_OCS"],
-        intrinsic[:, 4:29].T,
+        intrinsic[:, noll_indices].T,
+        noll_indices,
         cmap="seismic",
         s=2,
     )

python/lsst/donut/viz/zernike_pyramid.py

@@ -22,7 +22,7 @@ def zernikePyramid(
     xs,
     ys,
     zs,
-    jmin=4,
+    noll_indices,
     figsize=(13, 8),
     vmin=-1,
     vmax=1,
@@ -48,8 +48,8 @@ def zernikePyramid(
         Zernike values.  First index labels the particular Zernike coefficient,
         second index labels spatial coordinate.  First index implicitly starts
         at j=jmin (defocus by default).
-    jmin: int, optional
-        Minimum Zernike to plot.  Default 4 (defocus).
+    noll_indices: np.ndarray
+        Noll indices for zernikes in zs.
     figsize: tuple of float, optional
         Figure size in inches.  Default (13, 8).
     vmin, vmax: float, optional
@@ -78,7 +78,9 @@ def zernikePyramid(
     fig: matplotlib Figure
         The figure.
     """
-    jmax = zs.shape[0] + jmin - 1
+    jmin = min(noll_indices)
+    jmax = max(noll_indices)
+    jdict = {x: y for x, y in zip(noll_indices-jmin, range(len(noll_indices)))}
     nmax, _ = galsim.zernike.noll_to_zern(jmax)
     nmin, _ = galsim.zernike.noll_to_zern(jmin)
 
@@ -98,7 +100,7 @@ def shiftAxes(axes, amt):
     axes = {}
     shiftLeft = []
     shiftRight = []
-    for j in range(jmin, jmax + 1):
+    for j in noll_indices:
         n, m = galsim.zernike.noll_to_zern(j)
         if n % 2 == 0:
             row, col = n - nmin, m // 2 + ncol // 2
@@ -114,6 +116,8 @@ def shiftAxes(axes, amt):
 
     cbar = {}
     for j, ax in axes.items():
+        if j not in noll_indices:
+            continue
         n, _ = galsim.zernike.noll_to_zern(j)
         ax.set_title("Z{}".format(j))
         if vdim:
@@ -125,7 +129,7 @@ def shiftAxes(axes, amt):
         scat = ax.scatter(
             xs,
             ys,
-            c=zs[j - jmin],
+            c=zs[jdict[j - jmin]],
             s=s,
             linewidths=0.5,
             cmap=cmap,

tests/test_donut_viz_pipeline_science_sensors.py

@@ -27,6 +27,7 @@ def setUpClass(cls):
             "rotAngle",
             "rotTelPos",
             "visit",
+            "nollIndices"
         ]
 
         cls.butler = Butler(cls.test_repo_dir)
@@ -106,6 +107,7 @@ def testAggregateDonuts(self):
             agg_donut_table.meta.keys(), ["extra", "intra", "average"]
         )
         donut_meta_keys = self.meta_keys + ["focusZ"]
+        donut_meta_keys.remove("nollIndices")
         for key in ["extra", "intra"]:
             self.assertCountEqual(agg_donut_table.meta[key].keys(), donut_meta_keys)
         donut_meta_keys.remove("focusZ")