RCAL-862 Refactor L3 Injection

romanisim/image.py

@@ -29,10 +29,7 @@
 from romanisim import log
 
 import roman_datamodels
-try:
-    import roman_datamodels.maker_utils as maker_utils
-except ImportError:
-    import roman_datamodels.testing.utils as maker_utils
+import roman_datamodels.maker_utils as maker_utils
 
 
 # galsim fluxes are in photons / cm^2 / s
@@ -217,7 +214,8 @@ def trim_objlist(objlist, image):
 
 
 def add_objects_to_image(image, objlist, xpos, ypos, psf,
-                         flux_to_counts_factor, bandpass=None, filter_name=None,
+                         flux_to_counts_factor, convtimes=None,
+                         bandpass=None, filter_name=None,
                          rng=None, seed=None):
     """Add sources to an image.
 
@@ -227,17 +225,19 @@ def add_objects_to_image(image, objlist, xpos, ypos, psf,
     Parameters
     ----------
     image : galsim.Image
-        Image to which to add sources with associated WCS.
+        Image to which to add sources with associated WCS. Updated in place.
     objlist : list[CatalogObject]
         Objects to add to image
     xpos, ypos : array_like
         x & y positions of sources (pixel) at which sources should be added
     psf : galsim.Profile
         PSF for image
-    flux_to_counts_factor : float
+    flux_to_counts_factor : float or list
         physical fluxes in objlist (whether in profile SEDs or flux arrays)
         should be multiplied by this factor to convert to total counts in the
         image
+    convtimes: array_like
+        Exposure times with unit scaling to convert to output rate units
     bandpass : galsim.Bandpass
         bandpass in which image is being rendered.  This is used only in cases
         where chromatic profiles & PSFs are being used.
@@ -279,6 +279,7 @@ def add_objects_to_image(image, objlist, xpos, ypos, psf,
     for i, obj in enumerate(objlist):
         t0 = time.time()
         image_pos = galsim.PositionD(xpos[i], ypos[i])
+        pwcs = image.wcs.local(image_pos)
         profile = obj.profile
         if not chromatic:
             if obj.flux is None:
@@ -289,18 +290,22 @@ def add_objects_to_image(image, objlist, xpos, ypos, psf,
             psf0 = psf.at_position(xpos[i], ypos[i])
         else:
             psf0 = psf
-        final = galsim.Convolve(profile * flux_to_counts_factor, psf0)
+        factor = flux_to_counts_factor[i] if isinstance(flux_to_counts_factor, list) else flux_to_counts_factor
+        final = galsim.Convolve(profile * factor, psf0)
         if chromatic:
             stamp = final.drawImage(
-                bandpass, center=image_pos, wcs=image.wcs.local(image_pos),
+                bandpass, center=image_pos, wcs=pwcs,
                 method='phot', rng=rng)
         else:
             try:
                 stamp = final.drawImage(center=image_pos,
-                                        wcs=image.wcs.local(image_pos))
+                                        wcs=pwcs)
             except galsim.GalSimFFTSizeError:
                 log.warning(f'Skipping source {i} due to too '
                             f'large FFT needed for desired accuracy.')
+        if convtimes is not None:
+            stamp /= convtimes[i]
+
         bounds = stamp.bounds & image.bounds
         if bounds.area() > 0:
             image[bounds] += stamp[bounds]

romanisim/l3.py

@@ -1,80 +1,104 @@
-"""Function for Level 3-like images.
+"""Roman WFI simulator functions for Level 3 mosaics.
 
+Based on galsim's implementation of Roman image simulation.  Uses galsim Roman modules
+for most of the real work.
 """
 
 import numpy as np
+import math
+import astropy.time
+from astropy import table
 import galsim
-from . import log
-from . import image, wcs, psf
-
-
-def add_objects_to_l3(l3_mos, source_cat, rng=None, seed=None):
+from galsim import roman
+
+from . import parameters
+from . import util
+import romanisim.wcs
+import romanisim.l1
+import romanisim.bandpass
+import romanisim.psf
+import romanisim.image
+import romanisim.persistence
+from romanisim import log
+import roman_datamodels.maker_utils as maker_utils
+import roman_datamodels.datamodels as rdm
+from roman_datamodels.stnode import WfiMosaic
+
+
+def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords=None, cps_conv=1.0, unit_factor=1.0,
+                      filter_name=None, coords_unit='rad', wcs=None, psf=None, rng=None, seed=None):
     """Add objects to a Level 3 mosaic
 
     Parameters
     ----------
-    l3_mos : MosaicModel
+    l3_mos : MosaicModel or galsim.Image
         Mosaic of images
     source_cat : list
         List of catalog objects to add to l3_mos
+    exptimes : list
+        Exposure times to scale back to rate units
+    xpos, ypos : array_like
+        x & y positions of sources (pixel) at which sources should be added
+    coords : array_like
+        ra & dec positions of sources (coords_unit) at which sources should be added
+    cps_conv : float
+        Factor to convert data to cps
+    unit_factor: float
+        Factor to convert counts data to MJy / sr
+    coords_unit : string
+        units of coords
+    wcs : galsim.GSFitsWCS
+        WCS corresponding to image
+    psf : galsim.Profile
+        PSF for image
+    rng : galsim.BaseDeviate
+        random number generator to use
+    seed : int
+        seed to use for random number generator
 
     Returns
     -------
     None
         l3_mos is updated in place
     """
-
-    if rng is None and seed is None:
-        seed = 143
-        log.warning(
-            'No RNG set, constructing a new default RNG from default seed.')
-    if rng is None:
-        rng = galsim.UniformDeviate(seed)
-
     # Obtain optical element
-    filter_name = l3_mos.meta.basic.optical_element
-
-    # Generate WCS
-    twcs = wcs.get_mosaic_wcs(l3_mos.meta)
-
-    # Create PSF
-    l3_psf = psf.make_psf(filter_name=filter_name, sca=2, chromatic=False, webbpsf=True)
-
-    # Generate x,y positions for sources
-    coords = np.array([[o.sky_pos.ra.rad, o.sky_pos.dec.rad]
-                       for o in source_cat])
-    sourcecountsall = galsim.ImageF(l3_mos.data.shape[0], l3_mos.data.shape[1], wcs=twcs, xmin=0, ymin=0)
-    xpos, ypos = sourcecountsall.wcs._xy(coords[:, 0], coords[:, 1])
-    xpos_idx = [round(x) for x in xpos]
-    ypos_idx = [round(y) for y in ypos]
-
-    # Create overall scaling factor map
-    Ct_all = (l3_mos.data.value / l3_mos.var_poisson)
-
-    # Cycle over sources and add them to the mosaic
-    for idx, (x, y) in enumerate(zip(xpos_idx, ypos_idx)):
-        # Set scaling factor for injected sources
-        # Flux / sigma_p^2
-        Ct = (l3_mos.data[x][y].value / l3_mos.var_poisson[x][y].value)
-
-        # Create empty postage stamp galsim source image
-        sourcecounts = galsim.ImageF(l3_mos.data.shape[0], l3_mos.data.shape[1], wcs=twcs, xmin=0, ymin=0)
-
-        # Simulate source postage stamp
-        image.add_objects_to_image(sourcecounts, [source_cat[idx]], xpos=[xpos[idx]], ypos=[ypos[idx]],
-                                   psf=l3_psf, flux_to_counts_factor=Ct, bandpass=[filter_name],
-                                   filter_name=filter_name, rng=rng)
-
-        # Scale the source image back by its flux ratios
-        sourcecounts /= Ct
-
-        # Add sources to the original mosaic data array
-        l3_mos.data = (l3_mos.data.value + sourcecounts.array) * l3_mos.data.unit
-
-        # Note for the future - other noise sources (read and flat) need to be implemented
-
-    # Set new poisson variance
-    l3_mos.var_poisson = l3_mos.data.value / Ct_all
+    if filter_name is None:
+        filter_name = l3_mos.meta.basic.optical_element
+
+    # Generate WCS (if needed)
+    if wcs is None:
+        wcs = romanisim.wcs.get_mosaic_wcs(l3_mos.meta, shape=l3_mos.data.shape)
+
+    # Create PSF (if needed)
+    if psf is None:
+        psf = romanisim.psf.make_psf(filter_name=filter_name, sca=parameters.default_sca, chromatic=False, webbpsf=True)
+
+    # Create Image canvas to add objects to
+    if isinstance(l3_mos, (rdm.MosaicModel, WfiMosaic)):
+        sourcecountsall = galsim.ImageF(l3_mos.data.value, wcs=wcs, xmin=0, ymin=0)
+    else:
+        sourcecountsall = l3_mos
+
+    # Create position arrays (if needed)
+    if any(pos is None for pos in [xpos, ypos]):
+        # Create coordinates (if needed)
+        if coords is None:
+            coords = np.array([[o.sky_pos.ra.rad, o.sky_pos.dec.rad]
+                               for o in source_cat])
+            coords_unit = 'rad'
+        # Generate x,y positions for sources
+        xpos, ypos = sourcecountsall.wcs.radecToxy(coords[:, 0], coords[:, 1], coords_unit)
+
+    # Add sources to the original mosaic data array
+    romanisim.image.add_objects_to_image(sourcecountsall, source_cat, xpos=xpos, ypos=ypos,
+                                         psf=psf, flux_to_counts_factor=[xpt * cps_conv for xpt in exptimes],
+                                         convtimes=[xpt / unit_factor for xpt in exptimes],
+                                         bandpass=[filter_name], filter_name=filter_name,
+                                         rng=rng, seed=seed)
+
+    # Save array with added sources
+    if isinstance(l3_mos, (rdm.MosaicModel, WfiMosaic)):
+        l3_mos.data = sourcecountsall.array * l3_mos.data.unit
 
     # l3_mos is updated in place, so no return
     return None

romanisim/parameters.py

@@ -84,6 +84,17 @@
     "linearity": None,
     "readnoise": 5.0 * u.DN,
     "saturation": 55000 * u.DN,
+    # Taken from roman_wfi_photom_0046.asdf
+    "photom": {"pixelareasr": {"F062": 2.640600009241359e-13 * u.sr,
+                               "F087": 2.640600009241359e-13 * u.sr,
+                               "F106": 2.640600009241359e-13 * u.sr,
+                               "F129": 2.640600009241359e-13 * u.sr,
+                               "F146": 2.640600009241359e-13 * u.sr,
+                               "F158": 2.640600009241359e-13 * u.sr,
+                               "F184": 2.640600009241359e-13 * u.sr,
+                               "F213": 2.640600009241359e-13 * u.sr,
+                               }
+               }
 }
 
 nborder = 4  # number of border pixels used for reference pixels.
@@ -140,6 +151,9 @@
     "pixel_depth": 5,
 }
 
+# Centermost PSF to use for mosaic creation
+default_sca = 2
+
 # Persistence defaults
 # delete persistence records fainter than 0.01 electron / s
 # e.g., MA table 1 has ~144 s, so this is ~1 electron over the whole exposure.

romanisim/psf.py

@@ -201,8 +201,8 @@ def at_position(self, x, y):
         # x = [0, off] -> 1
         # x = [npix, infinity] -> 0
         # linearly between those, likewise for y.
-        out = (self.psf['ll'] * wleft * wlow +
-               self.psf['lr'] * (1 - wleft) * wlow +
-               self.psf['ul'] * wleft * (1 - wlow) +
-               self.psf['ur'] * (1 - wleft) * (1 - wlow))
+        out = (self.psf['ll'] * wleft * wlow
+               + self.psf['lr'] * (1 - wleft) * wlow
+               + self.psf['ul'] * wleft * (1 - wlow)
+               + self.psf['ur'] * (1 - wleft) * (1 - wlow))
         return out

romanisim/tests/test_image.py

@@ -18,7 +18,7 @@
 import numpy as np
 import galsim
 from galsim import roman
-from romanisim import image, parameters, catalog, psf, util, wcs, persistence, ramp, l1, l3
+from romanisim import image, parameters, catalog, psf, util, wcs, persistence, ramp, l1
 from astropy.coordinates import SkyCoord
 from astropy import units as u
 from astropy.time import Time
@@ -30,7 +30,6 @@
 from metrics_logger.decorators import metrics_logger
 from romanisim import log
 from roman_datamodels.stnode import WfiScienceRaw, WfiImage
-import roman_datamodels.maker_utils as maker_utils
 import romanisim.bandpass
 
 
@@ -732,98 +731,6 @@ def test_inject_source_into_image():
         af.write_to(os.path.join(artifactdir, 'dms231.asdf'))
 
 
-@metrics_logger("DMS232")
-@pytest.mark.soctests
-def test_inject_sources_into_mosaic():
-    """Inject sources into a mosaic.
-    """
-
-    # Set constants and metadata
-    galsim.roman.n_pix = 200
-    rng_seed = 42
-    metadata = copy.deepcopy(parameters.default_mosaic_parameters_dictionary)
-    metadata['basic']['optical_element'] = 'F158'
-
-    # Create WCS
-    twcs = wcs.get_mosaic_wcs(metadata)
-
-    # Create initial Level 3 mosaic
-
-    # Create Four-quadrant pattern of gaussian noise, centered around one
-    # Each quadrant's gaussian noise scales like total exposure time
-    # (total files contributed to each quadrant)
-
-    # Create gaussian noise generators
-    g1 = galsim.GaussianDeviate(rng_seed, mean=1.0, sigma=0.01)
-    g2 = galsim.GaussianDeviate(rng_seed, mean=1.0, sigma=0.02)
-    g3 = galsim.GaussianDeviate(rng_seed, mean=1.0, sigma=0.05)
-    g4 = galsim.GaussianDeviate(rng_seed, mean=1.0, sigma=0.1)
-
-    # Create level 3 mosaic model
-    l3_mos = maker_utils.mk_level3_mosaic(shape=(galsim.roman.n_pix, galsim.roman.n_pix))
-
-    # Update metadata in the l3 model
-    for key in metadata.keys():
-        if key in l3_mos.meta:
-            l3_mos.meta[key].update(metadata[key])
-
-    # Populate the mosaic data array with gaussian noise from generators
-    g1.generate(l3_mos.data.value[0:100, 0:100])
-    g2.generate(l3_mos.data.value[0:100, 100:200])
-    g3.generate(l3_mos.data.value[100:200, 0:100])
-    g4.generate(l3_mos.data.value[100:200, 100:200])
-
-    # Define Poisson Noise of mosaic
-    l3_mos.var_poisson.value[0:100, 0:100] = 0.01**2
-    l3_mos.var_poisson.value[0:100, 100:200] = 0.02**2
-    l3_mos.var_poisson.value[100:200, 0:100] = 0.05**2
-    l3_mos.var_poisson.value[100:200, 100:200] = 0.1**2
-
-    # Create normalized psf source catalog (same source in each quadrant)
-    sc_dict = {"ra": 4 * [0.0], "dec": 4 * [0.0], "type": 4 * ["PSF"], "n": 4 * [-1.0],
-               "half_light_radius": 4 * [0.0], "pa": 4 * [0.0], "ba": 4 * [1.0], "F158": 4 * [1.0]}
-    sc_table = table.Table(sc_dict)
-
-    xpos, ypos = 50, 50
-    sc_table["ra"][0], sc_table["dec"][0] = (twcs._radec(xpos, ypos) * u.rad).to(u.deg).value
-    xpos, ypos = 50, 150
-    sc_table['ra'][1], sc_table['dec'][1] = (twcs._radec(xpos, ypos) * u.rad).to(u.deg).value
-    xpos, ypos = 150, 50
-    sc_table['ra'][2], sc_table['dec'][2] = (twcs._radec(xpos, ypos) * u.rad).to(u.deg).value
-    xpos, ypos = 150, 150
-    sc_table['ra'][3], sc_table['dec'][3] = (twcs._radec(xpos, ypos) * u.rad).to(u.deg).value
-
-    source_cat = catalog.table_to_catalog(sc_table, ["F158"])
-
-    # Copy original Mosaic before adding sources
-    l3_mos_orig = l3_mos.copy()
-
-    # Add source_cat objects to mosaic
-    l3.add_objects_to_l3(l3_mos, source_cat, seed=rng_seed)
-
-    # Ensure that every data pixel value has increased or
-    # remained the same with the new sources injected
-    assert np.all(l3_mos.data.value >= l3_mos_orig.data.value)
-
-    # Ensure that every pixel's poisson variance has increased or
-    # remained the same with the new sources injected
-    # Numpy isclose is needed to determine equality, due to float precision issues
-    close_mask = np.isclose(l3_mos.var_poisson.value, l3_mos_orig.var_poisson.value, rtol=1e-06)
-    assert np.all(l3_mos.var_poisson.value[~close_mask] > l3_mos_orig.var_poisson.value[~close_mask])
-
-    # Create log entry and artifacts
-    log.info('DMS232 successfully injected sources into a mosiac at points (50,50), (50,150), (150,50), (150,150).')
-
-    artifactdir = os.environ.get('TEST_ARTIFACT_DIR', None)
-    if artifactdir is not None:
-        af = asdf.AsdfFile()
-        af.tree = {'l3_mos': l3_mos,
-                   'l3_mos_orig': l3_mos_orig,
-                   'source_cat_table': sc_table,
-                   }
-        af.write_to(os.path.join(artifactdir, 'dms232.asdf'))
-
-
 @metrics_logger("DMS228")
 @pytest.mark.soctests
 def test_image_input(tmpdir):