tickets/DM-47480: sparse Zernike estimation

doc/versionHistory.rst

@@ -6,6 +6,17 @@
 Version History
 ##################
 
+.. _lsst.ts.wep-13.0.0:
+
+-------------
+13.0.0
+-------------
+
+* enabled sparse Zernike estimation
+* removed most jmax and return4Up configs in favor of nollIndices configs
+* removed return4Up from estimator WfEstimator and WfAlgorithm
+* added makeSparse and makeDense to Zernike utils
+
 .. _lsst.ts.wep-12.7.0:
 
 -------------

pipelines/production/comCamRapidAnalysisPipeline.yaml

@@ -17,7 +17,7 @@ tasks:
   calcZernikesTask:
     class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask
     config:
-      estimateZernikes.maxNollIndex: 22
+      estimateZernikes.nollIndices: [4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 20, 21, 27, 28]
       estimateZernikes.convergeTol: 10.0e-9
       estimateZernikes.requireConverge: True
       estimateZernikes.saveHistory: False

pipelines/production/comCamSimRapidAnalysisPipeline.yaml

@@ -5,21 +5,24 @@ tasks:
     class: lsst.ts.wep.task.generateDonutDirectDetectTask.GenerateDonutDirectDetectTask
     config:
       donutSelector.useCustomMagLimit: True
-      donutSelector.sourceLimit: 5
+      donutSelector.sourceLimit: 20
   cutOutDonutsScienceSensorGroupTask:
     class: lsst.ts.wep.task.cutOutDonutsScienceSensorTask.CutOutDonutsScienceSensorTask
     config:
       python: |
         from lsst.ts.wep.task.pairTask import GroupPairer
         config.pairer.retarget(GroupPairer)
-      donutStampSize: 160
+      donutStampSize: 200
       initialCutoutPadding: 40
   calcZernikesTask:
     class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask
     config:
-      estimateZernikes.maxNollIndex: 28
+      estimateZernikes.nollIndices: [4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 20, 21, 27, 28]
+      estimateZernikes.convergeTol: 10.0e-9
+      estimateZernikes.requireConverge: True
       estimateZernikes.saveHistory: False
-      estimateZernikes.maskKwargs: {'doMaskBlends': False}
+      estimateZernikes.maskKwargs: { "doMaskBlends": False }
+      donutStampSelector.maxSelect: 5
   isr:
     class: lsst.ip.isr.IsrTask
     config:

python/lsst/ts/wep/estimation/danish.py

@@ -124,6 +124,7 @@ def history(self) -> dict:
         following entries
             - "image" - the image that is being fit
             - "variance" - the background variance that was used for fitting
+            - "nollIndices" - Noll indices for which Zernikes are estimated
             - "zkStart" - the starting Zernike coefficients
             - "lstsqResult" - dictionary of results returned by least_squares
             - "zkFit" - the Zernike coefficients fit to the donut
@@ -140,6 +141,7 @@ def _estimateSingleZk(
         self,
         image: Image,
         zkStart: np.ndarray,
+        nollIndices: np.ndarray,
         instrument: Instrument,
         factory: danish.DonutFactory,
         saveHistory: bool,
@@ -152,6 +154,8 @@ def _estimateSingleZk(
             The ts_wep image of the donut stamp
         zkStart : np.ndarray
             The starting point for the Zernikes
+        nollIndices : np.ndarray
+            Noll indices for which to estimate Zernikes
         instrument : Instrument
             The ts_wep Instrument
         factory : danish.DonutFactory
@@ -170,22 +174,16 @@ def _estimateSingleZk(
         if np.isfinite(image.blendOffsets).sum() > 0:
             warnings.warn("Danish is currently only setup for non-blended donuts.")
 
-        # Determine the maximum Noll index
-        jmax = len(zkStart) + 3
-
         # Create reference Zernikes by adding off-axis coefficients to zkStart
-        zkStart = np.pad(zkStart, (4, 0))
         offAxisCoeff = instrument.getOffAxisCoeff(
             *image.fieldAngle,
             image.defocalType,
             image.bandLabel,
-            jmaxIntrinsic=jmax,
-            return4Up=False,
+            nollIndicesModel=np.arange(0, 79),
+            nollIndicesIntr=nollIndices,
         )
-        size = max(zkStart.size, offAxisCoeff.size)
-        zkRef = np.zeros(size)
-        zkRef[: zkStart.size] = zkStart
-        zkRef[: offAxisCoeff.size] += offAxisCoeff
+        zkRef = offAxisCoeff.copy()
+        zkRef[nollIndices] += zkStart
 
         # Create the background mask if it does not exist
         if image.maskBackground is None:
@@ -213,14 +211,14 @@ def _estimateSingleZk(
         model = danish.SingleDonutModel(
             factory,
             z_ref=zkRef,
-            z_terms=np.arange(4, jmax + 1),
+            z_terms=nollIndices,
             thx=np.deg2rad(image.fieldAngle[0]),
             thy=np.deg2rad(image.fieldAngle[1]),
             npix=img.shape[0],
         )
 
         # Create the initial guess for the model parameters
-        x0 = [0.0, 0.0, 1.0] + [0.0] * (jmax - 3)
+        x0 = [0.0, 0.0, 1.0] + [0.0] * len(nollIndices)
 
         # Use scipy to optimize the parameters
         try:
@@ -238,7 +236,7 @@ def _estimateSingleZk(
             zkFit = np.array(zkFit)
 
             # Add the starting zernikes back into the result
-            zkSum = zkFit + zkStart[4:]
+            zkSum = zkFit + zkStart
 
             # Flag that we didn't hit GalSimFFTSizeError
             galSimFFTSizeError = False
@@ -258,8 +256,8 @@ def _estimateSingleZk(
         except GalSimFFTSizeError:
             # Fill dummy objects
             result = None
-            zkFit = np.full(jmax - 3, np.nan)
-            zkSum = np.full(jmax - 3, np.nan)
+            zkFit = np.full_like(zkStart, np.nan)
+            zkSum = np.full_like(zkStart, np.nan)
             if saveHistory:
                 modelImage = np.full_like(img, np.nan)
 
@@ -271,6 +269,7 @@ def _estimateSingleZk(
             hist = {
                 "image": img.copy(),
                 "variance": backgroundStd**2,
+                "nollIndices": nollIndices.copy(),
                 "zkStart": zkStart.copy(),
                 "lstsqResult": result,
                 "zkFit": zkFit.copy(),
@@ -290,6 +289,7 @@ def _estimateZk(
         I2: Optional[Image],
         zkStartI1: np.ndarray,
         zkStartI2: Optional[np.ndarray],
+        nollIndices: np.ndarray,
         instrument: Instrument,
         saveHistory: bool,
     ) -> np.ndarray:
@@ -305,6 +305,8 @@ def _estimateZk(
             The starting Zernikes for I1 (in meters, for Noll indices >= 4)
         zkStartI2 : np.ndarray or None
             The starting Zernikes for I2 (in meters, for Noll indices >= 4)
+        nollIndices : np.ndarray
+            Noll indices for which you wish to estimate Zernike coefficients.
         instrument : Instrument
             The Instrument object associated with the DonutStamps.
         saveHistory : bool
@@ -334,6 +336,7 @@ def _estimateZk(
         zk1, hist[I1.defocalType.value] = self._estimateSingleZk(
             I1,
             zkStartI1,
+            nollIndices,
             instrument,
             factory,
             saveHistory,
@@ -344,6 +347,7 @@ def _estimateZk(
             zk2, hist[I2.defocalType.value] = self._estimateSingleZk(
                 I2,
                 zkStartI2,
+                nollIndices,
                 instrument,
                 factory,
                 saveHistory,

python/lsst/ts/wep/estimation/tie.py

@@ -32,6 +32,7 @@
     binArray,
     createZernikeBasis,
     createZernikeGradBasis,
+    makeSparse,
 )
 from scipy.ndimage import gaussian_filter
 
@@ -472,6 +473,8 @@ def history(self) -> dict:
                               full OPD.
             - "pupil" - model pupil image, in the case where only one image
                         was passed to the estimator.
+            - "nollIndices" - array of Noll indices corresponding to the
+                              estimated Zernikes values.
 
         Each iteration of the solver is then stored under indices >= 1.
         The entry for each iteration is also a dictionary, containing
@@ -595,7 +598,7 @@ def _expSolve(
         self,
         I0: np.ndarray,
         dIdz: np.ndarray,
-        jmax: int,
+        nollIndices: np.ndarray,
         instrument: Instrument,
     ) -> np.ndarray:
         """Solve the TIE directly using a Zernike expansion.
@@ -606,8 +609,8 @@ def _expSolve(
             The beam intensity at the exit pupil
         dIdz : np.ndarray
             The z-derivative of the beam intensity across the exit pupil
-        jmax : int
-            The maximum Zernike Noll index to estimate
+        nollIndices : np.ndarray
+            Noll indices for which you wish to estimate Zernike coefficients.
         instrument : Instrument, optional
             The Instrument object associated with the DonutStamps.
             (the default is the default Instrument)
@@ -618,8 +621,11 @@ def _expSolve(
             Numpy array of the Zernike coefficients estimated from the image
             or pair of images, in nm.
         """
-        # Get Zernike Bases
+        # Get the pupil grid
         uPupil, vPupil = instrument.createPupilGrid()
+
+        # Get Zernike Bases
+        jmax = nollIndices.max()
         zk = createZernikeBasis(
             uPupil,
             vPupil,
@@ -633,6 +639,11 @@ def _expSolve(
             obscuration=instrument.obscuration,
         )
 
+        # Down-select to the Zernikes we are solving for
+        zk = makeSparse(zk, nollIndices)
+        dzkdu = makeSparse(dzkdu, nollIndices)
+        dzkdv = makeSparse(dzkdv, nollIndices)
+
         # Calculate quantities for the linear system
         # See Equations 43-45 of https://sitcomtn-111.lsst.io
         b = np.einsum("ab,jab->j", dIdz, zk, optimize=True)
@@ -651,6 +662,7 @@ def _estimateZk(
         I2: Optional[Image],  # type: ignore[override]
         zkStartI1: np.ndarray,
         zkStartI2: Optional[np.ndarray],
+        nollIndices: np.ndarray,
         instrument: Instrument,
         saveHistory: bool,
     ) -> np.ndarray:
@@ -667,6 +679,8 @@ def _estimateZk(
         zkStartI2 : np.ndarray or None
             The starting Zernikes for I2 (in meters, for Noll indices >= 4)
             Can be None.
+        nollIndices : np.ndarray
+            Noll indices for which you wish to estimate Zernike coefficients.
         instrument : Instrument
             The Instrument object associated with the DonutStamps.
         saveHistory : bool
@@ -685,9 +699,11 @@ def _estimateZk(
         RuntimeError
             If the solver is not supported
         """
+        # Rescale instrument pixel size to account for binning
         if self.binning > 1:
             instrument = instrument.copy()
             instrument.pixelSize *= self.binning
+
         # Create the ImageMapper for centering and image compensation
         imageMapper = ImageMapper(instConfig=instrument, opticalModel=self.opticalModel)
 
@@ -699,6 +715,7 @@ def _estimateZk(
             zkStartI2,
         )
 
+        # Bin images to reduce resolution?
         if self.binning > 1:
             if intra is not None:
                 intra.image = binArray(intra.image, self.binning)
@@ -731,6 +748,7 @@ def _estimateZk(
                 "zkStartExtra": None if extra is None else zkStartExtra.copy(),
                 "zkStartMean": zkStartMean.copy(),
                 "pupil": None if pupil is None else pupil.copy(),
+                "nollIndices": nollIndices.copy(),
             }
 
         # Initialize Zernike arrays at zero
@@ -744,7 +762,7 @@ def _estimateZk(
         compSequence = iter(self.compSequence)
 
         # Determine the maximum Noll index we will solve for
-        jmax = len(zkStartMean) + 3
+        jmax = nollIndices.max()
 
         # Set the caustic and converged flags to False
         caustic = False
@@ -760,7 +778,7 @@ def _estimateZk(
             # The gain scales how much of previous residual we incorporate
             # Everything past jmaxComp is set to zero
             zkComp += self.compGain * zkResid
-            zkComp[(jmaxComp - 3) :] = 0
+            zkComp[nollIndices > jmaxComp] = 0
 
             # Center the images
             recenter = (i == 0) or (np.max(np.abs(zkComp - zkCenter)) > self.centerTol)
@@ -773,6 +791,7 @@ def _estimateZk(
                     else imageMapper.centerOnProjection(
                         intra,
                         zkCenter + zkStartIntra,
+                        nollIndices,
                         isBinary=self.centerBinary,
                         **self.maskKwargs,
                     )
@@ -783,6 +802,7 @@ def _estimateZk(
                     else imageMapper.centerOnProjection(
                         extra,
                         zkCenter + zkStartExtra,
+                        nollIndices,
                         isBinary=self.centerBinary,
                         **self.maskKwargs,
                     )
@@ -795,6 +815,7 @@ def _estimateZk(
                 else imageMapper.mapImageToPupil(
                     intraCent,
                     zkComp + zkStartIntra,
+                    nollIndices,
                     masks=None if i == 0 else intraComp.masks,  # noqa: F821
                     **self.maskKwargs,
                 )
@@ -805,6 +826,7 @@ def _estimateZk(
                 else imageMapper.mapImageToPupil(
                     extraCent,
                     zkComp + zkStartExtra,
+                    nollIndices,
                     masks=None if i == 0 else extraComp.masks,  # noqa: F821
                     **self.maskKwargs,
                 )
@@ -842,7 +864,7 @@ def _estimateZk(
                 )
 
                 # Estimate the Zernikes
-                zkResid = self._expSolve(I0, dIdz, jmax, instrument)
+                zkResid = self._expSolve(I0, dIdz, nollIndices, instrument)
 
                 # If estimating with a single donut, double the coefficients
                 # This is because the simulated pupil image is aberration free