Merge pull request #120 from PaulHuwe/RCAL-789_L3SourceInjection

RCAL-789: Source Injection - Level 3

romanisim/l3.py

@@ -0,0 +1,80 @@
+"""Function for Level 3-like images.
+
+"""
+
+import numpy as np
+import galsim
+from . import log
+from . import image, wcs, psf
+
+
+def add_objects_to_l3(l3_mos, source_cat, rng=None, seed=None):
+    """Add objects to a Level 3 mosaic
+
+    Parameters
+    ----------
+    l3_mos : MosaicModel
+        Mosaic of images
+    source_cat : list
+        List of catalog objects to add to l3_mos
+
+    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
+
+    # l3_mos is updated in place, so no return
+    return None

romanisim/parameters.py

@@ -55,6 +55,23 @@
                 },
     'aperture': {'name': 'WFI_CEN',
                  'position_angle': 0
+                 },
+}
+
+# Default metadata for level 3 mosaics
+default_mosaic_parameters_dictionary = {
+    'basic': {'time_mean_mjd': Time('2026-01-01T00:00:00').mjd,
+              'optical_element': 'F184',
+              },
+    'wcsinfo': {'ra_ref': 270.0,
+                'dec_ref': 66.0,
+                'v2_ref': 0,
+                'v3_ref': 0,
+                'roll_ref': 0,
+                'vparity': -1,
+                'v3yangle': -60.0,
+                # I don't know what vparity and v3yangle should really be,
+                # but they are always -1 and -60 in existing files.
                 },
 }
 

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
+from romanisim import image, parameters, catalog, psf, util, wcs, persistence, ramp, l1, l3
 from astropy.coordinates import SkyCoord
 from astropy import units as u
 from astropy.time import Time
@@ -30,6 +30,7 @@
 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
 
 
@@ -754,3 +755,95 @@ def test_inject_source_into_image():
                    'flux': flux,
                    'tij': tij}
         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'))

romanisim/tests/test_wcs.py

@@ -13,10 +13,7 @@
 import pytest
 
 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
 
 
 def make_fake_distortion_function():

romanisim/wcs.py

@@ -77,8 +77,8 @@ def fill_in_parameters(parameters, coord, pa_aper=0, boresight=True):
         parameters['pointing']['dec_v1'])
 
     # Romanisim uses ROLL_REF = PA_APER - V3IdlYAngle
-    V3IdlYAngle = -60 # this value should eventually be taken from the SIAF
-    parameters['wcsinfo']['roll_ref'] = pa_aper - V3IdlYAngle 
+    V3IdlYAngle = -60  # this value should eventually be taken from the SIAF
+    parameters['wcsinfo']['roll_ref'] = pa_aper - V3IdlYAngle
 
     if boresight:
         parameters['wcsinfo']['v2_ref'] = 0
@@ -406,3 +406,44 @@ def convert_wcs_to_gwcs(wcs):
     else:
         # make a gwcs WCS from a galsim.roman WCS
         return wcs_from_fits_header(wcs.header.header)
+
+
+def get_mosaic_wcs(mosaic):
+    """Get a WCS object for a given sca or set of CRDS parameters.
+
+    Parameters
+    ----------
+    mosaic : roman_datamodels.datamodels.MosaicModel or dict
+        Mosaic model or dictionary containing WCS parameters.
+
+    Returns
+    -------
+    galsim.CelestialWCS for the mosaic
+    """
+
+    # If sent a dictionary, create a temporary model for data interface
+    if (type(mosaic) is not roman_datamodels.datamodels.MosaicModel):
+        mosaic_node = maker_utils.mk_level3_mosaic()
+        for key in mosaic.keys():
+            if isinstance(mosaic[key], dict):
+                mosaic_node['meta'][key].update(mosaic[key])
+            else:
+                mosaic_node['meta'][key] = mosaic[key]
+        mosaic_node = roman_datamodels.datamodels.MosaicModel(mosaic_node)
+    else:
+        mosaic_node = mosaic
+
+    world_pos = astropy.coordinates.SkyCoord(
+        mosaic_node.meta.wcsinfo.ra_ref * u.deg,
+        mosaic_node.meta.wcsinfo.dec_ref * u.deg)
+
+    # Create a tangent plane WCS for the mosaic
+    # The affine parameters below should be reviewed and updated
+    affine = galsim.AffineTransform(
+        0.1, 0, 0, 0.1, origin=galsim.PositionI(mosaic_node.data.shape[0] / 2.0,
+                                                mosaic_node.data.shape[1] / 2.0,),
+        world_origin=galsim.PositionD(0, 0))
+    wcs = galsim.TanWCS(affine,
+                        util.celestialcoord(world_pos))
+
+    return wcs