def create_asdf_in_fits(dtype): """Test fixture to create AsdfInFits object to use for testing""" hdulist = fits.HDUList() hdulist.append(fits.ImageHDU(np.arange(512, dtype=dtype))) hdulist.append(fits.ImageHDU(np.arange(512, dtype=dtype))) hdulist.append(fits.ImageHDU(np.arange(512, dtype=dtype))) tree = { 'model': { 'sci': { 'data': hdulist[0].data, 'wcs': 'WCS info' }, 'dq': { 'data': hdulist[1].data, 'wcs': 'WCS info' }, 'err': { 'data': hdulist[2].data, 'wcs': 'WCS info' } } } return fits_embed.AsdfInFits(hdulist, tree)
def to_fits(tree, schema): hdulist = fits.HDUList() hdulist.append(fits.PrimaryHDU()) _save_from_schema(hdulist, tree, schema) _save_extra_fits(hdulist, tree) _save_history(hdulist, tree) asdf = fits_embed.AsdfInFits(hdulist, tree) return asdf
def test_embed_asdf_in_fits_file(tmpdir, backwards_compat, dtype): fits_testfile = str(tmpdir.join('test.fits')) asdf_testfile = str(tmpdir.join('test.asdf')) hdulist = fits.HDUList() hdulist.append(fits.ImageHDU(np.arange(512, dtype=dtype), name='SCI')) hdulist.append(fits.ImageHDU(np.arange(512, dtype=dtype), name='DQ')) # Test a name with underscores to make sure it works hdulist.append( fits.ImageHDU(np.arange(512, dtype=dtype), name='WITH_UNDERSCORE')) tree = { 'model': { 'sci': { 'data': hdulist['SCI'].data, 'wcs': 'WCS info' }, 'dq': { 'data': hdulist['DQ'].data, 'wcs': 'WCS info' }, 'with_underscore': { 'data': hdulist['WITH_UNDERSCORE'].data, 'wcs': 'WCS info' } } } ff = fits_embed.AsdfInFits(hdulist, tree) ff.write_to(fits_testfile, use_image_hdu=backwards_compat) with fits.open(fits_testfile) as hdulist2: assert len(hdulist2) == 4 assert [x.name for x in hdulist2] == ['SCI', 'DQ', 'WITH_UNDERSCORE', 'ASDF'] assert_array_equal(hdulist2[0].data, np.arange(512, dtype=dtype)) asdf_hdu = hdulist2['ASDF'] assert asdf_hdu.data.tobytes().startswith(b'#ASDF') # When in backwards compatibility mode, the ASDF file will be contained # in an ImageHDU if backwards_compat: assert isinstance(asdf_hdu, fits.ImageHDU) assert asdf_hdu.data.tobytes().strip().endswith(b'...') else: assert isinstance(asdf_hdu, fits.BinTableHDU) with fits_embed.AsdfInFits.open(hdulist2) as ff2: assert_tree_match(tree, ff2.tree) ff = asdf.AsdfFile(copy.deepcopy(ff2.tree)) ff.write_to(asdf_testfile) with asdf.open(asdf_testfile) as ff: assert_tree_match(tree, ff.tree)
def to_fits(tree, schema): hdulist = fits.HDUList() hdulist.append(fits.PrimaryHDU()) _save_from_schema(hdulist, tree, schema) _save_extra_fits(hdulist, tree) _save_history(hdulist, tree) # Store the FITS hash in the tree tree[FITS_HASH_KEY] = fits_hash(hdulist) asdf = fits_embed.AsdfInFits(hdulist, tree) return asdf
def test_serialize_table(tmpdir): tmpfile = str(tmpdir.join('table.fits')) data = np.random.random((10, 10)) table = Table(data) hdu = fits.BinTableHDU(table) hdulist = fits.HDUList() hdulist.append(hdu) tree = {'my_table': hdulist[1].data} with fits_embed.AsdfInFits(hdulist, tree) as ff: ff.write_to(tmpfile) with asdf.open(tmpfile) as ff: data = ff.tree['my_table'] assert data._source.startswith('fits:')
def test_create_in_tree_first(tmpdir, dtype): tree = { 'model': { 'sci': { 'data': np.arange(512, dtype=dtype), 'wcs': 'WCS info' }, 'dq': { 'data': np.arange(512, dtype=dtype), 'wcs': 'WCS info' }, 'err': { 'data': np.arange(512, dtype=dtype), 'wcs': 'WCS info' } } } hdulist = fits.HDUList() hdulist.append(fits.ImageHDU(tree['model']['sci']['data'])) hdulist.append(fits.ImageHDU(tree['model']['dq']['data'])) hdulist.append(fits.ImageHDU(tree['model']['err']['data'])) tmpfile = os.path.join(str(tmpdir), 'test.fits') with fits_embed.AsdfInFits(hdulist, tree) as ff: ff.write_to(tmpfile) with asdf.AsdfFile(tree) as ff: ff.write_to(os.path.join(str(tmpdir), 'plain.asdf')) with asdf.open(os.path.join(str(tmpdir), 'plain.asdf')) as ff: assert_array_equal(ff.tree['model']['sci']['data'], np.arange(512, dtype=dtype)) # This tests the changes that allow FITS files with ASDF extensions to be # opened directly by the top-level asdf.open API with asdf_open(tmpfile) as ff: assert_array_equal(ff.tree['model']['sci']['data'], np.arange(512, dtype=dtype))
def cube(tmpdir, tmp_asdf): """ Mock a JWST s3d cube """ hdulist = fits.HDUList() hdulist.append(fits.PrimaryHDU()) hdulist["PRIMARY"].header["TELESCOP"] = ("JWST", "comment") hdulist["PRIMARY"].header["FLUXEXT"] = ("ERR", "comment") hdulist["PRIMARY"].header["ERREXT"] = ("ERR", "comment") hdulist["PRIMARY"].header["MASKEXT"] = ("DQ", "comment") # Add ImageHDU for cubes shape = [30, 10, 10] hdulist.append(create_image_hdu(name='SCI', shape=shape, hdrs=[("BUNIT", 'MJy')])) hdulist.append(create_image_hdu(name='ERR', shape=shape, hdrs=[("BUNIT", 'MJy'), ('ERRTYPE', 'ERR')])) hdulist.append(create_image_hdu(name='DQ', shape=shape)) # Mock the ASDF extension hdulist.append(fits.BinTableHDU(name='ASDF')) ff = fits_embed.AsdfInFits(hdulist, tmp_asdf) tmpfile = str(tmpdir.join('jwst_embedded_asdf.fits')) ff.write_to(tmpfile) return hdulist
def perform_dither_distortion(plot=False): """ Dither the initial direct image and apply detector distortion. Args: plot (Bool): True if plots should be saved. """ # Make four dithered exposures following the standard template pattern # or whatever specified pattern. It reopens image just created and it # uses a reference rate image to create the right and complete header. # Dithers for WFSS (use the sim image just created) ref = fits.open('{0}-{1}.fits'.format(OUTROOT, NIS_FILTER.lower())) ref_model = grizli.jwst.hdu_to_imagemodel(ref[0]) blot_parent = gwcs_blot.GWCSBlot(ref_model) # Center position on the image. r0 = RA_CENTER * u.deg d0 = DEC_CENTER * u.deg # Dither pattern dx = DITHER_X dy = DITHER_Y if plot is True: plt.figure(figsize=(20, 20)) #Make the 4 dithers with distortion for pos in range(4): rate_im, aa = read_rate_im() aa['data'] = rate_im['SCI'].data aa['dq'] = rate_im['DQ'].data wcsinfo = aa['meta']['wcsinfo'] pointing = aa['meta']['pointing'] inst = aa['meta']['instrument'] print('Position {0}, {1}'.format(pos + 1, NIS_FILTER)) inst['filter'] = rate_im[0].header['FILTER'] = 'CLEAR' inst['pupil'] = rate_im[0].header['PUPIL'] = NIS_FILTER.upper() cosd = np.cos(d0 / 180 * np.pi) pos_ra = (r0 + dx[pos] / cosd * u.arcsec).value pos_dec = (d0 + dy[pos] * u.arcsec).value #Update header in rate image and wcf info rate_im[1].header['CRVAL1'] = rate_im[1].header['RA_REF'] = pos_ra rate_im[1].header['CRVAL2'] = rate_im[1].header['DEC_REF'] = pos_dec wcsinfo['crval1'] = wcsinfo['ra_ref'] = pointing['ra_v1'] = pos_ra wcsinfo['crval2'] = wcsinfo['dec_ref'] = pointing['dec_v1'] = pos_dec # Blot reference data to distorted frame rate_model = grizli.jwst.img_with_wcs(rate_im) blot_child = blot_parent.extract_image(rate_model, interp='poly5') rate_im['SCI'].data = blot_child # zero-out DQ array rate_im['DQ'].data *= 0 # rate_im[0].header['TARGNAME'] = aa['meta']['target'] = OUTROOT print('{0}-{1}-clear-pos{2}_rate.fits'.format(OUTROOT, NIS_FILTER, pos + 1)) ff = fits_embed.AsdfInFits(rate_im, {'meta': aa['meta']}) ff.write_to('{0}-{1}-clear-pos{2}_rate.fits'.format( OUTROOT, NIS_FILTER, pos + 1), overwrite=True) # Make grism files (sci extension is direct image for now) for grism in ['gr150r', 'gr150c']: inst['filter'] = rate_im[0].header['FILTER'] = grism.upper() out = '{0}-{1}-{grism}-pos{2}_rate.fits'.format(OUTROOT, NIS_FILTER, pos + 1, grism=grism) print(out) ff = fits_embed.AsdfInFits(rate_im, aa['meta']) ff.write_to(out, overwrite=True) rate_im.close() # Show dithered images if plot is True: plt.subplot(2, 2, pos + 1, projection=pywcs.WCS(rate_im[1].header)) plt.title("Dither=({}, {})".format(dx[pos], dy[pos])) plt.imshow(np.log10(rate_im['SCI'].data)) plt.grid(color='black', ls='solid') figname = "{}_dithers.png".format(OUTROOT) plt.savefig(figname) print("Wrote {}".format(figname))