Refactor l2 injection (#138)

* Initial commit.

* Make test_catalog_and_images test faster.

* Ruff formatting changes.

romanisim/image.py

@@ -915,3 +915,137 @@ def make_asdf(slope, slopevar_rn, slopevar_poisson, metadata=None,
         af.tree = {'roman': out, 'romanisim': extras}
         af.write_to(filepath)
     return out, extras
+
+
+def inject_sources_into_l2(model, cat, x=None, y=None, psf=None, rng=None,
+                           gain=None, webbpsf=True):
+    """Inject sources into an L2 image.
+
+    This routine allows sources to be injected onto an existing L2 image.
+    Source injection into an L2 image relies on knowing the objects'
+    x and y locations, the PSF, and the image gain; if these are not provided,
+    reasonable defaults are generated.
+
+    The simulation proceeds by (optionally) using the model WCS to generate the
+    x & y locations, grabbing the gain from
+    romanisim.parameters.reference_data, and grabbing the read_pattern from the
+    model_metadata.  The number of additional counts in each pixel are
+    simulated.  We create a "virtual" ramp that uses the input L2 image and
+    evenly apportions the measured DN/s along the ramp using the MA table.  We
+    apportion the new counts to a new ramp, and add the new ramp to the virtual
+    ramp.  We then refit the new ramp, and replace the old fit with the new
+    fit.
+
+    This simulation is not as complete as the full L2 simulation.  We do not
+    include non-linearity or saturation, for example.  Identified CR hits
+    are lost.  But it should do a decent job at providing realistic
+    uncertainties otherwise, including how read & Poisson noise influence
+    the choice of weights used in ramp fitting and how those influence the
+    final uncertainties.
+
+    Parameters
+    ----------
+    model: roman_datamodels.datamodels.ImageModel
+        model into which to inject sources
+    cat: astropy.table.Table
+        catalog of sources to inject into image
+    x: list[float] or None
+        x coordinates of catalog locations in image
+    y: list[float] or None
+        y coordinates of catalog locations in image
+    psf: galsim.gsobject.GSObject
+        PSF to use
+    rng: galsim.BaseDeviate
+        galsim random number generator to use
+    gain: float [electron / DN]
+        gain to use when converting simulated electrons to DN
+    webbpsf: bool
+        if True, use WebbPSF to model the PSF
+
+    Returns
+    -------
+    model_out: roman_datamodels.datamodel.ImageModel
+        model with additional sources
+    """
+    if rng is None:
+        rng = galsim.UniformDeviate(123)
+
+    if x is None or y is None:
+        x, y = model.meta.wcs.numerical_inverse(cat['ra'], cat['dec'],
+                                                with_bounding_box=False)
+
+    filter_name = model.meta.instrument.optical_element
+    cat = catalog.table_to_catalog(cat, [filter_name])
+
+    # are we doing photon shooting?
+    chromatic = False
+    if len(cat) > 0 and cat[0].profile.spectral:
+        chromatic = True
+
+    wcs = romanisim.wcs.GWCS(model.meta.wcs)
+    if psf is None:
+        psf = romanisim.psf.make_psf(
+            int(model.meta.instrument.detector[-2:]), filter_name, wcs=wcs,
+            chromatic=False, webbpsf=webbpsf)
+
+    if gain is None:
+        gain = parameters.reference_data['gain']
+
+    # assemble bits we need in order to add a source to an image
+    sourcecounts = galsim.ImageF(model.data.shape[0], model.data.shape[1],
+                                 wcs=wcs, xmin=0, ymin=0)
+    galsim_filter_name = romanisim.bandpass.roman2galsim_bandpass[filter_name]
+    bandpass = roman.getBandpasses(AB_zeropoint=True)[galsim_filter_name]
+    abflux = romanisim.bandpass.get_abflux(filter_name)
+    read_pattern = model.meta.exposure.read_pattern
+    exptime = parameters.read_time * read_pattern[-1][-1]
+    tij = romanisim.l1.read_pattern_to_tij(read_pattern)
+    tbar = ramp.read_pattern_to_tbar(read_pattern)
+    flux_to_counts_factor = exptime
+    if not chromatic:
+        flux_to_counts_factor *= abflux
+
+    # compute the total number of counts we got from the source
+    add_objects_to_image(
+        sourcecounts, cat,
+        xpos=x, ypos=y, psf=psf,
+        flux_to_counts_factor=flux_to_counts_factor,
+        bandpass=bandpass, filter_name=filter_name, rng=rng)
+
+    m = sourcecounts.array != 0
+
+    # add Poisson noise if we made a noiseless, not-photon-shooting
+    # image.
+    if not chromatic:
+        pd = galsim.PoissonDeviate(rng)
+        noise = np.clip(sourcecounts.array[m].copy(), 0, np.inf)
+        pd.generate_from_expectation(noise)
+        sourcecounts.array[m] = noise
+
+    sourcecounts.quantize()
+
+    m = sourcecounts.array != 0
+    # many pixels which may have received a small fraction of a count
+    # receive 0 after the Poisson sampling.
+
+    # create injected source ramp resultants
+    resultants, dq = romanisim.l1.apportion_counts_to_resultants(
+        sourcecounts.array[m], tij, rng=rng)
+    resultants *= u.electron
+
+    # Inject source to original image
+    newramp = model.data[None, :] * tbar[:, None, None] * u.s
+    newramp[:, m] += resultants / gain
+    # newramp has units of DN
+
+    # Make new image of the combination
+    newimage, readvar, poissonvar = make_l2(
+        newramp[:, m], read_pattern,
+        gain=gain, flat=1, darkrate=0)
+
+    res = model.copy()
+    res.data[m] = newimage
+    res.var_rnoise[m] = readvar
+    res.var_poisson[m] = poissonvar
+    res.err[m] = np.sqrt(readvar + poissonvar)
+    return res

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
 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
+from roman_datamodels import datamodels
 import romanisim.bandpass
 
 
@@ -577,7 +578,9 @@ def test_make_test_catalog_and_images():
     if fn is not None:
         fn = str(fn)
     res = image.make_test_catalog_and_images(usecrds=False,
-                                             galaxy_sample_file_name=fn)
+                                             galaxy_sample_file_name=fn,
+                                             webbpsf=False,
+                                             filters=['Y106'])
     assert len(res) > 0
 
 
@@ -633,101 +636,52 @@ def test_inject_source_into_image():
 
     # Set constants and metadata
     galsim.roman.n_pix = 100
-    radius = 0.005
-    flux = 10e-10
+    coord = SkyCoord(ra=270 * u.deg, dec=66 * u.deg)
+    filt = 'F158'
+    meta = util.default_image_meta(coord=coord, filter_name=filt,
+                                   detector='WFI07', ma_table=1)
     rng_seed = 42
     rng = galsim.UniformDeviate(rng_seed)
-    nobj = 10
-    metadata = copy.deepcopy(parameters.default_parameters_dictionary)
-    metadata['instrument']['detector'] = 'WFI07'
-    metadata['instrument']['optical_element'] = 'F158'
-    metadata['exposure']['ma_table_number'] = 1
-    bandpasses = [metadata['instrument']['optical_element']]
-
-    # Establish exposure timing parameters
-    read_pattern = parameters.read_pattern[metadata['exposure']['ma_table_number']]
-    tij = l1.read_pattern_to_tij(read_pattern)
-    tbar = ramp.read_pattern_to_tbar(read_pattern)
-    exptime = parameters.read_time * read_pattern[-1][1]
-
-    # Create catalog of original sources
-    twcs = wcs.get_wcs(metadata, usecrds=False)
-    rd_sca = twcs.toWorld(galsim.PositionD(
-        galsim.roman.n_pix / 2, galsim.roman.n_pix / 2))
-    cat = table.vstack(catalog.make_stars(coord=rd_sca, radius=radius, rng=rng, n=nobj,
-                                          bandpasses=bandpasses, truncation_radius=radius * 0.3))
-
-    # Set source fluxes
-    cat['F158'] = [flux] * len(cat['F158'])
-
-    # Create original image with sources
-    rng = galsim.UniformDeviate(None)
-    origimage, simcatobj = image.simulate(
-        metadata, cat, usecrds=False,
-        webbpsf=False, level=2,
-        rng=rng)
-
-    # Create source to inject
+    cat = catalog.make_dummy_table_catalog(coord, radius=0.1, bandpasses=[filt],
+                                           nobj=2000, rng=rng)
+    # Create starting image
+    im, simcatobj = image.simulate(
+        meta, cat, usecrds=False, webbpsf=True, level=2,
+        rng=rng, psf_keywords=dict(nlambda=1),
+        crparam=None)
+    im = datamodels.ImageModel(im)
 
     # Create catalog with one source for injection
     xpos, ypos = 10, 10
-    source_cat = cat.copy()
-    source_cat.remove_rows(slice(0, nobj - 1))
-    source_cat['ra'], source_cat['dec'] = (twcs._radec(xpos, ypos) * u.rad).to(u.deg).value
-
-    # Create empty galsim image
-    sourcecounts = galsim.ImageF(roman.n_pix, roman.n_pix, wcs=twcs, xmin=0, ymin=0)
-
-    # Set parameters for injection source simulation
-    filter_name = metadata['instrument']['optical_element']
-    galsim_filter_name = romanisim.bandpass.roman2galsim_bandpass[filter_name]
-    bandpass = roman.getBandpasses(AB_zeropoint=True)[galsim_filter_name]
-    abflux = romanisim.bandpass.get_abflux(filter_name)
-    sca = int(metadata['instrument']['detector'][3:])
-    flux_to_counts_factor = exptime * abflux
-
-    # Create PSF
-    psf = romanisim.psf.make_psf(sca, filter_name, wcs=twcs,
-                                 chromatic=False, webbpsf=True)
-
-    # Create injected source image
-    source_cat = catalog.table_to_catalog(source_cat, [filter_name])
-    image.add_objects_to_image(
-        sourcecounts, source_cat,
-        xpos=[xpos], ypos=[ypos], psf=psf, flux_to_counts_factor=flux_to_counts_factor,
-        bandpass=bandpass, filter_name=filter_name, rng=rng)
-    sourcecounts.quantize()
-
-    # Create injected source ramp resultants
-    resultants, dq = l1.apportion_counts_to_resultants(sourcecounts.array, tij, rng=rng)
-
-    # Inject source to original image
-    newramp = (origimage.data[np.newaxis, :] * tbar[:, np.newaxis, np.newaxis]).value + resultants
+    catinj = cat[:1]
+    flux = 1e-7
+    catinj[filt] = flux
+    catinj['type'] = 'PSF'
 
-    # Make new image of the combination
-    newimage, readvar, poissonvar = image.make_l2(
-        newramp * u.DN, read_pattern,
-        gain=1 * u.electron / u.DN, flat=1, darkrate=0)
+    iminj = image.inject_sources_into_l2(im, catinj, x=[xpos], y=[ypos])
 
-    # Create mask of PSF
-    nonzero = (sourcecounts.array != 0)
+    # Test that all pixels near the PSF are different from the original values
+    assert np.all((im.data.value[ypos - 1:ypos + 2, xpos - 1:xpos + 2] !=
+                   iminj.data.value[ypos - 1:ypos + 2, xpos - 1: xpos + 2]))
 
-    # Test that all pixels inside of the psf of the injected source are different from the original image
-    assert np.any((origimage.data.value != newimage.value), where=nonzero)
+    # Test that pixels far from the injected source are close to the original image
+    assert np.all(im.data.value[-10:, -10:] == iminj.data.value[-10:, -10:])
 
-    # Test that all pixels outside of the psf of the injected source are close to the original image
-    assert np.allclose(origimage.data.value[~nonzero], newimage.value[~nonzero], rtol=1e-05, atol=1e-08)
+    # Test that the amount of added flux makes sense
+    fluxeps = flux * romanisim.bandpass.get_abflux('F158') * u.electron / u.s
+    assert np.abs(np.sum(iminj.data - im.data) * parameters.reference_data['gain'] /
+                  fluxeps - 1) < 0.1
 
     # Create log entry and artifacts
     log.info(f'DMS231: successfully injected a source into an image at x,y = {xpos},{ypos}.')
 
     artifactdir = os.environ.get('TEST_ARTIFACT_DIR', None)
     if artifactdir is not None:
         af = asdf.AsdfFile()
-        af.tree = {'originalimage': origimage,
-                   'newimage': newimage,
-                   'flux': flux,
-                   'tij': tij}
+        af.tree = {'originalimage': im,
+                   'newimage': iminj,
+                   'catinj': catinj,
+                   'cat': cat}
         af.write_to(os.path.join(artifactdir, 'dms231.asdf'))