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