From bc3636eb1af3f5188116c7879f34e02ffe8ab813 Mon Sep 17 00:00:00 2001
From: John Franklin Crenshaw <jfcrenshaw@gmail.com>
Date: Thu, 21 Mar 2024 15:09:11 -0700
Subject: [PATCH] Fixed offAxis coeff for Danish. Catch GalsimFFT error.

---
 python/lsst/ts/wep/estimation/danish.py       | 45 +++++++++++++------
 tests/task/test_calcZernikesDanishTaskCwfs.py |  2 +-
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/python/lsst/ts/wep/estimation/danish.py b/python/lsst/ts/wep/estimation/danish.py
index 8264a4740..54f95bb3d 100644
--- a/python/lsst/ts/wep/estimation/danish.py
+++ b/python/lsst/ts/wep/estimation/danish.py
@@ -26,6 +26,7 @@
 
 import danish
 import numpy as np
+from galsim import GalSimFFTSizeError
 from lsst.ts.wep import Image, ImageMapper, Instrument
 from lsst.ts.wep.estimation.wfAlgorithm import WfAlgorithm
 from scipy.ndimage import binary_erosion
@@ -101,6 +102,7 @@ def history(self) -> dict:
             - "zkResid" - the Zernike coefficients fit to the donut
             - "zkSum" - zkResid + the intrinsic Zernikes
             - "model" - the final forward-modeled donut image
+            - "GalSimFFTSizeError" - whether this was hit during least_squares
 
         Note the units for all Zernikes are in meters, and all Zernikes start
         with Noll index 4.
@@ -150,6 +152,7 @@ def _estimateSingleZk(
             *image.fieldAngle,
             image.defocalType,
             image.bandLabel,
+            jmaxIntrinsic=jmax,
             return4Up=False,
         )
         size = max(zkStart.size, offAxisCoeff.size)
@@ -189,20 +192,35 @@ def _estimateSingleZk(
         x0 = [0.0, 0.0, 1.0] + [0.0] * (jmax - 3)
 
         # Use scipy to optimize the parameters
-        result = least_squares(
-            model.chi,
-            jac=model.jac,
-            x0=x0,
-            args=(img, backgroundStd**2),
-            **self.lstsqKwargs,
-        )
+        try:
+            result = least_squares(
+                model.chi,
+                jac=model.jac,
+                x0=x0,
+                args=(img, backgroundStd**2),
+                **self.lstsqKwargs,
+            )
+            result = dict(result)
+
+            # Unpack the parameters
+            dx, dy, fwhm, *zkResid = result["x"]
+            zkResid = np.array(zkResid)
+
+            # Add the starting zernikes back into the result
+            zkSum = zkResid + zkStart[4:]
+
+            # Flag that we didn't hit GalSimFFTSizeError
+            galSimFFTSizeError = False
 
-        # Unpack the parameters
-        dx, dy, fwhm, *zkResid = result.x
-        zkResid = np.array(zkResid)
+        # Sometimes this happens with Danish :(
+        except GalSimFFTSizeError:
+            # Fill dummy objects
+            result = None
+            zkResid = np.full(jmax - 3, np.nan)
+            zkSum = np.full(jmax - 3, np.nan)
 
-        # Add the starting zernikes back into the result
-        zkSum = zkResid + zkStart[4:]
+            # Flag the error
+            galSimFFTSizeError = True
 
         if saveHistory:
             # Save the history
@@ -210,7 +228,7 @@ def _estimateSingleZk(
                 "image": img.copy(),
                 "variance": backgroundStd**2,
                 "zkStart": zkStart.copy(),
-                "lstsqResult": dict(result),
+                "lstsqResult": result,
                 "zkResid": zkResid.copy(),
                 "zkSum": zkSum.copy(),
                 "model": model.model(
@@ -221,6 +239,7 @@ def _estimateSingleZk(
                     sky_level=backgroundStd,
                     flux=img.sum(),
                 ),
+                "GalSimFFTSizeError": galSimFFTSizeError,
             }
 
         else:
diff --git a/tests/task/test_calcZernikesDanishTaskCwfs.py b/tests/task/test_calcZernikesDanishTaskCwfs.py
index a91969ef7..70df6e6a8 100644
--- a/tests/task/test_calcZernikesDanishTaskCwfs.py
+++ b/tests/task/test_calcZernikesDanishTaskCwfs.py
@@ -212,7 +212,7 @@ def testEstimateCornerZernikes(self):
         # Make sure the total rms error is less than 0.35 microns off
         # from the OPD truth as a sanity check
         self.assertLess(
-            np.sqrt(np.sum(np.square(zernCoeffAvgR40 - trueZernCoeffR40))), 0.39
+            np.sqrt(np.sum(np.square(zernCoeffAvgR40 - trueZernCoeffR40))), 0.35
         )
 
     def testGetCombinedZernikes(self):