From 064ce028e4f2e21b9081503c3d5fad8ca469f5ac Mon Sep 17 00:00:00 2001 From: Eddie Schlafly Date: Tue, 13 Aug 2024 05:57:37 -0400 Subject: [PATCH] Refactor l2 injection (#138) * Initial commit. * Make test_catalog_and_images test faster. * Ruff formatting changes. --- romanisim/image.py | 134 ++++++++++++++++++++++++++++++++++ romanisim/tests/test_image.py | 116 +++++++++-------------------- 2 files changed, 169 insertions(+), 81 deletions(-) diff --git a/romanisim/image.py b/romanisim/image.py index ba32a3ae..8b726b38 100644 --- a/romanisim/image.py +++ b/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 diff --git a/romanisim/tests/test_image.py b/romanisim/tests/test_image.py index 4a76b833..38f1fc82 100644 --- a/romanisim/tests/test_image.py +++ b/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,90 +636,41 @@ 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}.') @@ -724,10 +678,10 @@ def test_inject_source_into_image(): 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'))