diff --git a/romanisim/image.py b/romanisim/image.py index b2aa2818..b9ded592 100644 --- a/romanisim/image.py +++ b/romanisim/image.py @@ -206,7 +206,7 @@ def trim_objlist(objlist, image): def add_objects_to_image(image, objlist, xpos, ypos, psf, flux_to_counts_factor, convtimes=None, - bandpass=None, filter_name=None, + bandpass=None, filter_name=None, add_noise=False, rng=None, seed=None): """Add sources to an image. @@ -287,6 +287,8 @@ def add_objects_to_image(image, objlist, xpos, ypos, psf, stamp = final.drawImage( bandpass, center=image_pos, wcs=pwcs, method='phot', rng=rng) + if add_noise: + stamp.addNoise(galsim.PoissonNoise(rng)) else: try: stamp = final.drawImage(center=image_pos, @@ -314,7 +316,8 @@ def simulate_counts_generic(image, exptime, objlist=None, psf=None, sky=None, dark=None, flat=None, xpos=None, ypos=None, ignore_distant_sources=10, bandpass=None, - filter_name=None, rng=None, seed=None): + filter_name=None, rng=None, seed=None, + **kwargs): """Add some simulated counts to an image. No Roman specific code allowed! To do this, we need to have an image diff --git a/romanisim/l3.py b/romanisim/l3.py index b7bcfbac..c7708d53 100644 --- a/romanisim/l3.py +++ b/romanisim/l3.py @@ -4,22 +4,30 @@ for most of the real work. """ +import math import numpy as np import galsim -from . import parameters +from . import log +import romanisim.catalog import romanisim.wcs import romanisim.l1 import romanisim.bandpass import romanisim.psf import romanisim.image import romanisim.persistence +import romanisim.parameters import roman_datamodels.datamodels as rdm from roman_datamodels.stnode import WfiMosaic +import astropy.units as u +import roman_datamodels.maker_utils as maker_utils +import astropy +from galsim import roman +from astropy import table 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): + filter_name=None, bandpass=None, coords_unit='rad', wcs=None, psf=None, rng=None, seed=None): """Add objects to a Level 3 mosaic Parameters @@ -38,6 +46,12 @@ def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords Factor to convert data to cps unit_factor: float Factor to convert counts data to MJy / sr + filter_name : str + Filter to use to select appropriate flux from objlist. This is only + used when achromatic PSFs and sources are being rendered. + bandpass : galsim.Bandpass + Bandpass in which mosaic is being rendered. This is used only in cases + where chromatic profiles & PSFs are being used. coords_unit : string units of coords wcs : galsim.GSFitsWCS @@ -57,6 +71,8 @@ def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords # Obtain optical element if filter_name is None: filter_name = l3_mos.meta.basic.optical_element + if bandpass is None: + bandpass = [filter_name] # Generate WCS (if needed) if wcs is None: @@ -64,7 +80,7 @@ def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords # 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) + psf = romanisim.psf.make_psf(filter_name=filter_name, sca=romanisim.parameters.default_sca, chromatic=False, webbpsf=True) # Create Image canvas to add objects to if isinstance(l3_mos, (rdm.MosaicModel, WfiMosaic)): @@ -83,10 +99,10 @@ def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords 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, + outinfo = 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, + bandpass=bandpass, filter_name=filter_name, rng=rng, seed=seed) # Save array with added sources @@ -94,4 +110,491 @@ def add_objects_to_l3(l3_mos, source_cat, exptimes, xpos=None, ypos=None, coords l3_mos.data = sourcecountsall.array * l3_mos.data.unit # l3_mos is updated in place, so no return - return None + return outinfo + + +def generate_mosaic_geometry(): + """Generate a geometry map (context) for a mosaic dither pattern / set of pointings and roll angles + TBD + """ + pass + + +def generate_exptime_array(cat, meta): + """ To be updated. + Function to ascertain / set exposure time in each pixel + Present code is placeholder to be built upon + """ + + # Get wcs for this metadata + twcs = romanisim.wcs.get_mosaic_wcs(meta) + + # Obtain sky positions for objects + coords = np.array([[o.sky_pos.ra.rad, o.sky_pos.dec.rad] + for o in cat]) + + # Calculate x,y positions for objects + allx, ally = twcs.radecToxy(coords[:, 0], coords[:, 1], 'rad') + + # Obtain the sample extremums + xmin = min(allx) + xmax = max(allx) + ymin = min(ally) + ymax = max(ally) + + # Obtain WCS center + xcen, ycen = twcs.radecToxy(twcs.center.ra, twcs.center.dec, 'rad') + + # Determine maximum extremums from WCS center + xdiff = max([math.ceil(xmax - xcen), math.ceil(xcen - xmin)]) + ydiff = max([math.ceil(ymax - ycen), math.ceil(ycen - ymin)]) + + # Create context map preserving WCS center + context = np.ones((1, 2 * ydiff, 2 * xdiff), dtype=np.uint32) + + return context + + +def simulate(shape, wcs, efftimes, filter, catalog, metadata={}, + effreadnoise=None, sky=None, psf=None, coords_unit='rad', + cps_conv=None, unit_factor=None, + bandpass=None, seed=None, rng=None, + **kwargs): + """Simulate a sequence of observations on a field in different bandpasses. + + Parameters + ---------- + shape : tuple + Array shape of mosaic to simulate. + wcs : galsim.GSFitsWCS + WCS corresponding to image + efftimes : np.ndarray or float + Effective exposure time of reach pixel in mosaic. + If an array, shape must match shape parameter. + filter : str + Filter to use to select appropriate flux from objlist. This is only + used when achromatic PSFs and sources are being rendered. + catalog : list[CatalogObject] or Table + List of catalog objects to add to l3_mos + metadata : dict + Metadata structure for Roman asdf file. + effreadnoise : float + Effective read noise for mosaic. + sky : float or array_like + Image or constant with the counts / pix / sec from sky. If None, then + sky will be generated from galsim's getSkyLevel for Roman for the + date provided in metadata[basic][time_mean_mjd]. + psf : galsim.Profile + PSF for image + coords_unit : string + units of coords + cps_conv : float + Factor to convert data to cps + unit_factor: float + Factor to convert counts data to MJy / sr + bandpass : galsim.Bandpass + Bandpass in which mosaic is being rendered. This is used only in cases + where chromatic profiles & PSFs are being used. + rng : galsim.BaseDeviate + random number generator to use + seed : int + seed to use for random number generator + + Returns + ------- + mosaic_mdl : roman_datamodels model + simulated mosaic + extras : dict + Dictionary of additionally valuable quantities. Includes at least + simcatobj, the image positions and fluxes of simulated objects. It may + also include information on persistence and cosmic ray hits. + """ + + # Create metadata object + meta = maker_utils.mk_mosaic_meta() + for key in romanisim.parameters.default_mosaic_parameters_dictionary.keys(): + meta[key].update(romanisim.parameters.default_mosaic_parameters_dictionary[key]) + meta['wcs'] = wcs + meta['basic']['optical_element'] = filter + for key in metadata.keys(): + if key not in meta: + meta[key] = metadata[key] + else: + meta[key].update(metadata[key]) + # May need an equivalent of this this for L3? + # util.add_more_metadata(meta) + + log.info('Simulating filter {0}...'.format(filter)) + + # Get filter and bandpass + galsim_filter_name = romanisim.bandpass.roman2galsim_bandpass[filter] + if bandpass is None: + bandpass = roman.getBandpasses(AB_zeropoint=True)[galsim_filter_name] + + # Create initial galsim image (x & y are flipped) + image = galsim.ImageF(shape[1], shape[0], wcs=wcs, xmin=0, ymin=0) + + # Create sky for this mosaic, if not provided (in cps) + if sky is None: + date = meta['basic']['time_mean_mjd'] + if not isinstance(date, astropy.time.Time): + date = astropy.time.Time(date, format='mjd') + + # Examine this in more detail to ensure it is correct + # Create base sky level + mos_cent_pos = wcs.toWorld(image.true_center) + sky_level = roman.getSkyLevel(bandpass, world_pos=mos_cent_pos, exptime=1) + sky_level *= (1.0 + roman.stray_light_fraction) + sky_mosaic = image * 0 + wcs.makeSkyImage(sky_mosaic, sky_level) + sky_mosaic += roman.thermal_backgrounds[galsim_filter_name] + sky = sky_mosaic + + # Obtain unit conversion factors + # Need to convert from counts / pixel to MJy / sr + # Flux to counts + if cps_conv is None: + cps_conv = romanisim.bandpass.get_abflux(filter) + # Unit factor + if unit_factor is None: + unit_factor = ((3631 * u.Jy) / (romanisim.bandpass.get_abflux(filter) * 10e6 + * romanisim.parameters.reference_data['photom']["pixelareasr"][filter])).to(u.MJy / u.sr) + + # Set effective read noise + if effreadnoise is None: + effreadnoise = efftimes / romanisim.parameters.read_time + + # Simulate mosaic cps + mosaic, simcatobj = simulate_cps( + image, meta, efftimes, objlist=catalog, psf=psf, + sky=sky, + effreadnoise=effreadnoise, cps_conv=cps_conv, coords_unit=coords_unit, + bandpass=bandpass, + wcs=wcs, rng=rng, seed=seed, unit_factor=unit_factor, + **kwargs) + + # Set poisson variance + if "var_poisson" in simcatobj: + var_poisson = simcatobj["var_poisson"] + else: + var_poisson = None + + # Set read noise variance + if "var_readnoise" in simcatobj: + var_readnoise = simcatobj["var_readnoise"].value + else: + var_readnoise = None + + # Create Mosaic Model + mosaic_mdl = make_l3(mosaic, metadata, efftimes, unit_factor=unit_factor, + var_poisson=var_poisson, var_readnoise=var_readnoise) + + # Add simulation artifacts + extras = {} + extras['simcatobj'] = simcatobj + extras['wcs'] = mosaic.wcs + + log.info('Simulation complete.') + # return mosaic, extras + return mosaic_mdl, extras + + +def simulate_cps(image, metadata, efftimes, objlist=None, psf=None, + wcs=None, xpos=None, ypos=None, coord=None, sky=0, + effreadnoise=None, cps_conv=1, unit_factor=1.0 * (u.MJy / u.sr), + bandpass=None, coords_unit='rad', + rng=None, seed=None, ignore_distant_sources=10,): + """Simulate average MegaJankies per steradian in a single SCA. + + Parameters + ---------- + image : galsim.Image + Image onto which other effects should be added, with associated WCS. + metadata : dict + Metadata structure for Roman asdf file. + efftimes : np.ndarray or float + Effective exposure time of reach pixel in mosaic. + If an array, shape must match shape parameter. + objlist : list[CatalogObject], Table, or None + Sources to render + psf : galsim.Profile + PSF to use when rendering sources + wcs : galsim.GSFitsWCS + WCS corresponding to image + xpos, ypos : array_like (float) + x, y positions of each source in objlist + coord : array_like (float) + ra, dec positions of each source in objlist + sky : float or array_like + Image or constant with the counts / pix / sec from sky. + effreadnoise : float + Effective read noise for mosaic. + cps_conv : float + Factor to convert data to cps + unit_factor: float + Factor to convert counts data to MJy / sr + bandpass : galsim.Bandpass + bandpass to use for rendering chromatic objects + coords_unit : string + units of coords + rng : galsim.BaseDeviate + random number generator + seed : int + seed for random number generator + ignore_distant_sources : int + Ignore sources more than this distance off image. + + Returns + ------- + objinfo : np.ndarray + Information on position and flux of each rendered source. + + Returns + ------- + image : galsim.Image + Idealized image of scene as seen by Roman + extras : dict + catalog of simulated objects in image, noise, and misc. debug + """ + + if rng is None and seed is None: + seed = 144 + log.warning( + 'No RNG set, constructing a new default RNG from default seed.') + if rng is None: + rng = galsim.UniformDeviate(seed) + + # Dictionary to hold simulation artifacts + extras = {} + + # Objlist wcs check + if (objlist is not None and len(objlist) > 0 + and image.wcs is None and (xpos is None or ypos is None)): + raise ValueError('xpos and ypos must be set if rendering objects ' + 'without a WCS.') + if objlist is None: + objlist = [] + if len(objlist) > 0 and xpos is None: + if isinstance(objlist, table.Table): + coord = np.array([[o['ra'], np.radians(o['dec'])] for o in objlist]) + else: + if (coords_unit == 'rad'): + coord = np.array([[o.sky_pos.ra.rad, o.sky_pos.dec.rad] + for o in objlist]) + else: + coord = np.array([[o.sky_pos.ra.deg, o.sky_pos.dec.deg] + for o in objlist]) + xpos, ypos = image.wcs.radecToxy(coord[:, 0], coord[:, 1], 'rad') + # use private vectorized transformation + if xpos is not None: + xpos = np.array(xpos) + if ypos is not None: + ypos = np.array(ypos) + + chromatic = False + if len(objlist) > 0 and objlist[0].profile.spectral: + chromatic = True + + # Check for objects outside the image boundary (+ consideration) + if len(objlist) > 0: + keep = romanisim.image.in_bounds(xpos, ypos, image.bounds, ignore_distant_sources) + + objlist = np.array(objlist)[keep].tolist() + xpos = xpos[keep] + ypos = ypos[keep] + # Pixelized object locations + xpos_idx = [round(x) for x in xpos] + ypos_idx = [round(y) for y in ypos] + + offedge = romanisim.image.in_bounds(xpos, ypos, image.bounds, 0) + # Set exposure time per source + if isinstance(efftimes, np.ndarray): + src_exptimes = [efftimes[y, x] for x, y in zip(xpos_idx, ypos_idx)] + else: + src_exptimes = [efftimes] * len(xpos) + src_exptimes = np.array(src_exptimes) + + # Set the average exposure time to objects lacking one + if True in offedge: + avg_exptime = np.average(src_exptimes) + src_exptimes[offedge] = avg_exptime + + else: + src_exptimes = [] + + # Generate GWCS compatible wcs + if wcs is None: + wcs = romanisim.wcs.get_mosaic_wcs(metadata, shape=image.array.shape, xpos=xpos, ypos=ypos, coord=coord) + image.wcs = wcs + + # Add objects to mosaic + if len(src_exptimes) > 0: + objinfokeep = add_objects_to_l3( + image, objlist, src_exptimes, wcs=wcs, xpos=xpos, ypos=ypos, + filter_name=metadata['basic']['optical_element'], bandpass=bandpass, psf=psf, rng=rng, + cps_conv=cps_conv, unit_factor=unit_factor) + + # Add object info artifacts + objinfo = {} + objinfo['array'] = np.zeros( + len(objlist), + dtype=[('x', 'f4'), ('y', 'f4'), ('counts', 'f4'), ('time', 'f4')]) + if len(objlist) > 0: + objinfo['x'] = xpos + objinfo['y'] = ypos + objinfo['counts'] = objinfokeep['counts'] + objinfo['time'] = objinfokeep['time'] + else: + objinfo['counts'] = np.array([]) + objinfo['time'] = np.array([]) + extras['objinfo'] = objinfo + + # Noise + im_no_noise = image.array.copy() + + # add Poisson noise if we made a noiseless, not-photon-shooting + # image. + poisson_noise = galsim.PoissonNoise(rng) + if not chromatic: + # This works in ADU, so need to convert back to counts first.. add this, then convert back + image *= efftimes / unit_factor.value + image.addNoise(poisson_noise) + image /= efftimes / unit_factor.value + + if sky is not None: + if isinstance(sky, (galsim.Image)): + sky_arr = sky.array + elif not isinstance(sky, (np.ndarray)): + sky_arr = sky * np.ones(image.array.shape, dtype=float) + else: + sky_arr = sky + + workim = image * 0 + workim += sky + workim *= efftimes + workim.addNoise(poisson_noise) + + image += (workim * unit_factor.value / efftimes) + + else: + sky_arr = np.zeros(image.array.shape, dtype=float) + + extras['var_poisson'] = np.abs(image.array - (im_no_noise + sky_arr)) + + # Add readnoise + if effreadnoise is not None: + # This works in ADU, so need to convert back to counts first.. add this, then convert back + readnoise = np.zeros(image.array.shape, dtype='f4') + rn_rng = galsim.GaussianDeviate(seed) + rn_rng.generate(readnoise) + readnoise = readnoise * effreadnoise + readnoise /= np.sqrt(efftimes) + + image += readnoise * unit_factor.value / efftimes + extras['readnoise'] = readnoise * unit_factor / efftimes + + extras['var_readnoise'] = (rn_rng.sigma * (effreadnoise / np.sqrt(efftimes)) * (unit_factor / efftimes))**2 + + # Return image and artifacts + return image, extras + + +def make_l3(image, metadata, efftimes, rng=None, seed=None, var_poisson=None, + var_flat=None, var_readnoise=None, context=None, + unit_factor=(1.0 * u.MJy / u.sr)): + """ + Create and populate MosaicModel of image and noises. + + Parameters + ---------- + image : galsim.Image + Image containing mosaic data. + metadata : dict + Metadata structure for Roman asdf file. + efftimes : np.ndarray or float + Effective exposure time of reach pixel in mosaic. + If an array, shape must match shape parameter. + rng : galsim.BaseDeviate + Random number generator. + seed : int + Seed for random number generator. + var_poisson : np.ndarray + Poisson variance for each pixel + var_flat : np.ndarray + Flat variance for each pixel + var_readnoise : np.ndarray + Read Noise variance for each pixel + context : np.ndarray + File number(s) for each pixel + unit_factor: float + Factor to convert counts data to MJy / sr + + Returns + ------- + objinfo : np.ndarray + Information on position and flux of each rendered source. + + Returns + ------- + image : rdm.MosaicModel + Mosaic model object containing the data, metadata, variances, weight, and context. + """ + + # Create mosaic data object + mosaic = image.array.copy() + + # Set rng for creating readnoise, flat noise + if rng is None and seed is None: + seed = 46 + log.warning( + 'No RNG set, constructing a new default RNG from default seed.') + if rng is None: + rng = galsim.GaussianDeviate(seed) + + # Ensure that effective times are an array + if isinstance(efftimes, np.ndarray): + efftimes_arr = efftimes + else: + efftimes_arr = efftimes * np.ones(mosaic.shape, dtype=np.float32) + + # Set mosaic to be a mosaic node + mosaic_node = maker_utils.mk_level3_mosaic(shape=mosaic.shape, meta=metadata) + + # Set data + mosaic_node.data = mosaic * unit_factor.unit + + # Poisson noise variance + if var_poisson is None: + mosaic_node.var_poisson = (unit_factor**2 / efftimes_arr**2).astype(np.float32) + else: + mosaic_node.var_poisson = u.Quantity(var_poisson.astype(np.float32), unit_factor.unit ** 2) + + # Read noise variance + if var_readnoise is None: + mosaic_node.var_rnoise = u.Quantity(np.zeros(mosaic.shape, dtype=np.float32), unit_factor.unit ** 2) + else: + mosaic_node.var_rnoise = u.Quantity((var_readnoise * np.ones(mosaic.shape)).astype(np.float32), unit_factor.unit ** 2) + + # Flat variance + if var_flat is None: + mosaic_node.var_flat = u.Quantity(np.zeros(mosaic.shape, dtype=np.float32), unit_factor.unit ** 2) + else: + mosaic_node.var_flat = var_flat + + # Total error + mosaic_node.err = np.sqrt(mosaic_node.var_rnoise + mosaic_node.var_flat + mosaic_node.var_poisson).astype(np.float32) + + # Weight + # Use exptime weight + mosaic_node.weight = efftimes_arr.astype(np.float32) + + # Context + # Binary index of images that contribute to the pixel + # Defined by geometry.. if not, all non zero = 1. + if context is None: + mosaic_node.context = np.ones((1,) + mosaic.shape, dtype=np.uint32) + else: + mosaic_node.context = context + + # Return mosaic + return mosaic_node diff --git a/romanisim/parameters.py b/romanisim/parameters.py index bb993970..6619ed3b 100644 --- a/romanisim/parameters.py +++ b/romanisim/parameters.py @@ -5,6 +5,8 @@ from astropy.time import Time from astropy import units as u +pixel_scale = 0.11 + read_pattern = {1: [[1 + x for x in range(8)], [9 + x for x in range(8)], [17 + x for x in range(8)], diff --git a/romanisim/ramp.py b/romanisim/ramp.py index caf94ed3..6ed10bb3 100644 --- a/romanisim/ramp.py +++ b/romanisim/ramp.py @@ -248,7 +248,7 @@ class RampFitInterpolator: resultants. The weights of this linear combination are a single parameter family in the flux in the ramp divided by the read variance. So rather than explicitly calculating those weights for each pixel, we can up front calculate - them overa grid in the flux over the read variance, and interpolate off that + them over a grid in the flux over the read variance, and interpolate off that grid for each point. That can all be done in a vectorized way, allowing one to avoid doing something like a matrix inverse for each of a 16 million pixels. diff --git a/romanisim/tests/test_image.py b/romanisim/tests/test_image.py index 38f1fc82..717e23c6 100644 --- a/romanisim/tests/test_image.py +++ b/romanisim/tests/test_image.py @@ -394,18 +394,23 @@ def test_simulate_counts_generic(): # to figure out where to render objects. That would require setting # up a real PSF and is annoying. Skipping that. # render some objects - image.simulate_counts_generic( - im, exptime, objlist=imdict['graycatalog'], psf=imdict['impsfgray'], + im6 = im.copy() + objinfo = image.simulate_counts_generic( + im6, exptime, objlist=imdict['graycatalog'], psf=imdict['impsfgray'], xpos=[50, 50], ypos=[50, 50], zpflux=zpflux, filter_name=imdict['filter_name']) + assert np.sum(im6.array) > 0 # at least verify that we added some sources... + assert len(objinfo) == 2 # two sources were added + im7 = im.copy() objinfo = image.simulate_counts_generic( - im, exptime, objlist=imdict['chromcatalog'], + im7, exptime, objlist=imdict['chromcatalog'], xpos=[50, 50], ypos=[50, 50], psf=imdict['impsfchromatic'], bandpass=imdict['bandpass']) - assert np.sum(im.array) > 0 # at least verify that we added some sources... + assert np.sum(im7.array) > 0 # at least verify that we added some sources... assert len(objinfo) == 2 # two sources were added + im8 = im.copy() objinfo = image.simulate_counts_generic( - im, exptime, objlist=imdict['chromcatalog'], + im8, exptime, objlist=imdict['chromcatalog'], psf=imdict['impsfchromatic'], xpos=[1000, 1000], ypos=[1000, 1000], zpflux=zpflux) assert np.sum(objinfo['counts'] > 0) == 0 diff --git a/romanisim/tests/test_l3.py b/romanisim/tests/test_l3.py index 71225edb..a1238330 100644 --- a/romanisim/tests/test_l3.py +++ b/romanisim/tests/test_l3.py @@ -3,19 +3,20 @@ """ import os -import copy +from copy import deepcopy import math import numpy as np import galsim -from romanisim import parameters, catalog, wcs, l3 +from romanisim import parameters, catalog, wcs, l3, psf, util, log from astropy import units as u from astropy import table import asdf import pytest from metrics_logger.decorators import metrics_logger -from romanisim import log import roman_datamodels.maker_utils as maker_utils import romanisim.bandpass +from galsim import roman +from astropy.coordinates import SkyCoord @metrics_logger("DMS232") @@ -27,7 +28,7 @@ def test_inject_sources_into_mosaic(): # Set constants and metadata galsim.roman.n_pix = 200 rng_seed = 42 - metadata = copy.deepcopy(parameters.default_mosaic_parameters_dictionary) + metadata = deepcopy(parameters.default_mosaic_parameters_dictionary) filter_name = 'F158' metadata['basic']['optical_element'] = filter_name @@ -142,3 +143,364 @@ def test_inject_sources_into_mosaic(): 'source_cat_table': sc_table, } af.write_to(os.path.join(artifactdir, 'dms232.asdf')) + + +@metrics_logger("DMS219") +@pytest.mark.soctests +def test_sim_mosaic(): + """Generating mosaic from catalog file. + """ + # Define random seed + rng_seed = 42 + + # Obtain pointing + ra_ref = parameters.default_mosaic_parameters_dictionary['wcsinfo']['ra_ref'] + dec_ref = parameters.default_mosaic_parameters_dictionary['wcsinfo']['dec_ref'] + + # Set metadata and capture filter + metadata = deepcopy(parameters.default_mosaic_parameters_dictionary) + filter_name = metadata['basic']['optical_element'] + + # Set exposure time + exptimes = [600] + + # Create catalog of objects + cen = SkyCoord(ra=(ra_ref * u.deg).to(u.rad), dec=(dec_ref * u.deg).to(u.rad)) + cat = catalog.make_dummy_table_catalog(cen, radius=0.02, nobj=100, seed=rng_seed) + # Make the first 10 bright for tests + cat[filter_name][0:10] *= 1e4 + source_cat = catalog.table_to_catalog(cat, [filter_name]) + + # Create bounds from the object list + twcs = romanisim.wcs.get_mosaic_wcs(metadata) + coords = np.array([[o.sky_pos.ra.rad, o.sky_pos.dec.rad] + for o in source_cat]) + allx, ally = twcs.radecToxy(coords[:, 0], coords[:, 1], 'rad') + + # Obtain the sample extremums + xmin = min(allx) + xmax = max(allx) + ymin = min(ally) + ymax = max(ally) + + # Obtain WCS center + xcen, ycen = twcs.radecToxy(twcs.center.ra, twcs.center.dec, 'rad') + + # Determine maximum extremums from WCS center + xdiff = max([math.ceil(xmax - xcen), math.ceil(xcen - xmin)]) + 1 + ydiff = max([math.ceil(ymax - ycen), math.ceil(ycen - ymin)]) + 1 + + # Create context map preserving WCS center + context = np.ones((1, 2 * ydiff, 2 * xdiff), dtype=np.uint32) + + # Generate properly sized WCS + moswcs = romanisim.wcs.get_mosaic_wcs(metadata, shape=(context.shape[1:])) + + # Simulate mosaic + mosaic, extras = l3.simulate(context.shape[1:], moswcs, exptimes[0], + filter_name, source_cat, metadata=metadata, seed=rng_seed) + + # Ensure center pixel of bright objects is bright + x_all, y_all = moswcs.radecToxy(coords[:10, 0], coords[:10, 1], 'rad') + for x, y in zip(x_all, y_all): + x = int(x) + y = int(y) + assert mosaic.data.value[y, x] > (np.median(mosaic.data.value) * 5) + + # Did we get all the flux? + # Convert to CPS for comparison + # Unit factor + unit_factor = ((3631 * u.Jy) / (romanisim.bandpass.get_abflux(filter_name) * 10e6 + * parameters.reference_data['photom']["pixelareasr"][filter_name])).to(u.MJy / u.sr) + totflux = np.sum(mosaic.data.value - np.median(mosaic.data.value)) / unit_factor.value + + # Flux to counts + cps_conv = romanisim.bandpass.get_abflux(filter_name) + expectedflux = np.sum(cat[filter_name]) * cps_conv + + # Ensure that the measured flux is close to the expected flux + assert np.abs(np.log(expectedflux) / np.log(totflux) - 1) < 0.1 + + # Add log entries and artifacts + log.info('DMS219 successfully created mosaic file with sources rendered ' + 'at correct locations with matching flux and added noise.') + + artifactdir = os.environ.get('TEST_ARTIFACT_DIR', None) + if artifactdir is not None: + af = asdf.AsdfFile() + af.tree = {'l3_mosaic': mosaic, + 'source_cat_table': cat, + } + af.write_to(os.path.join(artifactdir, 'dms219.asdf')) + + +def set_up_image_rendering_things(): + """ Function to set up objects for use in tests + """ + # Create sample image, filter, etc. + im = galsim.ImageF(100, 100, scale=0.11, xmin=0, ymin=0) + filter_name = 'F158' + impsfgray = psf.make_psf(1, filter_name, webbpsf=True, chromatic=False, + nlambda=1) # nlambda = 1 speeds tests + impsfchromatic = psf.make_psf(1, filter_name, webbpsf=False, + chromatic=True) + bandpass = roman.getBandpasses(AB_zeropoint=True)['H158'] + counts = 1000 + fluxdict = {filter_name: counts} + + # Sample catalogs + graycatalog = [ + catalog.CatalogObject(None, galsim.DeltaFunction(), deepcopy(fluxdict)), + catalog.CatalogObject(None, galsim.Sersic(1, half_light_radius=0.2), + deepcopy(fluxdict)) + ] + vega_sed = galsim.SED('vega.txt', 'nm', 'flambda') + vega_sed = vega_sed.withFlux(counts, bandpass) + chromcatalog = [ + catalog.CatalogObject(None, galsim.DeltaFunction() * vega_sed, None), + catalog.CatalogObject( + None, galsim.Sersic(1, half_light_radius=0.2) * vega_sed, None) + ] + tabcat = table.Table() + tabcat['ra'] = [270.0] + tabcat['dec'] = [66.0] + tabcat[filter_name] = counts + tabcat['type'] = 'PSF' + tabcat['n'] = -1 + tabcat['half_light_radius'] = -1 + tabcat['pa'] = -1 + tabcat['ba'] = -1 + + # Return dictionary with the above values + return dict(im=im, impsfgray=impsfgray, + impsfchromatic=impsfchromatic, + bandpass=bandpass, counts=counts, fluxdict=fluxdict, + graycatalog=graycatalog, + chromcatalog=chromcatalog, filter_name=filter_name, + tabcatalog=tabcat) + + +def test_simulate_vs_cps(): + """ Tests to ensure that simulate runs match simulate_cps output + """ + # Set random seed + rng_seed = 42 + + # Set image, catalog, etc. + imdict = set_up_image_rendering_things() + chromcat = imdict['chromcatalog'] + graycat = imdict['graycatalog'] + coord = SkyCoord(270 * u.deg, 66 * u.deg) + for o in chromcat: + o.sky_pos = coord + for o in graycat: + o.sky_pos = coord + # these are all dumb coordinates; the coord sent to simulate_counts + # is the coordinate of the boresight, but that doesn't need to be on SCA 1. + # But at least they'll exercise some machinery if the ignore_distant_sources + # argument is high enough! + roman.n_pix = 100 + exptime = 600 + + # Create metadata + meta = util.default_image_meta(filter_name='F158') + wcs.fill_in_parameters(meta, coord) + metadata = deepcopy(parameters.default_mosaic_parameters_dictionary) + filter_name = 'F158' + metadata['basic']['optical_element'] = filter_name + metadata['wcsinfo']['ra_ref'] = 270 + metadata['wcsinfo']['dec_ref'] = 66 + + # Set up blank image + im = imdict['im'].copy() + im.array[:] = 0 + + # Set WCS + twcs = romanisim.wcs.get_mosaic_wcs(metadata, shape=(roman.n_pix, roman.n_pix)) + + # Create chromatic data in simulate_cps + im1 = im.copy() + im1, extras1 = l3.simulate_cps(im1, metadata, exptime, objlist=chromcat, + xpos=[50] * len(chromcat), ypos=[50] * len(chromcat), + bandpass=imdict['bandpass'], seed=rng_seed, + ignore_distant_sources=100) + + # Create filter data in simulate_cps + im2 = im.copy() + im2, extras2 = l3.simulate_cps(im2, metadata, exptime, objlist=graycat, + xpos=[50] * len(chromcat), ypos=[50] * len(chromcat), + psf=imdict['impsfgray'], seed=rng_seed, + ignore_distant_sources=100) + + # Create chromatic data in simulate + im3, extras3 = l3.simulate((roman.n_pix, roman.n_pix), twcs, exptime, filter_name, chromcat, + bandpass=imdict['bandpass'], seed=rng_seed, + cps_conv=1, unit_factor=(1 * u.MJy / u.sr), + metadata=metadata, sky=0, + ignore_distant_sources=100, effreadnoise=0, + ) + + # Create filter data in simulate + im4, extras4 = l3.simulate((roman.n_pix, roman.n_pix), twcs, exptime, filter_name, graycat, + psf=imdict['impsfgray'], + cps_conv=1, unit_factor=(1 * u.MJy / u.sr), seed=rng_seed, + metadata=metadata, sky=0, + ignore_distant_sources=100, + effreadnoise=0, + ) + + # Ensure that the simulate and simulate_cps output matches for each type + assert np.allclose(im1.array, im3['data'].value) + assert np.allclose(im2.array, im4['data'].value) + + +def test_simulate_cps(): + """ Test various simulation options + """ + # Set random seed + rng_seed = 42 + + # Set image, catalog, etc. + imdict = set_up_image_rendering_things() + im = imdict['im'].copy() + im.array[:] = 0 + npix = np.prod(im.array.shape) + exptime = 100 + + # Create metadata + metadata = deepcopy(parameters.default_mosaic_parameters_dictionary) + filter_name = 'F158' + metadata['basic']['optical_element'] = filter_name + metadata['wcsinfo']['ra_ref'] = 270 + metadata['wcsinfo']['dec_ref'] = 66 + coord = SkyCoord(270 * u.deg, 66 * u.deg) + wcs.fill_in_parameters(metadata, coord) + + # Test empty image + l3.simulate_cps( + im, metadata, exptime, objlist=[], psf=imdict['impsfgray'], + sky=0) + assert np.all(im.array == 0) # verify nothing in -> nothing out + + # Test flat sky + sky = im.copy() + skycountspersecond = 1 + sky.array[:] = skycountspersecond + im2 = im.copy() + l3.simulate_cps(im2, metadata, exptime, sky=sky, seed=rng_seed) + # verify adding the sky increases the counts + assert np.all(im2.array >= im.array) + # verify that the count rate is about right. + # poisson_rate = skycountspersecond * exptime + poisson_rate = skycountspersecond + assert (np.abs(np.mean(im2.array) - poisson_rate) + < 10 * np.sqrt(poisson_rate / npix)) + + # verify that Poisson noise is included + # pearson chi2 test is probably best here, but it's finicky to get + # right---one needs to choose the right bins so that the convergence + # to the chi^2 distribution is close enough. + # For a Poisson distribution, the variance is equal to the mean rate; + # let's verify that in fact the variance matches expectations within + # some tolerance. + # the variance on the sample variance for a Gaussian is 2*sigma^4/(N-1) + # this isn't a Gaussian but should be close with 100 counts? + var_of_var = 2 * (poisson_rate ** 2) * exptime / (npix - 1) + assert (np.abs(poisson_rate - np.var(im2.array)) + < 10 * np.sqrt(var_of_var)) + + # there are a few WCS bits where we use the positions in the catalog + # to figure out where to render objects. That would require setting + # up a real PSF and is annoying. Skipping that. + # render some objects + im3 = im.copy() + _, objinfo = l3.simulate_cps( + im3, metadata, exptime, objlist=imdict['graycatalog'], psf=imdict['impsfgray'], + xpos=[50, 50], ypos=[50, 50], seed=rng_seed) + + assert np.sum(im3.array) > 0 # at least verify that we added some sources... + assert len(objinfo['objinfo']['array']) == 2 # two sources were added + + im4 = im.copy() + _, objinfo = l3.simulate_cps( + im4, metadata, exptime, objlist=imdict['chromcatalog'], + xpos=[50, 50], ypos=[50, 50], + seed=rng_seed, + psf=imdict['impsfchromatic'], bandpass=imdict['bandpass']) + assert np.sum(im4.array) > 0 # at least verify that we added some sources... + assert len(objinfo['objinfo']['array']) == 2 # two sources were added + + im5 = im.copy() + _, objinfo = l3.simulate_cps( + im5, metadata, exptime, objlist=imdict['chromcatalog'], + psf=imdict['impsfchromatic'], xpos=[1000, 1000], + seed=rng_seed, + ypos=[1000, 1000]) + assert np.sum(objinfo['objinfo']['counts'] > 0) == 0 + # these sources should be out of bounds + + +def test_exptime_array(): + """ Test variable exposure times + """ + # Set random seed + rng_seed = 42 + + # Set image, catalog, etc. + imdict = set_up_image_rendering_things() + im = imdict['im'].copy() + im.array[:] = 0 + chromcat = [imdict['chromcatalog'][0]] + graycat = [imdict['graycatalog'][0]] + coord = SkyCoord(270 * u.deg, 66 * u.deg) + for o in chromcat: + o.sky_pos = coord + for o in graycat: + o.sky_pos = coord + # these are all dumb coordinates; the coord sent to simulate_counts + # is the coordinate of the boresight, but that doesn't need to be on SCA 1. + # But at least they'll exercise some machinery if the ignore_distant_sources + # argument is high enough! + roman.n_pix = 100 + + # Create metadata + meta = util.default_image_meta(filter_name='F158') + wcs.fill_in_parameters(meta, coord) + metadata = deepcopy(parameters.default_mosaic_parameters_dictionary) + filter_name = 'F158' + metadata['basic']['optical_element'] = filter_name + metadata['wcsinfo']['ra_ref'] = 270 + metadata['wcsinfo']['dec_ref'] = 66 + + # Set variable exposure time array + exptime = np.ones((roman.n_pix, roman.n_pix)) + exptime[0:50, :] = 300 + exptime[50:, :] = 600 + + # Set WCS + twcs = romanisim.wcs.get_mosaic_wcs(metadata, shape=(roman.n_pix, roman.n_pix)) + + # Create chromatic data simulation + im1, extras1 = l3.simulate((roman.n_pix, roman.n_pix), twcs, exptime, filter_name, chromcat, + bandpass=imdict['bandpass'], seed=rng_seed, + cps_conv=1, unit_factor=(1 * u.MJy / u.sr), + metadata=metadata, ignore_distant_sources=100, + ) + + # Create filter data simulation + im2, extras2 = l3.simulate((roman.n_pix, roman.n_pix), twcs, exptime, filter_name, graycat, + psf=imdict['impsfgray'], + cps_conv=1, unit_factor=(1 * u.MJy / u.sr), + seed=rng_seed, metadata=metadata, + ignore_distant_sources=100, + ) + + # Ensure that the poisson variance scales with exposure time difference + assert np.isclose((np.mean(im1['var_poisson'][0:50, :].value) / np.mean(im1['var_poisson'][50:, :].value)), 2**0.5, rtol=0.1) + assert np.isclose((np.mean(im2['var_poisson'][0:50, :].value) / np.mean(im2['var_poisson'][50:, :].value)), 2**0.5, rtol=0.1) + + # Ensure that the data remains consistent across the exposure times + assert np.isclose(np.mean(im1['data'][0:50, :].value), np.mean(im1['data'][50:, :].value), rtol=0.2) + assert np.isclose(np.mean(im2['data'][0:50, :].value), np.mean(im2['data'][50:, :].value), rtol=0.2) + +# TBD: Test of geometry construction diff --git a/romanisim/util.py b/romanisim/util.py index d98d7ab1..6dd4e282 100644 --- a/romanisim/util.py +++ b/romanisim/util.py @@ -420,16 +420,16 @@ def decode_context_times(context, exptimes): total_exptimes = np.zeros(context.shape[1:]) - for x in range(total_exptimes.shape[0]): - for y in range(total_exptimes.shape[1]): - files = [v & (1 << k) for v in context[:, x, y] for k in range(nbits)] + for y in range(total_exptimes.shape[0]): + for x in range(total_exptimes.shape[1]): + files = [v & (1 << k) for v in context[:, y, x] for k in range(nbits)] tot_time = 0 files = [file for file in files if (file != 0)] for im_idx in files: tot_time += exptimes[im_idx - 1] - total_exptimes[x][y] = tot_time + total_exptimes[y,x] = tot_time def sum_times(x): tot_time = 0 diff --git a/romanisim/wcs.py b/romanisim/wcs.py index 403f3ad4..8c0f6abf 100644 --- a/romanisim/wcs.py +++ b/romanisim/wcs.py @@ -414,13 +414,22 @@ def convert_wcs_to_gwcs(wcs): return wcs_from_fits_header(wcs.header.header) -def get_mosaic_wcs(mosaic, shape=None): +def get_mosaic_wcs(mosaic, shape=None, xpos=None, ypos=None, coord=None): """Get a WCS object for a given sca or set of CRDS parameters. + - if xpos, ypos, and coords are provided, then a GWCS compatible object will be created (and meta updated with it) + - if not, a functional CelestialWCS is created [useful for quick computation, + but GWCS needed for validation of a final simulation] Parameters ---------- mosaic : roman_datamodels.datamodels.MosaicModel or dict Mosaic model or dictionary containing WCS parameters. + shape: list + Length of dimensions of mosaic + xpos, ypos : array_like (float) + x, y positions of each source in objlist + coord : array_like (float) + ra, dec positions of each source in objlist Returns ------- @@ -434,11 +443,10 @@ def get_mosaic_wcs(mosaic, shape=None): else: mosaic_node = maker_utils.mk_level3_mosaic(shape=shape) for key in mosaic.keys(): - if isinstance(mosaic[key], dict): + if isinstance(mosaic[key], dict) and key in mosaic_node['meta'].keys(): 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 @@ -449,15 +457,22 @@ def get_mosaic_wcs(mosaic, shape=None): if shape is None: shape = (mosaic_node.data.shape[0], mosaic_node.data.shape[1]) - - # 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(math.ceil(shape[0] / 2.0), math.ceil(shape[1] / 2.0)), - world_origin=galsim.PositionD(0, 0)) - wcs = galsim.TanWCS(affine, - util.celestialcoord(world_pos)) - + shape = mosaic_node.data.shape + + if (elem is None for elem in [xpos,ypos,coord]): + # Create a tangent plane WCS for the mosaic + # The affine parameters below should be reviewed and updated + affine = galsim.fitswcs.AffineTransform( + romanisim.parameters.pixel_scale, 0, 0, romanisim.parameters.pixel_scale, origin=galsim.PositionI(x=math.ceil(shape[1] / 2.0), y=math.ceil(shape[0] / 2.0)), + world_origin=galsim.PositionD(0, 0)) + wcs = galsim.fitswcs.TanWCS(affine, + util.celestialcoord(world_pos)) + else: + # Create GWCS compatible tangent plane WCS + header = {} + wcs = galsim.FittedSIPWCS(xpos, ypos, coord[:, 0], coord[:, 1], wcs_type='TAN', center=util.celestialcoord(world_pos)) + wcs._writeHeader(header, galsim.BoundsI(0, shape[0], 0, shape[1])) + # metadata['wcs'] = romanisim.wcs.wcs_from_fits_header(header) return wcs