Merge pull request #67 from DukeCosmology/enable_achromatic_drawing

enable achromatic drawing

Author: aguinot

roman_imsim/skycat.py

@@ -208,14 +208,10 @@ def getObj(self, index, gsparams=None, rng=None, exptime=30):
         for component in gsobjs:
             if faint:
                 gsobjs[component] = gsobjs[component].evaluateAtWavelength(self.bandpass)
-                gs_obj_list.append(
-                    gsobjs[component] * self._trivial_sed * self.exptime * roman.collecting_area
-                )
+                gs_obj_list.append(gsobjs[component] * self._trivial_sed)
             else:
                 if component in seds:
-                    gs_obj_list.append(
-                        gsobjs[component] * seds[component] * self.exptime * roman.collecting_area
-                    )
+                    gs_obj_list.append(gsobjs[component] * seds[component])
 
         if not gs_obj_list:
             return None
@@ -225,14 +221,22 @@ def getObj(self, index, gsparams=None, rng=None, exptime=30):
         else:
             gs_object = galsim.Add(gs_obj_list)
 
+        # This should catch both "star" and "gaia_star" objects
+        if "star" in skycat_obj.object_type:
+            # Cap (star) flux at 30M photons to avoid gross artifacts when trying
+            # to draw the Roman PSF in finite time and memory
+            flux_cap = 3e7
+            if flux > flux_cap:
+                flux = flux_cap
+
         # Give the object the right flux
+        gs_object = gs_object.withFlux(flux, self.bandpass)
         gs_object.flux = flux
-        gs_object.withFlux(gs_object.flux, self.bandpass)
 
         # Get the object type
         if (skycat_obj.object_type == "diffsky_galaxy") | (skycat_obj.object_type == "galaxy"):
             gs_object.object_type = "galaxy"
-        if skycat_obj.object_type == "star":
+        if skycat_obj.object_type in {"star", "gaia_star"}:
             gs_object.object_type = "star"
         if skycat_obj.object_type == "snana":
             gs_object.object_type = "transient"

roman_imsim/stamp.py

@@ -60,18 +60,6 @@ def setup(self, config, base, xsize, ysize, ignore, logger):
             # or cached by the skyCatalogs code.
             gal.flux = gal.calculateFlux(bandpass)
         self.flux = gal.flux
-        # Cap (star) flux at 30M photons to avoid gross artifacts when trying
-        # to draw the Roman PSF in finite time and memory
-        flux_cap = 3e7
-        if self.flux > flux_cap:
-            if (
-                hasattr(gal, "original")
-                and hasattr(gal.original, "original")
-                and isinstance(gal.original.original, galsim.DeltaFunction)
-            ) or (isinstance(gal, galsim.DeltaFunction)):
-                gal = gal.withFlux(flux_cap, bandpass)
-                self.flux = flux_cap
-                gal.flux = flux_cap
         base["flux"] = gal.flux
         base["mag"] = -2.5 * np.log10(gal.flux) + bandpass.zeropoint
         # print('stamp setup2',process.memory_info().rss)
@@ -284,9 +272,6 @@ def draw(self, prof, image, method, offset, config, base, logger):
         # print('stamp draw2',process.memory_info().rss)
 
         if method == "phot":
-            # We already calculated realized_flux above.  Use that now and tell GalSim not
-            # recalculate the Poisson realization of the flux.
-            gal = gal.withFlux(self.realized_flux, bandpass)
             # print('stamp draw3b ',process.memory_info().rss)
 
             if not faint and "photon_ops" in config:
@@ -306,7 +291,7 @@ def draw(self, prof, image, method, offset, config, base, logger):
 
             # print('stamp draw3a',process.memory_info().rss)
             gal.drawImage(
-                bandpass,
+                bandpass=bandpass,
                 method="phot",
                 offset=offset,
                 rng=self.rng,
@@ -316,7 +301,7 @@ def draw(self, prof, image, method, offset, config, base, logger):
                 photon_ops=photon_ops,
                 sensor=None,
                 add_to_image=True,
-                poisson_flux=False,
+                poisson_flux=True,
             )
         else:
             fft_image = image.copy()
@@ -337,10 +322,9 @@ def draw(self, prof, image, method, offset, config, base, logger):
                 )
 
             # Go back to a combined convolution for fft drawing.
-            gal = gal.withFlux(self.flux, bandpass)
             prof = galsim.Convolve([gal] + psfs)
             try:
-                prof.drawImage(bandpass, **kwargs)
+                prof.drawImage(bandpass=bandpass, **kwargs)
             except galsim.errors.GalSimFFTSizeError as e:
                 # I think this shouldn't happen with the updates I made to how the image size
                 # is calculated, even for extremely bright things.  So it should be ok to