diff --git a/romanisim/l3.py b/romanisim/l3.py new file mode 100644 index 00000000..0b47022c --- /dev/null +++ b/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 diff --git a/romanisim/parameters.py b/romanisim/parameters.py index bf37d481..221a14f2 100644 --- a/romanisim/parameters.py +++ b/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. }, } diff --git a/romanisim/tests/test_image.py b/romanisim/tests/test_image.py index 99319e85..554b2227 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, 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')) diff --git a/romanisim/tests/test_wcs.py b/romanisim/tests/test_wcs.py index 96236f44..9b25067c 100644 --- a/romanisim/tests/test_wcs.py +++ b/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(): diff --git a/romanisim/wcs.py b/romanisim/wcs.py index ffdc9ba2..04830824 100644 --- a/romanisim/wcs.py +++ b/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