def test_indexing_on_instance(): """Test indexing on compound model instances.""" m = Gaussian1D(1, 0, 0.1) + Const1D(2) assert isinstance(m[0], Gaussian1D) assert isinstance(m[1], Const1D) assert m.param_names == ('amplitude_0', 'mean_0', 'stddev_0', 'amplitude_1') # Test parameter equivalence assert m[0].amplitude == 1 == m.amplitude_0 assert m[0].mean == 0 == m.mean_0 assert m[0].stddev == 0.1 == m.stddev_0 assert m[1].amplitude == 2 == m.amplitude_1 # Test that parameter value updates are symmetric between the compound # model and the submodel returned by indexing const = m[1] m.amplitude_1 = 42 assert const.amplitude == 42 const.amplitude = 137 assert m.amplitude_1 == 137 # Similar couple of tests, but now where the compound model was created # from model instances g = Gaussian1D(1, 2, 3, name='g') p = Polynomial1D(2, name='p') m = g + p assert m[0].name == 'g' assert m[1].name == 'p' assert m['g'].name == 'g' assert m['p'].name == 'p' poly = m[1] m.c0_1 = 12345 assert poly.c0 == 12345 poly.c1 = 6789 assert m.c1_1 == 6789 # Test negative indexing assert isinstance(m[-1], Polynomial1D) assert isinstance(m[-2], Gaussian1D) with pytest.raises(IndexError): m[42] with pytest.raises(IndexError): m['foobar'] # Confirm index-by-name works with fix_inputs g = Gaussian2D(1, 2, 3, 4, 5, name='g') m = fix_inputs(g, {0: 1}) assert m['g'].name == 'g' # Test string slicing A = Const1D(1.1, name='A') B = Const1D(2.1, name='B') C = Const1D(3.1, name='C') M = A + B * C assert_allclose(M['B':'C'](1), 6.510000000000001)
def test_model_set_raises_value_error(expr, result): """Check that creating model sets with components whose _n_models are different raise a value error """ with pytest.raises(ValueError): s = expr(Const1D((2, 2), n_models=2), Const1D(3, n_models=1))
def gwa_to_ifuslit(slits, disperser, wrange, order, reference_files): """ GWA to SLIT transform. Parameters ---------- slits : list A list of slit IDs for all IFU slits 0-29. disperser : dict A corrected disperser ASDF object. filter : str The filter used. grating : str The grating used in the observation. reference_files: dict Dictionary with reference files returned by CRDS. Returns ------- model : `~jwst_lib.pipeline_models.Gwa2Slit` model. Transform from GWA frame to SLIT frame. """ agreq = AngleFromGratingEquation(disperser['groove_density'], order, name='alpha_from_greq') lgreq = WavelengthFromGratingEquation(disperser['groove_density'], order, name='lambda_from_greq') collimator2gwa = collimator_to_gwa(reference_files, disperser) lam = (wrange[1] - wrange[0]) / 2 + wrange[0] ifuslicer = AsdfFile.open(reference_files['ifuslicer']) ifupost = AsdfFile.open(reference_files['ifupost']) slit_models = {} ifuslicer_model = ifuslicer.tree['model'] for slit in slits: slitdata = ifuslicer.tree['data'][slit] slitdata_model = get_slit_location_model(slitdata) ifuslicer_transform = (slitdata_model | ifuslicer_model) ifupost_transform = ifupost.tree[slit]['model'] msa2gwa = ifuslicer_transform | ifupost_transform | collimator2gwa gwa2msa = gwa_to_ymsa(msa2gwa) # TODO: Use model sets here bgwa2msa = Mapping((0, 1, 0, 1), n_inputs=3) | \ Const1D(0) * Identity(1) & Const1D(-1) * Identity(1) & Identity(2) | \ Identity(1) & gwa2msa & Identity(2) | \ Mapping((0, 1, 0, 1, 2, 3)) | Identity(2) & msa2gwa & Identity(2) | \ Mapping((0, 1, 2, 5), n_inputs=7)| Identity(2) & lgreq # msa to before_gwa #msa2bgwa = Mapping((0, 1, 2, 2)) | msa2gwa & Identity(1) | Mapping((3, 0, 1, 2)) | agreq msa2bgwa = msa2gwa & Identity(1) | Mapping((3, 0, 1, 2)) | agreq bgwa2msa.inverse = msa2bgwa slit_models[slit] = bgwa2msa ifuslicer.close() ifupost.close() return Gwa2Slit(slit_models)
def imaging(input_model, reference_files): """ Imaging pipeline frames : detector, gwa, msa, sky """ # Get the corrected disperser model disperser = get_disperser(input_model, reference_files['disperser']) # DETECTOR to GWA transform det2gwa = detector_to_gwa(reference_files, input_model.meta.instrument.detector, disperser) gwa_through = Const1D(-1) * Identity(1) & Const1D(-1) * Identity( 1) & Identity(1) angles = [ disperser['theta_x'], disperser['theta_y'], disperser['theta_z'], disperser['tilt_y'] ] rotation = Rotation3DToGWA(angles, axes_order="xyzy", name='rotaton').inverse dircos2unitless = DirCos2Unitless(name='directional_cosines2unitless') col = AsdfFile.open(reference_files['collimator']).tree['model'] # Get the default spectral order and wavelength range and record them in the model. sporder, wrange = get_spectral_order_wrange( input_model, reference_files['wavelengthrange']) input_model.meta.wcsinfo.waverange_start = wrange[0] input_model.meta.wcsinfo.waverange_end = wrange[1] input_model.meta.wcsinfo.spectral_order = sporder lam = wrange[0] + (wrange[1] - wrange[0]) * .5 lam_model = Mapping((0, 1, 1)) | Identity(2) & Const1D(lam) gwa2msa = gwa_through | rotation | dircos2unitless | col | lam_model gwa2msa.inverse = col.inverse | dircos2unitless.inverse | rotation.inverse | gwa_through # MSA to OTEIP transform msa2ote = msa_to_oteip(reference_files) msa2oteip = msa2ote | Mapping((0, 1), n_inputs=3) msa2oteip.inverse = Mapping((0, 1, 0, 1)) | msa2ote.inverse | Mapping( (0, 1), n_inputs=3) # OTEIP to V2,V3 transform with AsdfFile.open(reference_files['ote']) as f: oteip2v23 = f.tree['model'].copy() # Create coordinate frames in the NIRSPEC WCS pipeline # "detector", "gwa", "msa", "oteip", "v2v3", "world" det, gwa, msa_frame, oteip, v2v3 = create_imaging_frames() imaging_pipeline = [(det, det2gwa), (gwa, gwa2msa), (msa_frame, msa2oteip), (oteip, oteip2v23), (v2v3, None)] return imaging_pipeline
def gwa_to_slit(slits_id, disperser, wrange, order, reference_files): """ GWA to SLIT transform. Parameters ---------- slits_id : list A list of slit IDs for all open shutters/slitlets. disperser : dict A corrected disperser ASDF object. filter : str The filter used. grating : str The grating used in the observation. reference_files: dict Dictionary with reference files returned by CRDS. Returns ------- model : `~jwst_lib.pipeline_models.Gwa2Slit` model. Transform from GWA frame to SLIT frame. """ agreq = AngleFromGratingEquation(disperser['groove_density'], order, name='alpha_from_greq') lgreq = WavelengthFromGratingEquation(disperser['groove_density'], order, name='lambda_from_greq') collimator2gwa = collimator_to_gwa(reference_files, disperser) msa = AsdfFile.open(reference_files['msa']) slit_models = {} for i in range(1, 6): slit_names = slits_id[slits_id[:, 0] == i] if slit_names.any(): msa_model = msa.tree[i]['model'] for slit in slit_names: index = slit[1] - 1 slitdata = msa.tree[slit[0]]['data'][index] slitdata_model = get_slit_location_model(slitdata) msa_transform = slitdata_model | msa_model msa2gwa = (msa_transform | collimator2gwa) gwa2msa = gwa_to_ymsa(msa2gwa) # TODO: Use model sets here bgwa2msa = Mapping((0, 1, 0, 1), n_inputs=3) | \ Const1D(0) * Identity(1) & Const1D(-1) * Identity(1) & Identity(2) | \ Identity(1) & gwa2msa & Identity(2) | \ Mapping((0, 1, 0, 1, 2, 3)) | Identity(2) & msa2gwa & Identity(2) | \ Mapping((0, 1, 2, 5), n_inputs=7)| Identity(2) & lgreq # msa to before_gwa msa2bgwa = msa2gwa & Identity(1) | Mapping( (3, 0, 1, 2)) | agreq bgwa2msa.inverse = msa2bgwa s = slitid_to_slit(np.array([slit]))[0] slit_models[s] = bgwa2msa msa.close() return Gwa2Slit(slit_models)
def test_input_shape_1d(): m1 = Const1D() m2 = Const1D() model = convolve_models_fft(m1, m2, (-1, 1), 0.01) results = model(0) assert results.shape == (1, ) x = np.arange(-1, 1, 0.1) results = model(x) assert results.shape == x.shape
def test_two_model_instance_arithmetic_1d(expr, result): """ Like test_two_model_class_arithmetic_1d, but creates a new model from two model *instances* with fixed parameters. """ s = expr(Const1D(2), Const1D(3)) assert isinstance(s, Model) assert s.n_inputs == 1 assert s.n_outputs == 1 out = s(0) assert out == result assert isinstance(out, float)
def create_slit(model, x0, y0, order): """ Create a SlitModel representing a grism slit.""" ymin = 0 xmin = 0 # ymax = 58 # xmax = 1323 model = Mapping((0, 1, 0, 0, 0)) | (Shift(xmin) & Shift(ymin) & Const1D(x0) & Const1D(y0) & Const1D(order)) | model wcsobj = wcs.WCS([('det', model), ('world', None)]) wcsobj.bounding_box = ((20, 25), (800, 805)) slit = SlitModel() slit.meta.wcs = wcsobj slit.source_xpos = x0 slit.source_ypos = y0 return slit
def get_astropy_model(model): astropy_model = None if model["type"] == "Const": astropy_model = Const1D(model["amplitude"]) elif model["type"] == "Gaussian": astropy_model = Gaussian1D(model["amplitude"], model["mean"], model["stddev"]) elif model["type"] == "Lorentz": astropy_model = Lorentz1D(model["amplitude"], model["x_0"], model["fwhm"]) elif model["type"] == "PowerLaw": astropy_model = PowerLaw1D(model["amplitude"], model["x_0"], model["alpha"]) elif model["type"] == "BrokenPowerLaw": astropy_model = BrokenPowerLaw1D(model["amplitude"], model["x_break"], model["alpha_1"], model["alpha_2"]) if astropy_model: astropy_model = fix_parameters_to_astropy_model(astropy_model, model) return astropy_model
def test_clear_cache(): m1 = Const1D() m2 = Const1D() model = convolve_models_fft(m1, m2, (-1, 1), 0.01) assert model._kwargs is None assert model._convolution is None results = model(0) assert results.all() == np.array([1.]).all() assert model._kwargs is not None assert model._convolution is not None model.clear_cache() assert model._kwargs is None assert model._convolution is None
def oneD_gaussian(shape=200, mean=0., std=10., amp=1., bkg=0.): from astropy.modeling.models import Gaussian1D, Const1D mod = Gaussian1D(mean=mean, stddev=std, amplitude=amp) + \ Const1D(amplitude=bkg) return mod
def create_spectral_wcs(ra, dec, wavelength): """Assign a WCS for sky coordinates and a table of wavelengths Parameters ---------- ra: float The right ascension (in degrees) at the nominal location of the entrance aperture (slit). dec: float The declination (in degrees) at the nominal location of the entrance aperture. wavelength: ndarray The wavelength in microns at each pixel of the extracted spectrum. Returns: -------- wcs: a gwcs.wcs.WCS object This takes a float or sequence of float and returns a tuple of the right ascension, declination, and wavelength (or sequence of wavelengths) at the pixel(s) specified by the input argument. """ # Only the first coordinate is used. input_frame = cf.Frame2D(axes_order=(0, 1), unit=(u.pix, u.pix), name="pixel_frame") sky = cf.CelestialFrame(name='sky', axes_order=(0, 1), reference_frame=coord.ICRS()) spec = cf.SpectralFrame(name='spectral', axes_order=(2, ), unit=(u.micron, ), axes_names=('wavelength', )) pixel = np.arange(len(wavelength), dtype=np.float) tab = Mapping((0, 0, 0)) | \ Const1D(ra) & Const1D(dec) & Tabular1D(pixel, wavelength) world = cf.CompositeFrame([sky, spec], name='world') pipeline = [(input_frame, tab), (world, None)] return WCS(pipeline)
def test_get_metadata(): m1 = Const1D() m1.meta['description'] = 'a constant' m1.meta['foo'] = 42.0 m2 = Const1D() m2.meta['description'] = 'another constant' m3 = Const1D() m3.meta[42] = 'answer' m = (m1 + m2) * m3 meta = get_metadata(m) keys = list(meta.keys()) assert len(keys) == 3 assert meta['description'] == 'another constant' assert meta['foo'] == 42 assert meta[42] == 'answer'
def _make_reference_gwcs_wcs(fits_hdr): hdr = fits.Header.fromfile(get_pkg_data_filename(fits_hdr)) fw = fitswcs.WCS(hdr) unit_conv = Scale(1.0 / 3600.0, name='arcsec_to_deg_1D') unit_conv = unit_conv & unit_conv unit_conv.name = 'arcsec_to_deg_2D' unit_conv_inv = Scale(3600.0, name='deg_to_arcsec_1D') unit_conv_inv = unit_conv_inv & unit_conv_inv unit_conv_inv.name = 'deg_to_arcsec_2D' c2s = CartesianToSpherical(name='c2s', wrap_lon_at=180) s2c = SphericalToCartesian(name='s2c', wrap_lon_at=180) c2tan = ((Mapping((0, 1, 2), name='xyz') / Mapping( (0, 0, 0), n_inputs=3, name='xxx')) | Mapping((1, 2), name='xtyt')) c2tan.name = 'Cartesian 3D to TAN' tan2c = (Mapping((0, 0, 1), n_inputs=2, name='xtyt2xyz') | (Const1D(1, name='one') & Identity(2, name='I(2D)'))) tan2c.name = 'TAN to cartesian 3D' tan2c.inverse = c2tan c2tan.inverse = tan2c aff = AffineTransformation2D(matrix=fw.wcs.cd) offx = Shift(-fw.wcs.crpix[0]) offy = Shift(-fw.wcs.crpix[1]) s = 5e-6 scale = Scale(s) & Scale(s) det2tan = (offx & offy) | scale | tan2c | c2s | unit_conv_inv taninv = s2c | c2tan tan = Pix2Sky_TAN() n2c = RotateNative2Celestial(fw.wcs.crval[0], fw.wcs.crval[1], 180) wcslin = unit_conv | taninv | scale.inverse | aff | tan | n2c sky_frm = cf.CelestialFrame(reference_frame=coord.ICRS()) det_frm = cf.Frame2D(name='detector') v2v3_frm = cf.Frame2D(name="v2v3", unit=(u.arcsec, u.arcsec), axes_names=('x', 'y'), axes_order=(0, 1)) pipeline = [(det_frm, det2tan), (v2v3_frm, wcslin), (sky_frm, None)] gw = gwcs.WCS(input_frame=det_frm, output_frame=sky_frm, forward_transform=pipeline) gw.crpix = fw.wcs.crpix gw.crval = fw.wcs.crval gw.bounding_box = ((-0.5, fw.pixel_shape[0] - 0.5), (-0.5, fw.pixel_shape[1] - 0.5)) return gw
def _get_wavelength_calibration(hdr): from astropy.modeling.models import Linear1D, Const1D _wcal_model = ( Const1D(hdr.get('CRVAL1')) + Linear1D(slope=hdr.get('CD1_1'), intercept=hdr.get('CRPIX1') - 1)) assert _wcal_model(0) == hdr.get('CRVAL1') return _wcal_model
def test_indexing_on_instance(): """Test indexing on compound model instances.""" m = Gaussian1D(1, 0, 0.1) + Const1D(2) assert isinstance(m[0], Gaussian1D) assert isinstance(m[1], Const1D) # assert isinstance(m['Gaussian1D'], Gaussian1D) # assert isinstance(m['Const1D'], Const1D) # Test parameter equivalence assert m[0].amplitude == 1 assert m[0].mean == 0 assert m[0].stddev == 0.1 assert m[1].amplitude == 2 # Similar couple of tests, but now where the compound model was created # from model instances g = Gaussian1D(1, 2, 3, name='g') p = Polynomial1D(2, name='p') m = g + p assert m[0].name == 'g' assert m[1].name == 'p' assert m['g'].name == 'g' assert m['p'].name == 'p' # Test negative indexing assert isinstance(m[-1], Polynomial1D) assert isinstance(m[-2], Gaussian1D) with pytest.raises(IndexError): m[42] with pytest.raises(IndexError): m['foobar'] # Test string slicing A = Const1D(1.1, name='A') B = Const1D(2.1, name='B') C = Const1D(3.1, name='C') M = A + B * C assert_allclose(M['B':'C'](1), 6.510000000000001)
def _init_bgmodel(lorentz_mean=15.0): """Initialize model for background: constant plus Lorentzian""" lorentz_fixed = {'x_0': True, 'fwhm': True} lorentz_bounds = {'amplitude': [0, None]} constant_bounds = {'amplitude': [0, CORE_CMAX]} bgmodel = (Lorentz1D(0.1, lorentz_mean, 100.0, name='Lorentz', bounds=lorentz_bounds, fixed=lorentz_fixed) + Const1D(1.0e-4, bounds=constant_bounds, name='Constant')) return bgmodel
def _v2v3_to_tpcorr_from_full(tpcorr): s2c = tpcorr['s2c'] c2s = tpcorr['c2s'] # unit_conv = _get_submodel(tpcorr, 'arcsec_to_deg_2D') # unit_conv_inv = _get_submodel(tpcorr, 'deg_to_arcsec_2D') # # The code below is a work-around to the code above not working. # TODO: understand why _get_submodel is unable to retrieve # some submodels in a compound model. # unit_conv = _get_submodel(tpcorr, 'arcsec_to_deg_1D') unit_conv = unit_conv & unit_conv unit_conv.name = 'arcsec_to_deg_2D' unit_conv_inv = _get_submodel(tpcorr, 'deg_to_arcsec_1D') unit_conv_inv = unit_conv_inv & unit_conv_inv unit_conv_inv.name = 'deg_to_arcsec_2D' affine = tpcorr['tp_affine'] affine_inv = affine.inverse affine_inv.name = 'tp_affine_inv' rot = tpcorr['det_to_optic_axis'] rot_inv = rot.inverse rot_inv.name = 'optic_axis_to_det' # c2tan = _get_submodel(tpcorr, 'Cartesian 3D to TAN') # tan2c = _get_submodel(tpcorr, 'TAN to cartesian 3D') # # The code below is a work-around to the code above not working. # TODO: understand why _get_submodel is unable to retrieve # some submodels in a compound model. # c2tan = ((Mapping((0, 1, 2), name='xyz') / Mapping( (0, 0, 0), n_inputs=3, name='xxx')) | Mapping((1, 2), name='xtyt')) c2tan.name = 'Cartesian 3D to TAN' tan2c = (Mapping((0, 0, 1), n_inputs=2, name='xtyt2xyz') | (Const1D(1, name='one') & Identity(2, name='I(2D)'))) tan2c.name = 'TAN to cartesian 3D' v2v3_to_tpcorr = unit_conv | s2c | rot | c2tan | affine v2v3_to_tpcorr.name = 'jwst_v2v3_to_tpcorr' tpcorr_to_v2v3 = affine_inv | tan2c | rot_inv | c2s | unit_conv_inv tpcorr_to_v2v3.name = 'jwst_tpcorr_to_v2v3' v2v3_to_tpcorr.inverse = tpcorr_to_v2v3 tpcorr_to_v2v3.inverse = v2v3_to_tpcorr return v2v3_to_tpcorr
def fit_data_with_lorentz_and_const(x_values, y_values): amplitude = 5. x_0 = 1 fwhm = 0.5 const = 5. g_init = Lorentz1D(amplitude, x_0, fwhm) g_init += Const1D(const) lpost = PSDLogLikelihood(x_values, y_values, g_init) parest = ParameterEstimation() res = parest.fit(lpost, [amplitude, x_0, fwhm, const], neg=True) opt_amplitude = res.p_opt[0] opt_x_0 = res.p_opt[1] opt_fwhm = res.p_opt[2] opt_const = res.p_opt[3] return opt_amplitude, opt_x_0, opt_fwhm, opt_const
def test_two_model_mixed_arithmetic_1d(expr, result): """ Like test_two_model_class_arithmetic_1d, but creates a new model from an expression of one model class with one model instance (and vice-versa). """ S1 = expr(Const1D, Const1D(3)) S2 = expr(Const1D(2), Const1D) for cls in (S1, S2): assert issubclass(cls, Model) assert cls.n_inputs == 1 assert cls.n_outputs == 1 # Requires values for both amplitudes even though one of them them has a # default # TODO: We may wish to fix that eventually, so that if a parameter has a # default it doesn't *have* to be given in the init s1 = S1(2, 3) s2 = S2(2, 3) for out in (s1(0), s2(0)): assert out == result assert isinstance(out, float)
def test_slicing_on_class(): """ Test slicing a simple compound model class using integers. """ A = Const1D.rename('A') B = Const1D.rename('B') C = Const1D.rename('C') D = Const1D.rename('D') E = Const1D.rename('E') F = Const1D.rename('F') M = A + B - C * D / E ** F assert M[0:1] is A # This test will also check that the correct parameter names are generated # for each slice (fairly trivial in this case since all the submodels have # the same parameter, but if any corner cases are found that aren't covered # by this test we can do something different...) assert M[0:1].param_names == ('amplitude',) # This looks goofy but if you slice by name to the sub-model of the same # name it should just return that model, logically. assert M['A':'A'] is A assert M['A':'A'].param_names == ('amplitude',) assert M[5:6] is F assert M[5:6].param_names == ('amplitude',) assert M['F':'F'] is F assert M['F':'F'].param_names == ('amplitude',) # 1 + 2 assert M[:2](1, 2)(0) == 3 assert M[:2].param_names == ('amplitude_0', 'amplitude_1') assert M[:'B'](1, 2)(0) == 3 assert M[:'B'].param_names == ('amplitude_0', 'amplitude_1') # 2 - 3 assert M[1:3](2, 3)(0) == -1 assert M[1:3].param_names == ('amplitude_1', 'amplitude_2') assert M['B':'C'](2, 3)(0) == -1 assert M['B':'C'].param_names == ('amplitude_1', 'amplitude_2') # 3 * 4 assert M[2:4](3, 4)(0) == 12 assert M[2:4].param_names == ('amplitude_2', 'amplitude_3') assert M['C':'D'](3, 4)(0) == 12 assert M['C':'D'].param_names == ('amplitude_2', 'amplitude_3') # 4 / 5 assert M[3:5](4, 5)(0) == 0.8 assert M[3:5].param_names == ('amplitude_3', 'amplitude_4') assert M['D':'E'](4, 5)(0) == 0.8 assert M['D':'E'].param_names == ('amplitude_3', 'amplitude_4') # 5 ** 6 assert M[4:6](5, 6)(0) == 15625 assert M[4:6].param_names == ('amplitude_4', 'amplitude_5') assert M['E':'F'](5, 6)(0) == 15625 assert M['E':'F'].param_names == ('amplitude_4', 'amplitude_5')
def test_slicing_on_class(): """ Test slicing a simple compound model class using integers. """ A = Const1D.rename('A') B = Const1D.rename('B') C = Const1D.rename('C') D = Const1D.rename('D') E = Const1D.rename('E') F = Const1D.rename('F') M = A + B - C * D / E**F assert M[0:1] is A # This test will also check that the correct parameter names are generated # for each slice (fairly trivial in this case since all the submodels have # the same parameter, but if any corner cases are found that aren't covered # by this test we can do something different...) assert M[0:1].param_names == ('amplitude', ) # This looks goofy but if you slice by name to the sub-model of the same # name it should just return that model, logically. assert M['A':'A'] is A assert M['A':'A'].param_names == ('amplitude', ) assert M[5:6] is F assert M[5:6].param_names == ('amplitude', ) assert M['F':'F'] is F assert M['F':'F'].param_names == ('amplitude', ) # 1 + 2 assert M[:2](1, 2)(0) == 3 assert M[:2].param_names == ('amplitude_0', 'amplitude_1') assert M[:'B'](1, 2)(0) == 3 assert M[:'B'].param_names == ('amplitude_0', 'amplitude_1') # 2 - 3 assert M[1:3](2, 3)(0) == -1 assert M[1:3].param_names == ('amplitude_1', 'amplitude_2') assert M['B':'C'](2, 3)(0) == -1 assert M['B':'C'].param_names == ('amplitude_1', 'amplitude_2') # 3 * 4 assert M[2:4](3, 4)(0) == 12 assert M[2:4].param_names == ('amplitude_2', 'amplitude_3') assert M['C':'D'](3, 4)(0) == 12 assert M['C':'D'].param_names == ('amplitude_2', 'amplitude_3') # 4 / 5 assert M[3:5](4, 5)(0) == 0.8 assert M[3:5].param_names == ('amplitude_3', 'amplitude_4') assert M['D':'E'](4, 5)(0) == 0.8 assert M['D':'E'].param_names == ('amplitude_3', 'amplitude_4') # 5 ** 6 assert M[4:6](5, 6)(0) == 15625 assert M[4:6].param_names == ('amplitude_4', 'amplitude_5') assert M['E':'F'](5, 6)(0) == 15625 assert M['E':'F'].param_names == ('amplitude_4', 'amplitude_5')
def test_compound_evaluate_power(): """ Tests that compound evaluate function produces the same result as the models with the power operator applied """ x = np.linspace(-5, 5, 10) p1 = np.array([1, 0, 0.2]) p2 = np.array([3]) model1 = Gaussian1D(2, 1, 5) model2 = Const1D(2) compound = model1 ** model2 assert_array_equal( compound.evaluate(x, *p1, *p2), model1.evaluate(x, *p1) ** model2.evaluate(x, *p2), )
def _tpcorr_init(v2_ref, v3_ref, roll_ref): s2c = SphericalToCartesian(name='s2c', wrap_lon_at=180) c2s = CartesianToSpherical(name='c2s', wrap_lon_at=180) unit_conv = Scale(1.0 / 3600.0, name='arcsec_to_deg_1D') unit_conv = unit_conv & unit_conv unit_conv.name = 'arcsec_to_deg_2D' unit_conv_inv = Scale(3600.0, name='deg_to_arcsec_1D') unit_conv_inv = unit_conv_inv & unit_conv_inv unit_conv_inv.name = 'deg_to_arcsec_2D' affine = AffineTransformation2D(name='tp_affine') affine_inv = AffineTransformation2D(name='tp_affine_inv') rot = RotationSequence3D([v2_ref, -v3_ref, roll_ref], 'zyx', name='det_to_optic_axis') rot_inv = rot.inverse rot_inv.name = 'optic_axis_to_det' # projection submodels: c2tan = ((Mapping((0, 1, 2), name='xyz') / Mapping( (0, 0, 0), n_inputs=3, name='xxx')) | Mapping((1, 2), name='xtyt')) c2tan.name = 'Cartesian 3D to TAN' tan2c = (Mapping((0, 0, 1), n_inputs=2, name='xtyt2xyz') | (Const1D(1, name='one') & Identity(2, name='I(2D)'))) tan2c.name = 'TAN to cartesian 3D' total_corr = (unit_conv | s2c | rot | c2tan | affine | tan2c | rot_inv | c2s | unit_conv_inv) total_corr.name = 'JWST tangent-plane linear correction. v1' inv_total_corr = (unit_conv | s2c | rot | c2tan | affine_inv | tan2c | rot_inv | c2s | unit_conv_inv) inv_total_corr.name = 'Inverse JWST tangent-plane linear correction. v1' # TODO # re-enable circular inverse definitions once # https://github.com/spacetelescope/asdf/issues/744 or # https://github.com/spacetelescope/asdf/issues/745 are resolved. # # inv_total_corr.inverse = total_corr total_corr.inverse = inv_total_corr return total_corr
def from_tree_transform(cls, node, ctx): mapping = node['mapping'] n_inputs = node.get('n_inputs') if all([isinstance(x, int) for x in mapping]): return Mapping(tuple(mapping), n_inputs) if n_inputs is None: n_inputs = max([x for x in mapping if isinstance(x, int)]) + 1 transform = Identity(n_inputs) new_mapping = [] i = n_inputs for entry in mapping: if isinstance(entry, int): new_mapping.append(entry) else: new_mapping.append(i) transform = transform & Const1D(entry.value) i += 1 return transform | Mapping(new_mapping)
def oteip_to_v23(reference_files): """ Transform from the OTEIP frame to the V2V3 frame. Parameters ---------- reference_files: dict Dictionary with reference files returned by CRDS. Returns ------- model : `~astropy.modeling.core.Model` model. Transform from OTEIP to V2V3. """ with AsdfFile.open(reference_files['ote']) as f: ote = f.tree['model'].copy() fore2ote_mapping = Identity(3, name='fore2ote_mapping') fore2ote_mapping.inverse = Mapping((0, 1, 2, 2)) # Convert the wavelength to microns return fore2ote_mapping | (ote & Identity(1) / Const1D(1e-6))
def niriss_soss(input_model, reference_files): """ The NIRISS SOSS pipeline includes 3 coordinate frames - detector, focal plane and sky reference_files={'specwcs': 'soss_wavelengths_configuration.asdf'} """ # Get the target RA and DEC, they will be used for setting the WCS RA and DEC based on a conversation # with Kevin Volk. try: target_ra = float(input_model['meta.target.ra']) target_dec = float(input_model['meta.target.dec']) except: # There was an error getting the target RA and DEC, so we are not going to continue. raise ValueError('Problem getting the TARG_RA or TARG_DEC from input model {}'.format(input_model)) # Define the frames detector = cf.Frame2D(name='detector', axes_order=(0, 1), unit=(u.pix, u.pix)) spec = cf.SpectralFrame(name='spectral', axes_order=(2,), unit=(u.micron,), axes_names=('wavelength',)) sky = cf.CelestialFrame(reference_frame=coord.ICRS(), axes_names=('ra', 'dec'), axes_order=(0, 1), unit=(u.deg, u.deg), name='sky') world = cf.CompositeFrame([sky, spec], name='world') try: with AsdfFile.open(reference_files['specwcs']) as wl: wl1 = wl.tree[1].copy() wl2 = wl.tree[2].copy() wl3 = wl.tree[3].copy() except Exception as e: raise IOError('Error reading wavelength correction from {}'.format(reference_files['specwcs'])) cm_order1 = (Mapping((0, 1, 0, 1)) | (Const1D(target_ra) & Const1D(target_dec) & wl1)).rename('Order1') cm_order2 = (Mapping((0, 1, 0, 1)) | (Const1D(target_ra) & Const1D(target_dec) & wl2)).rename('Order2') cm_order3 = (Mapping((0, 1, 0, 1)) | (Const1D(target_ra) & Const1D(target_dec) & wl3)).rename('Order3') # Define the transforms, they should accept (x,y) and return (ra, dec, lambda) soss_model = NirissSOSSModel([1, 2, 3], [cm_order1, cm_order2, cm_order3]).rename('3-order SOSS Model') # Define the pipeline based on the frames and models above. pipeline = [(detector, soss_model), (world, None) ] return pipeline
def test_slicing_on_instance(): """ Test slicing a simple compound model class using integers. """ A = Const1D.rename('A') B = Const1D.rename('B') C = Const1D.rename('C') D = Const1D.rename('D') E = Const1D.rename('E') F = Const1D.rename('F') M = A + B - C * D / E ** F m = M(1, 2, 3, 4, 5, 6) assert isinstance(m[0:1], A) assert isinstance(m['A':'A'], A) assert isinstance(m[5:6], F) assert isinstance(m['F':'F'], F) # 1 + 2 assert m[:'B'](0) == 3 assert m[:'B'].param_names == ('amplitude_0', 'amplitude_1') assert np.all(m[:'B'].parameters == [1, 2]) # 2 - 3 assert m['B':'C'](0) == -1 assert m['B':'C'].param_names == ('amplitude_1', 'amplitude_2') assert np.all(m['B':'C'].parameters == [2, 3]) # 3 * 4 assert m['C':'D'](0) == 12 assert m['C':'D'].param_names == ('amplitude_2', 'amplitude_3') assert np.all(m['C':'D'].parameters == [3, 4]) # 4 / 5 assert m['D':'E'](0) == 0.8 assert m['D':'E'].param_names == ('amplitude_3', 'amplitude_4') assert np.all(m['D':'E'].parameters == [4, 5]) # 5 ** 6 assert m['E':'F'](0) == 15625 assert m['E':'F'].param_names == ('amplitude_4', 'amplitude_5') assert np.all(m['E':'F'].parameters == [5, 6])
def test_slicing_on_instance(): """ Test slicing a simple compound model class using integers. """ A = Const1D.rename('A') B = Const1D.rename('B') C = Const1D.rename('C') D = Const1D.rename('D') E = Const1D.rename('E') F = Const1D.rename('F') M = A + B - C * D / E**F m = M(1, 2, 3, 4, 5, 6) assert isinstance(m[0:1], A) assert isinstance(m['A':'A'], A) assert isinstance(m[5:6], F) assert isinstance(m['F':'F'], F) # 1 + 2 assert m[:'B'](0) == 3 assert m[:'B'].param_names == ('amplitude_0', 'amplitude_1') assert np.all(m[:'B'].parameters == [1, 2]) # 2 - 3 assert m['B':'C'](0) == -1 assert m['B':'C'].param_names == ('amplitude_1', 'amplitude_2') assert np.all(m['B':'C'].parameters == [2, 3]) # 3 * 4 assert m['C':'D'](0) == 12 assert m['C':'D'].param_names == ('amplitude_2', 'amplitude_3') assert np.all(m['C':'D'].parameters == [3, 4]) # 4 / 5 assert m['D':'E'](0) == 0.8 assert m['D':'E'].param_names == ('amplitude_3', 'amplitude_4') assert np.all(m['D':'E'].parameters == [4, 5]) # 5 ** 6 assert m['E':'F'](0) == 15625 assert m['E':'F'].param_names == ('amplitude_4', 'amplitude_5') assert np.all(m['E':'F'].parameters == [5, 6])
def centroid_1dg(data, error=None, mask=None): """ Calculate the centroid of a 2D array by fitting 1D Gaussians to the marginal ``x`` and ``y`` distributions of the array. Invalid values (e.g. NaNs or infs) in the ``data`` or ``error`` arrays are automatically masked. The mask for invalid values represents the combination of the invalid-value masks for the ``data`` and ``error`` arrays. Parameters ---------- data : array_like The 2D data array. error : array_like, optional The 2D array of the 1-sigma errors of the input ``data``. mask : array_like (bool), optional A boolean mask, with the same shape as ``data``, where a `True` value indicates the corresponding element of ``data`` is masked. Returns ------- centroid : `~numpy.ndarray` The ``x, y`` coordinates of the centroid. """ data = np.ma.asanyarray(data) if mask is not None and mask is not np.ma.nomask: mask = np.asanyarray(mask) if data.shape != mask.shape: raise ValueError('data and mask must have the same shape.') data.mask |= mask if np.any(~np.isfinite(data)): data = np.ma.masked_invalid(data) warnings.warn( 'Input data contains input values (e.g. NaNs or infs), ' 'which were automatically masked.', AstropyUserWarning) if error is not None: error = np.ma.masked_invalid(error) if data.shape != error.shape: raise ValueError('data and error must have the same shape.') data.mask |= error.mask error.mask = data.mask xy_error = np.array( [np.sqrt(np.ma.sum(error**2, axis=i)) for i in [0, 1]]) xy_weights = [(1.0 / xy_error[i].clip(min=1.e-30)) for i in [0, 1]] else: xy_weights = [np.ones(data.shape[i]) for i in [1, 0]] # assign zero weight to masked pixels if data.mask is not np.ma.nomask: bad_idx = [np.all(data.mask, axis=i) for i in [0, 1]] for i in [0, 1]: xy_weights[i][bad_idx[i]] = 0. xy_data = np.array([np.ma.sum(data, axis=i) for i in [0, 1]]) constant_init = np.ma.min(data) centroid = [] for (data_i, weights_i) in zip(xy_data, xy_weights): params_init = gaussian1d_moments(data_i) g_init = Const1D(constant_init) + Gaussian1D(*params_init) fitter = LevMarLSQFitter() x = np.arange(data_i.size) g_fit = fitter(g_init, x, data_i, weights=weights_i) centroid.append(g_fit.mean_1.value) return np.array(centroid)
def extract_grism_objects(input_model, grism_objects=None, reference_files=None, extract_orders=None, mmag_extract=99., compute_wavelength=True, wfss_extract_half_height=None): """ Extract 2d boxes around each objects spectra for each order. Parameters ---------- input_model : `~jwst.datamodels.ImageModel` An instance of an ImageModel, this is the grism image grism_objects : list(GrismObject) A list of GrismObjects reference_files : dict Needs to include the name of the wavelengthrange reference file extract_orders : int Spectral orders to extract mmag_extract : float Sources with magnitudes fainter than this minimum magnitude extraction cutoff will not be extracted compute_wavelength : bool Compute a wavelength array for the datamodel. wfss_extract_half_height : int, (optional) Cross-dispersion extraction half height in pixels, WFSS mode. Overwrites the computed extraction height. Returns ------- output_model : `~jwst.datamodels.MultiSlitModel` Notes ----- This method supports NRC_WFSS and NIS_WFSS only GrismObject is a named tuple which contains distilled information about each catalog object. It can be created by calling jwst.assign_wcs.util.create_grism_bbox() which will return a list of GrismObjects that countain the bounding boxes that will be used to define the 2d extraction area. For each spectral order, the configuration file contains a magnitude-cutoff value. Sources with magnitudes fainter than the extraction cutoff (MMAG_EXTRACT) will not be extracted, but are accounted for when computing the spectral contamination and background estimates. The default extraction value is 99 right now. The sensitivity information from the original aXe style configuration file needs to be modified by the passband of the filter used for the direct image to get the min and max wavelengths which correspond to t=0 and t=1, this currently has been done by the team and the min and max wavelengths to use to calculate t are stored in the grism reference file as wavelengthrange. Step 1: Convert the source catalog from the reference frame of the uberimage to that of the dispersed image. For the Vanilla Pipeline we assume that the pointing information in the file headers is sufficient. This will be strictly true if all images were obtained in a single visit (same guide stars). Step 2: Record source information for each object in the catalog: position (RA and Dec), shape (A_IMAGE, B_IMAGE, THETA_IMAGE), and all available magnitudes, and minimum bounding boxes Step 3: Compute the trace and wavelength solutions for each object in the catalog and for each spectral order. Record this information. Step 4: Compute the WIDTH of each spectral subwindow, which may be fixed or variable. The cross-dispersion size is taken from the minimum bounding box. """ if reference_files is None or not reference_files: raise TypeError("Expected a dictionary for reference_files") if grism_objects is None: # get the wavelengthrange reference file from the input_model if ('wavelengthrange' not in reference_files or reference_files['wavelengthrange'] in ['N/A', '']): raise ValueError("Expected name of wavelengthrange reference file") else: grism_objects = util.create_grism_bbox(input_model, reference_files, extract_orders=extract_orders, mmag_extract=mmag_extract, wfss_extract_half_height=wfss_extract_half_height) log.info("Grism object list created from source catalog: {0:s}" .format(input_model.meta.source_catalog)) if not isinstance(grism_objects, list): raise TypeError("Expected input grism objects to be a list") if len(grism_objects) == 0: raise ValueError("No grism objects created from source catalog") log.info("Extracting grism objects into MultiSlitModel") output_model = datamodels.MultiSlitModel() output_model.update(input_model) # One WCS model can be used to govern all the extractions # and in fact the model transforms rely on the full frame # coordinates of the input pixel location. So the WCS # attached to the extraction is just a copy of the # input_model WCS with a shift transform to the corner # of the subarray. They also depend on the source object # center, this information will be saved to the meta of # the output model as source_[x/y]pos inwcs = input_model.meta.wcs # For easy reference here, GrismObjects has: # # xcenter,ycenter: in direct image pixels # order_bounding in grism_detector pixels # sky_centroid: SkyCoord of object center # sky_bbox_ :lower and upper bounding box in SkyCoord # sid: catalog ID of the object slits = [] for obj in grism_objects: for order in obj.order_bounding.keys(): # Add the shift to the lower corner to each subarray WCS object # The shift should just be the lower bounding box corner # also replace the object center location inputs to the GrismDispersion # model with the known object center and order information (in pixels of direct image) # This is changes the user input to the model from (x,y,x0,y0,order) -> (x,y) # # The bounding boxes here are also limited to the size of the detector # The check for boxes entirely off the detector is done in create_grism_bbox right now y, x = obj.order_bounding[order] # limit the boxes to the detector ymin = clamp(y[0], 0, input_model.meta.subarray.ysize) ymax = clamp(y[1], 0, input_model.meta.subarray.ysize) xmin = clamp(x[0], 0, input_model.meta.subarray.xsize) xmax = clamp(x[1], 0, input_model.meta.subarray.xsize) # don't extract anything that ended up with zero dimensions in one axis # this means that it was identified as a partial order but only on one # row or column of the detector if (((ymax - ymin) > 0) and ((xmax - xmin) > 0)): subwcs = copy.deepcopy(inwcs) log.info("Subarray extracted for obj: {} order: {}:".format(obj.sid, order)) log.info("Subarray extents are: " "(xmin:{}, xmax:{}), (ymin:{}, ymax:{})".format(xmin, xmax, ymin, ymax)) # only the first two numbers in the Mapping are used # the order and source position are put directly into # the new wcs for the subarray for the forward transform xcenter_model = Const1D(obj.xcentroid) xcenter_model.inverse = Const1D(obj.xcentroid) ycenter_model = Const1D(obj.ycentroid) ycenter_model.inverse = Const1D(obj.ycentroid) order_model = Const1D(order) order_model.inverse = Const1D(order) tr = inwcs.get_transform('grism_detector', 'detector') tr = Mapping((0, 1, 0, 0, 0)) | (Shift(xmin) & Shift(ymin) & xcenter_model & ycenter_model & order_model) | tr ext_data = input_model.data[ymin: ymax + 1, xmin: xmax + 1].copy() ext_err = input_model.err[ymin: ymax + 1, xmin: xmax + 1].copy() ext_dq = input_model.dq[ymin: ymax + 1, xmin: xmax + 1].copy() if input_model.var_poisson is not None and np.size(input_model.var_poisson) > 0: var_poisson = input_model.var_poisson[ymin:ymax + 1, xmin:xmax + 1].copy() else: var_poisson = None if input_model.var_rnoise is not None and np.size(input_model.var_rnoise) > 0: var_rnoise = input_model.var_rnoise[ymin:ymax + 1, xmin:xmax + 1].copy() else: var_rnoise = None if input_model.var_flat is not None and np.size(input_model.var_flat) > 0: var_flat = input_model.var_flat[ymin:ymax + 1, xmin:xmax + 1].copy() else: var_flat = None tr.bounding_box = util.transform_bbox_from_shape(ext_data.shape) subwcs.set_transform('grism_detector', 'detector', tr) new_slit = datamodels.SlitModel(data=ext_data, err=ext_err, dq=ext_dq, var_poisson=var_poisson, var_rnoise=var_rnoise, var_flat=var_flat) new_slit.meta.wcsinfo.spectral_order = order new_slit.meta.wcsinfo.dispersion_direction = \ input_model.meta.wcsinfo.dispersion_direction new_slit.meta.wcsinfo.specsys = input_model.meta.wcsinfo.specsys new_slit.meta.coordinates = input_model.meta.coordinates new_slit.meta.wcs = subwcs if compute_wavelength: log.debug("Computing wavelengths") new_slit.wavelength = compute_wavelength_array(new_slit) # set x/ystart values relative to the image (screen) frame. # The overall subarray offset is recorded in model.meta.subarray. # nslit = obj.sid - 1 # catalog id starts at zero new_slit.name = "{0}".format(obj.sid) new_slit.is_extended = obj.is_extended new_slit.xstart = xmin + 1 # fits pixels new_slit.xsize = ext_data.shape[1] new_slit.ystart = ymin + 1 # fits pixels new_slit.ysize = ext_data.shape[0] new_slit.source_xpos = float(obj.xcentroid) new_slit.source_ypos = float(obj.ycentroid) new_slit.source_id = obj.sid new_slit.bunit_data = input_model.meta.bunit_data new_slit.bunit_err = input_model.meta.bunit_err slits.append(new_slit) output_model.slits.extend(slits) # In the case that there are no spectra to extract deleting the variables # will fail so add the try block. try: del subwcs except UnboundLocalError: pass try: del new_slit except UnboundLocalError: pass # del subwcs # del new_slit log.info("Finished extractions") return output_model
def extract_tso_object(input_model, reference_files=None, tsgrism_extract_height=None, extract_orders=None, compute_wavelength=True): """ Extract the spectrum for a NIRCam TSO grism observation. Parameters ---------- input_model : `~jwst.datamodels.CubeModel` or `~jwst.datamodels.ImageModel` The input TSO data is an instance of a CubeModel (3D) or ImageModel (2D) reference_files : dict Needs to include the name of the wavelengthrange reference file tsgrism_extract_height : int The extraction height, in total, for the spectrum in the cross-dispersion direction. If this is other than None, it will override the team default of 64 pixels. The team wants the source centered near row 34, so the extraction height is not the same on either size of the central row. extract_orders : list[ints] This is an optional parameter that will override the orders specified for extraction in the wavelengthrange reference file. compute_wavelength : bool Compute a wavelength array for the datamodel. Returns ------- output_model : `~jwst.datamodels.SlitModel` Notes ----- This method supports NRC_TSGRISM only, where only one bright object is considered in the field, so there's no catalog of sources and the object is assumed to have been observed at the aperture reference position. The aperture reference location is read during level-1b (uncal) product creation by the "set_telescope_pointing" script from the SIAF entries XSciRef and YSciRef (reference location in the science frame) and saved as "meta.wcsinfo.siaf_xref_sci" and "meta.wcsinfo.siaf_yref_sci" (FITS header keywords XREF_SCI and YREF_SCI). Because this mode has a single known source location, the utilities used in the WFSS modes are overkill. Instead, similar structures are created during the extract2d process and then directly used. https://jwst-docs.stsci.edu/near-infrared-camera/nircam-observing-modes/nircam-time-series-observations/nircam-grism-time-series """ # Check for reference files if not isinstance(reference_files, dict): raise TypeError("Expected a dictionary for reference_files") # Check for wavelengthrange reference file if 'wavelengthrange' not in reference_files: raise KeyError("No wavelengthrange reference file specified") # If an extraction height is not supplied, default to entire # cross-dispersion size of the data array if tsgrism_extract_height is None: tsgrism_extract_height = input_model.meta.subarray.ysize log.info("Setting extraction height to {}".format(tsgrism_extract_height)) # Get the disperser parameters that have the wave limits with WavelengthrangeModel(reference_files['wavelengthrange']) as f: if (f.meta.instrument.name != 'NIRCAM' or f.meta.exposure.type != 'NRC_TSGRISM'): raise ValueError("Wavelengthrange reference file is not for NIRCAM TSGRISM mode!") wavelengthrange = f.wavelengthrange ref_extract_orders = f.extract_orders # If user-supplied spectral orders are not provided, # default to extracting only the 1st order if extract_orders is None: log.info("Using default order extraction from reference file") extract_orders = ref_extract_orders available_orders = [x[1] for x in extract_orders if x[0] == input_model.meta.instrument.filter].pop() else: if (not isinstance(extract_orders, list) or not all(isinstance(item, int) for item in extract_orders)): raise TypeError("Expected extract_orders to be a list of integers.") available_orders = extract_orders if len(available_orders) > 1: raise NotImplementedError("Multiple order extraction for TSO is " "not currently implemented.") # Check for the existence of the aperture reference location meta data if (input_model.meta.wcsinfo.siaf_xref_sci is None or input_model.meta.wcsinfo.siaf_yref_sci is None): raise ValueError('XREF_SCI and YREF_SCI are required for TSO mode.') # Create the extracted output as a SlitModel log.info("Extracting order: {}".format(available_orders)) output_model = datamodels.SlitModel() output_model.update(input_model) subwcs = copy.deepcopy(input_model.meta.wcs) # Loop over spectral orders for order in available_orders: range_select = [(x[2], x[3]) for x in wavelengthrange if (x[0] == order and x[1] == input_model.meta.instrument.filter)] # Use the filter that was in front of the grism for translation lmin, lmax = range_select.pop() # Create the order bounding box source_xpos = input_model.meta.wcsinfo.siaf_xref_sci - 1 # remove FITS 1-indexed offset source_ypos = input_model.meta.wcsinfo.siaf_yref_sci - 1 # remove FITS 1-indexed offset transform = input_model.meta.wcs.get_transform('direct_image', 'grism_detector') xmin, ymin, _ = transform(source_xpos, source_ypos, lmin, order) xmax, ymax, _ = transform(source_xpos, source_ypos, lmax, order) # Add the shift to the lower corner to the subarray WCS object. # The shift should just be the lower bounding box corner. # Also replace the object center location inputs to the GrismDispersion # model with the known object center and order information (in pixels of direct image) # This changes the user input to the model from (x,y,x0,y0,order) -> (x,y) # # The bounding box is limited to the size of the detector in the dispersion direction # and 64 pixels in the cross-dispersion direction (at request of instrument team). # # The team wants the object to fall near row 34 for all cutouts, but the default cutout # height is 64 pixels (32 on either side). So bump the extraction ycenter, when necessary, # so that the height is 30 above and 34 below (in full frame) the object center. bump = source_ypos - 34 extract_y_center = source_ypos - bump splitheight = int(tsgrism_extract_height / 2) below = extract_y_center - splitheight if below == 34: extract_y_min = 0 extract_y_max = extract_y_center + splitheight elif below < 0: extract_y_min = 0 extract_y_max = tsgrism_extract_height - 1 else: extract_y_min = extract_y_center - 34 # always return source at row 34 in cutout extract_y_max = extract_y_center + tsgrism_extract_height - 34 - 1 # Check for bad results if extract_y_min > extract_y_max: raise ValueError("Something bad happened calculating extraction y-size") # Limit the bounding box to the detector edges ymin, ymax = (max(extract_y_min, 0), min(extract_y_max, input_model.meta.subarray.ysize)) xmin, xmax = (max(xmin, 0), min(xmax, input_model.meta.subarray.xsize)) # The order and source position are put directly into the new WCS of the subarray # for the forward transform. # # NOTE NOTE NOTE 2020-02-14 # We would normally use x-axis (along dispersion) extraction limits calculated # above based on the min/max wavelength range and the source position to do the # subarray extraction and set the subarray WCS accordingly. HOWEVER, the NIRCam # team has asked for all data along the dispersion direction to be included in # subarray cutout, so here we override the xmin/xmax values calculated above and # instead hardwire the extraction limits for the x (dispersion) direction to # cover the entire range of the data and use this new minimum x value in the # subarray WCS transform. If the team ever decides to change the extraction limits, # the following two constants must be modified accordingly. xmin_ext = 0 # hardwire min x for extraction to zero xmax_ext = input_model.data.shape[-1] - 1 # hardwire max x for extraction to size of data order_model = Const1D(order) order_model.inverse = Const1D(order) tr = input_model.meta.wcs.get_transform('grism_detector', 'direct_image') tr = Mapping((0, 1, 0)) | Shift(xmin_ext) & Shift(ymin) & order_model | tr subwcs.set_transform('grism_detector', 'direct_image', tr) xmin = int(xmin) xmax = int(xmax) ymin = int(ymin) ymax = int(ymax) log.info("WCS made explicit for order: {}".format(order)) log.info("Spectral trace extents: (xmin: {}, ymin: {}), " "(xmax: {}, ymax: {})".format(xmin, ymin, xmax, ymax)) log.info("Extraction limits: (xmin: {}, ymin: {}), " "(xmax: {}, ymax: {})".format(xmin_ext, ymin, xmax_ext, ymax)) # Cut out the subarray from the input data arrays ext_data = input_model.data[..., ymin: ymax + 1, xmin_ext:xmax_ext + 1].copy() ext_err = input_model.err[..., ymin: ymax + 1, xmin_ext:xmax_ext + 1].copy() ext_dq = input_model.dq[..., ymin: ymax + 1, xmin_ext:xmax_ext + 1].copy() if input_model.var_poisson is not None and np.size(input_model.var_poisson) > 0: var_poisson = input_model.var_poisson[..., ymin:ymax + 1, xmin_ext:xmax_ext + 1].copy() else: var_poisson = None if input_model.var_rnoise is not None and np.size(input_model.var_rnoise) > 0: var_rnoise = input_model.var_rnoise[..., ymin:ymax + 1, xmin_ext:xmax_ext + 1].copy() else: var_rnoise = None if input_model.var_flat is not None and np.size(input_model.var_flat) > 0: var_flat = input_model.var_flat[..., ymin:ymax + 1, xmin_ext:xmax_ext + 1].copy() else: var_flat = None # Finish populating the output model and meta data if output_model.meta.model_type == "SlitModel": output_model.data = ext_data output_model.err = ext_err output_model.dq = ext_dq output_model.var_poisson = var_poisson output_model.var_rnoise = var_rnoise output_model.var_flat = var_flat output_model.meta.wcs = subwcs output_model.meta.wcs.bounding_box = util.wcs_bbox_from_shape(ext_data.shape) output_model.meta.wcsinfo.siaf_yref_sci = 34 # update for the move, vals are the same output_model.meta.wcsinfo.siaf_xref_sci = input_model.meta.wcsinfo.siaf_xref_sci output_model.meta.wcsinfo.spectral_order = order output_model.meta.wcsinfo.dispersion_direction = \ input_model.meta.wcsinfo.dispersion_direction if compute_wavelength: log.debug("Computing wavelengths (this takes a while ...)") output_model.wavelength = compute_wavelength_array(output_model) output_model.name = '1' output_model.source_type = input_model.meta.target.source_type output_model.source_name = input_model.meta.target.catalog_name output_model.source_alias = input_model.meta.target.proposer_name output_model.xstart = 1 # FITS pixels are 1-indexed output_model.xsize = ext_data.shape[-1] output_model.ystart = ymin + 1 # FITS pixels are 1-indexed output_model.ysize = ext_data.shape[-2] output_model.source_xpos = source_xpos output_model.source_ypos = 34 output_model.source_id = 1 output_model.bunit_data = input_model.meta.bunit_data output_model.bunit_err = input_model.meta.bunit_err if hasattr(input_model, 'int_times'): output_model.int_times = input_model.int_times.copy() del subwcs log.info("Finished extraction") return output_model