Fix scipy.fsolve failure when computing r0_500

python/desc/imsim/atmPSF.py

@@ -81,11 +81,14 @@ def __init__(self, airmass, rawSeeing, band, rng, t0=0.0, exptime=30.0, kcrit=0.
     @staticmethod
     def _seeing_resid(r0_500, wavelength, L0, target):
         """Residual function to use with `_r0_500` below."""
-        r0_500 = r0_500[0]
-        kolm_seeing = galsim.Kolmogorov(r0_500=r0_500, lam=wavelength).fwhm
-        r0 = r0_500 * (wavelength/500)**1.2
-        factor = np.sqrt(1. - 2.183*(r0/L0)**0.356)
-        return kolm_seeing*factor - target
+        r0_500 = np.atleast_1d(r0_500)
+        resids = np.empty_like(r0_500)
+        for i, this_r0_500 in enumerate(r0_500):
+            kolm_seeing = galsim.Kolmogorov(r0_500=this_r0_500, lam=wavelength).fwhm
+            r0 = this_r0_500 * (wavelength/500)**1.2
+            factor = np.sqrt(1. - 2.183*(r0/L0)**0.356)
+            resids[i] = kolm_seeing*factor - target
+        return resids
 
     @staticmethod
     def _r0_500(wavelength, L0, seeing):
@@ -110,10 +113,15 @@ def _getAtmKwargs(self):
         weights = np.clip(weights, 0.01, 0.8)  # keep weights from straying too far.
         weights /= np.sum(weights)  # renormalize
 
-        # Draw a single common outer scale for all layers from a log normal
+        # Draw outer scale from truncated log normal
         L0 = 0
         while L0 < 10.0 or L0 > 100:
             L0 = np.exp(gd() * 0.6 + np.log(25.0))
+        # Given the desired seeing500 and randomly selected L0, determine appropriate
+        # r0_500
+        r0_500 = AtmosphericPSF._r0_500(500.0, L0, self.seeing500)
+
+        # Broadcast common outer scale across all layers
         L0 = [L0 for _ in range(6)]
 
         # Uniformly draw layer speeds between 0 and max_speed.
@@ -122,10 +130,6 @@ def _getAtmKwargs(self):
         # Isotropically draw directions.
         directions = [ud()*360.0*galsim.degrees for _ in range(6)]
 
-        # Given the desired seeing500 and randomly selected L0, determine appropriate
-        # r0_500
-        r0_500 = AtmosphericPSF._r0_500(500.0, L0, self.seeing500)
-
         if self.logger:
             self.logger.debug("airmass = {}".format(self.airmass))
             self.logger.debug("seeing500 = {}".format(self.seeing500))

tests/test_atmPSF.py

@@ -0,0 +1,24 @@
+import unittest
+
+import numpy as np
+
+from desc.imsim.atmPSF import AtmosphericPSF
+
+class AtmPSF(unittest.TestCase):
+    def test_r0_500(self):
+        """Test that _seeing_resid has the API fsolve wants."""
+        for wavelength in [300.0, 500.0, 1100.0]:
+            for L0 in [10.0, 25.0, 100.0]:
+                for target_seeing in [0.5, 0.7, 1.0]:
+                    r0s = [0.1, 0.2]
+                    np.testing.assert_array_equal(
+                        np.hstack([
+                            AtmosphericPSF._seeing_resid(r0s[0], wavelength, L0, target_seeing),
+                            AtmosphericPSF._seeing_resid(r0s[1], wavelength, L0, target_seeing),
+                        ]),
+                        AtmosphericPSF._seeing_resid(r0s, wavelength, L0, target_seeing)
+                    )
+
+
+if __name__ == '__main__':
+    unittest.main()