def test_invalid_operands():
"""
Test that certain operators do not work with models whose inputs/outputs do
not match up correctly.
"""
with pytest.raises(ModelDefinitionError):
Rotation2D(90) | Gaussian1D(1, 0, 0.1)
with pytest.raises(ModelDefinitionError):
Rotation2D(90) + Gaussian1D(1, 0, 0.1)
def test_replace_submodel():
"""
Replace a model in a Compound model
"""
S1 = Shift(2, name='shift2') | Scale(
3, name='scale3') # First shift then scale
S2 = Scale(2, name='scale2') | Shift(
3, name='shift3') # First scale then shift
m = S1 & S2
assert m(1, 2) == (9, 7)
m2 = m.replace_submodel('scale3', Scale(4, name='scale4'))
assert m2(1, 2) == (12, 7)
assert m(1, 2) == (9, 7)
# Check the inverse has been updated
assert m2.inverse(12, 7) == (1, 2)
# Produce the same result by replacing a single model with a compound
m3 = m.replace_submodel('shift2', Shift(2) | Scale(2))
assert m(1, 2) == (9, 7)
assert m3(1, 2) == (18, 7)
# Check the inverse has been updated
assert m3.inverse(18, 7) == (1, 2)
# Test with arithmetic model compunding operator
m = S1 + S2
assert m(1) == 14
m2 = m.replace_submodel('scale2', Scale(4, name='scale4'))
assert m2(1) == 16
# Test with fix_inputs()
R = fix_inputs(Rotation2D(angle=90, name='rotate'), {0: 1})
m4 = S1 | R
assert_allclose(m4(0), (-6, 1))
m5 = m4.replace_submodel('rotate', Rotation2D(180))
assert_allclose(m5(0), (-1, -6))
# Check we get a value error when model name doesn't exist
with pytest.raises(ValueError):
m2 = m.replace_submodel('not_there', Scale(2))
# And now a model set
P = Polynomial1D(degree=1, n_models=2, name='poly')
S = Shift([1, 2], n_models=2)
m = P | S
assert_array_equal(m([0, 1]), (1, 2))
with pytest.raises(ValueError):
m2 = m.replace_submodel('poly', Polynomial1D(degree=1, c0=1))
m2 = m.replace_submodel('poly',
Polynomial1D(degree=1, c0=[1, 2], n_models=2))
assert_array_equal(m2([0, 1]), (2, 4))
def test_mapping_inverse():
"""Tests inverting a compound model that includes a `Mapping`."""
rs1 = Rotation2D(12.1) & Scale(13.2)
rs2 = Rotation2D(14.3) & Scale(15.4)
# Rotates 2 of the coordinates and scales the third--then rotates on a
# different axis and scales on the axis of rotation. No physical meaning
# here just a simple test
m = rs1 | Mapping([2, 0, 1]) | rs2
assert_allclose((0, 1, 2), m.inverse(*m(0, 1, 2)), atol=1e-08)
def test_parameters_compound_models():
tan = Pix2Sky_TAN()
sky_coords = coord.SkyCoord(ra=5.6, dec=-72, unit=u.deg)
lon_pole = 180 * u.deg
n2c = RotateNative2Celestial(sky_coords.ra, sky_coords.dec, lon_pole)
rot = Rotation2D(23)
m = rot | n2c
def compute_spec_transform(fiducial, refwcs):
"""
Compute a simple transform given a fidicial point in a spatial-spectral wcs.
"""
cdelt1 = refwcs.wcsinfo.cdelt1 / 3600.
cdelt2 = refwcs.wcsinfo.cdelt2 / 3600.
cdelt3 = refwcs.wcsinfo.cdelt3
roll_ref = refwcs.wcsinfo.roll_ref
y, x = grid_from_spec_domain(refwcs)
ra, dec, lam = refwcs(x, y)
min_lam = np.nanmin(lam)
offset = Shift(0.) & Shift(0.)
rot = Rotation2D(roll_ref)
scale = Scale(cdelt1) & Scale(cdelt2)
tan = Pix2Sky_TAN()
skyrot = RotateNative2Celestial(fiducial[0][0], fiducial[0][1], 180.0)
spatial = offset | rot | scale | tan | skyrot
spectral = Scale(cdelt3) | Shift(min_lam)
mapping = Mapping((1, 1, 0), )
mapping.inverse = Mapping((2, 1))
transform = mapping | spatial & spectral
transform.outputs = ('ra', 'dec', 'lamda')
return transform
def test_basic_compound_inverse():
"""
Test basic inversion of compound models in the limited sense supported for
models made from compositions and joins only.
"""
t = (Shift(2) & Shift(3)) | (Scale(2) & Scale(3)) | Rotation2D(90)
assert_allclose(t.inverse(*t(0, 1)), (0, 1))
def test_mapping_basic_permutations():
"""
Tests a couple basic examples of the Mapping model--specifically examples
that merely permute the outputs.
"""
x, y = Rotation2D(90)(1, 2)
rs = Rotation2D(90) | Mapping((1, 0))
x_prime, y_prime = rs(1, 2)
assert_allclose((x, y), (y_prime, x_prime))
# A more complicated permutation
m = Rotation2D(90) & Scale(2)
x, y, z = m(1, 2, 3)
ms = m | Mapping((2, 0, 1))
x_prime, y_prime, z_prime = ms(1, 2, 3)
assert_allclose((x, y, z), (y_prime, z_prime, x_prime))
def test_simple_two_model_compose_2d():
"""
A simple example consisting of two rotations.
"""
r1 = Rotation2D(45) | Rotation2D(45)
assert isinstance(r1, CompoundModel)
assert r1.n_inputs == 2
assert r1.n_outputs == 2
assert_allclose(r1(0, 1), (-1, 0), atol=1e-10)
r2 = Rotation2D(90) | Rotation2D(90) # Rotate twice by 90 degrees
assert_allclose(r2(0, 1), (0, -1), atol=1e-10)
# Compose R with itself to produce 4 rotations
r3 = r1 | r1
assert_allclose(r3(0, 1), (0, -1), atol=1e-10)
def evaluate(self, x, y, wavelength, order):
"""Return the dispersed pixel(s) given center x, y, lam and order
Parameters
----------
x : int,float
Input x location on the direct image
y : int,float
Input y location on the direct image
wavelength : float
Wavelength to disperse
order : list
The order to use
Returns
-------
x, y in the grism image for the pixel at x0, y0 that was
specified as input using the wavelength and order specified
Notes
-----
I kept the potential for rotation from NIRISS, unsure if it's actually
needed/useful for WFC3. Original note:
There's spatial dependence for NIRISS so the forward transform
dependes on x,y as well as the filter wheel rotation. Theta is
usu. taken to be the different between fwcpos_ref in the specwcs
reference file and fwcpos from the input image.
"""
if wavelength < 0:
raise ValueError("Wavelength should be greater than zero")
try:
iorder = self._order_mapping[int(order.flatten()[0])]
except AttributeError:
iorder = self._order_mapping[order]
except KeyError:
raise ValueError("Specified order is not available")
t = self.lmodels[iorder](wavelength)
xmodel = self.xmodels[iorder]
ymodel = self.ymodels[iorder]
dx = xmodel.evaluate(x, y, t)
dy = ymodel.evaluate(x, y, t)
## rotate by theta
if self.theta != 0.0:
rotate = Rotation2D(self.theta)
dx, dy = rotate(dx, dy)
return (x + dx, y + dy, x, y, order)
def test_identity_input():
"""
Test a case where an Identity (or Mapping) model is the first in a chain
of composite models and thus is responsible for handling input broadcasting
properly.
Regression test for https://github.com/astropy/astropy/pull/3362
"""
ident1 = Identity(1)
shift = Shift(1)
rotation = Rotation2D(angle=90)
model = ident1 & shift | rotation
assert_allclose(model(1, 2), [-3.0, 1.0])
def test_identity():
x = np.zeros((2, 3))
y = np.ones((2, 3))
ident1 = Identity(1)
shift = Shift(1)
rotation = Rotation2D(angle=60)
model = ident1 & shift | rotation
assert_allclose(model(1, 2), (-2.098076211353316, 2.3660254037844393))
res_x, res_y = model(x, y)
assert_allclose((res_x, res_y),
(np.array([[-1.73205081, -1.73205081, -1.73205081],
[-1.73205081, -1.73205081, -1.73205081]
]), np.array([[1., 1., 1.], [1., 1., 1.]])))
assert_allclose(model.inverse(res_x, res_y), (x, y), atol=1.e-10)
def test_levmar2x2_multivariate():
inputs = [np.array([10., 10., 20., 20.]), np.array([10., 20., 20., 10.])]
outputs = [
np.array([8.06101731, 0.98994949, 8.06101731, 15.13208512]),
np.array([12.16223664, 19.23330445, 26.30437226, 19.23330445])
]
rot = Rotation2D()
rot.fittable = True
model = (Shift() & Shift()) | rot
fitter = linearfit._LevMarLSQFitter2x2()
finfo = fitter(model, inputs, outputs)
assert np.allclose(finfo.parameters,
np.array([4.3, -7.1, 45.]),
rtol=1e-5,
atol=1e-5)
def test_tabular_in_compound():
"""
Issue #7411 - evaluate should not change the shape of the output.
"""
t = Tabular1D(points=([1, 5, 7],), lookup_table=[12, 15, 19],
bounds_error=False)
rot = Rotation2D(2)
p = Polynomial1D(1)
x = np.arange(12).reshape((3, 4))
# Create a compound model which does ot execute Tabular.__call__,
# but model.evaluate and is followed by a Rotation2D which
# checks the exact shapes.
model = p & t | rot
x1, y1 = model(x, x)
assert x1.ndim == 2
assert y1.ndim == 2
def test_slicing_on_instances_2():
"""
More slicing tests.
Regression test for https://github.com/embray/astropy/pull/10
"""
model_a = Shift(1, name='a')
model_b = Shift(2, name='b')
model_c = Rotation2D(3, name='c')
model_d = Scale(2, name='d')
model_e = Scale(3, name='e')
m = (model_a & model_b) | model_c | (model_d & model_e)
with pytest.raises(ModelDefinitionError):
# The slice can't actually be taken since the resulting model cannot be
# evaluated
assert m[1:].submodel_names == ('b', 'c', 'd', 'e')
assert m[:].submodel_names == ('a', 'b', 'c', 'd', 'e')
assert m['a':].submodel_names == ('a', 'b', 'c', 'd', 'e')
with pytest.raises(ModelDefinitionError):
assert m['c':'d'].submodel_names == ('c', 'd')
assert m[1:2].name == 'b'
assert m[2:7].submodel_names == ('c', 'd', 'e')
with pytest.raises(IndexError):
m['x']
with pytest.raises(IndexError):
m['a':'r']
with pytest.raises(ModelDefinitionError):
assert m[-4:4].submodel_names == ('b', 'c', 'd')
with pytest.raises(ModelDefinitionError):
assert m[-4:-2].submodel_names == ('b', 'c')
def build_miri_output_wcs(self, refwcs=None):
"""
Create a simple output wcs covering footprint of the input datamodels
"""
# TODO: generalize this for more than one input datamodel
# TODO: generalize this for imaging modes with distorted wcs
input_model = self.input_models[0]
if refwcs == None:
refwcs = input_model.meta.wcs
x, y = wcstools.grid_from_bounding_box(refwcs.bounding_box,
step=(1, 1),
center=True)
ra, dec, lam = refwcs(x.flatten(), y.flatten())
# TODO: once astropy.modeling._Tabular is fixed, take out the
# flatten() and reshape() code above and below
ra = ra.reshape(x.shape)
dec = dec.reshape(x.shape)
lam = lam.reshape(x.shape)
# Find rotation of the slit from y axis from the wcs forward transform
# TODO: figure out if angle is necessary for MIRI. See for discussion
# https://github.com/STScI-JWST/jwst/pull/347
rotation = [m for m in refwcs.forward_transform if \
isinstance(m, Rotation2D)]
if rotation:
rot_slit = functools.reduce(lambda x, y: x | y, rotation)
rot_angle = rot_slit.inverse.angle.value
unrotate = rot_slit.inverse
refwcs_minus_rot = refwcs.forward_transform | \
unrotate & Identity(1)
# Correct for this rotation in the wcs
ra, dec, lam = refwcs_minus_rot(x.flatten(), y.flatten())
ra = ra.reshape(x.shape)
dec = dec.reshape(x.shape)
lam = lam.reshape(x.shape)
# Get the slit size at the center of the dispersion
sky_coords = SkyCoord(ra=ra, dec=dec, unit=u.deg)
slit_coords = sky_coords[int(sky_coords.shape[0] / 2)]
slit_angular_size = slit_coords[0].separation(slit_coords[-1])
log.debug('Slit angular size: {0}'.format(slit_angular_size.arcsec))
# Compute slit center from bounding_box
dx0 = refwcs.bounding_box[0][0]
dx1 = refwcs.bounding_box[0][1]
dy0 = refwcs.bounding_box[1][0]
dy1 = refwcs.bounding_box[1][1]
slit_center_pix = (dx1 - dx0) / 2
dispersion_center_pix = (dy1 - dy0) / 2
slit_center = refwcs_minus_rot(dx0 + slit_center_pix,
dy0 + dispersion_center_pix)
slit_center_sky = SkyCoord(ra=slit_center[0],
dec=slit_center[1],
unit=u.deg)
log.debug('slit center: {0}'.format(slit_center))
# Compute spatial and spectral scales
spatial_scale = slit_angular_size / slit_coords.shape[0]
log.debug('Spatial scale: {0}'.format(spatial_scale.arcsec))
tcenter = int((dx1 - dx0) / 2)
trace = lam[:, tcenter]
trace = trace[~np.isnan(trace)]
spectral_scale = np.abs((trace[-1] - trace[0]) / trace.shape[0])
log.debug('spectral scale: {0}'.format(spectral_scale))
# Compute transform for output frame
log.debug('Slit center %s' % slit_center_pix)
offset = Shift(-slit_center_pix) & Shift(-slit_center_pix)
# TODO: double-check the signs on the following rotation angles
roll_ref = input_model.meta.wcsinfo.roll_ref * u.deg
rot = Rotation2D(roll_ref)
tan = Pix2Sky_TAN()
lon_pole = _compute_lon_pole(slit_center_sky, tan)
skyrot = RotateNative2Celestial(slit_center_sky.ra,
slit_center_sky.dec, lon_pole)
min_lam = np.nanmin(lam)
mapping = Mapping((0, 0, 1))
transform = Shift(-slit_center_pix) & Identity(1) | \
Scale(spatial_scale) & Scale(spectral_scale) | \
Identity(1) & Shift(min_lam) | mapping | \
(rot | tan | skyrot) & Identity(1)
transform.inputs = (x, y)
transform.outputs = ('ra', 'dec', 'lamda')
# Build the output wcs
input_frame = refwcs.input_frame
output_frame = refwcs.output_frame
wnew = WCS(output_frame=output_frame, forward_transform=transform)
# Build the bounding_box in the output frame wcs object
bounding_box_grid = wnew.backward_transform(ra, dec, lam)
bounding_box = []
for axis in input_frame.axes_order:
axis_min = np.nanmin(bounding_box_grid[axis])
axis_max = np.nanmax(bounding_box_grid[axis])
bounding_box.append((axis_min, axis_max))
wnew.bounding_box = tuple(bounding_box)
# Update class properties
self.output_spatial_scale = spatial_scale
self.output_spectral_scale = spectral_scale
self.output_wcs = wnew
def build_nirspec_output_wcs(self, refwcs=None):
"""
Create a simple output wcs covering footprint of the input datamodels
"""
# TODO: generalize this for more than one input datamodel
# TODO: generalize this for imaging modes with distorted wcs
input_model = self.input_models[0]
if refwcs == None:
refwcs = input_model.meta.wcs
# Generate grid of sky coordinates for area within bounding box
bb = refwcs.bounding_box
det = x, y = wcstools.grid_from_bounding_box(bb,
step=(1, 1),
center=True)
sky = ra, dec, lam = refwcs(*det)
x_center = int((bb[0][1] - bb[0][0]) / 2)
y_center = int((bb[1][1] - bb[1][0]) / 2)
log.debug("Center of bounding box: {} {}".format(x_center, y_center))
# Compute slit angular size, slit center sky coords
xpos = []
sz = 3
for row in lam:
if np.isnan(row[x_center]):
xpos.append(np.nan)
else:
f = interpolate.interp1d(row[x_center - sz + 1:x_center + sz],
x[y_center,
x_center - sz + 1:x_center + sz],
bounds_error=False,
fill_value='extrapolate')
xpos.append(f(lam[y_center, x_center]))
x_arg = np.array(xpos)[~np.isnan(lam[:, x_center])]
y_arg = y[~np.isnan(lam[:, x_center]), x_center]
# slit_coords, spect0 = refwcs(x_arg, y_arg, output='numericals_plus')
slit_ra, slit_dec, slit_spec_ref = refwcs(x_arg, y_arg)
slit_coords = SkyCoord(ra=slit_ra, dec=slit_dec, unit=u.deg)
pix_num = np.flipud(np.arange(len(slit_ra)))
# pix_num = np.arange(len(slit_ra))
interpol_ra = interpolate.interp1d(pix_num, slit_ra)
interpol_dec = interpolate.interp1d(pix_num, slit_dec)
slit_center_pix = len(slit_spec_ref) / 2. - 1
log.debug('Slit center pix: {0}'.format(slit_center_pix))
slit_center_sky = SkyCoord(ra=interpol_ra(slit_center_pix),
dec=interpol_dec(slit_center_pix),
unit=u.deg)
log.debug('Slit center: {0}'.format(slit_center_sky))
log.debug('Fiducial: {0}'.format(
resample_utils.compute_spec_fiducial([refwcs])))
angular_slit_size = np.abs(slit_coords[0].separation(slit_coords[-1]))
log.debug('Slit angular size: {0}'.format(angular_slit_size.arcsec))
dra, ddec = slit_coords[0].spherical_offsets_to(slit_coords[-1])
offset_up_slit = (dra.to(u.arcsec), ddec.to(u.arcsec))
log.debug('Offset up the slit: {0}'.format(offset_up_slit))
# Compute spatial and spectral scales
xposn = np.array(xpos)[~np.isnan(xpos)]
dx = xposn[-1] - xposn[0]
slit_npix = np.sqrt(dx**2 + np.array(len(xposn) - 1)**2)
spatial_scale = angular_slit_size / slit_npix
log.debug('Spatial scale: {0}'.format(spatial_scale.arcsec))
spectral_scale = lam[y_center, x_center] - lam[y_center, x_center - 1]
# Compute slit angle relative (clockwise) to y axis
slit_rot_angle = (np.arcsin(dx / slit_npix) * u.radian).to(u.degree)
slit_rot_angle = slit_rot_angle.value
log.debug('Slit rotation angle: {0}'.format(slit_rot_angle))
# Compute transform for output frame
roll_ref = input_model.meta.wcsinfo.roll_ref
min_lam = np.nanmin(lam)
offset = Shift(-slit_center_pix) & Shift(-slit_center_pix)
# TODO: double-check the signs on the following rotation angles
rot = Rotation2D(roll_ref + slit_rot_angle)
scale = Scale(spatial_scale.value) & Scale(spatial_scale.value)
tan = Pix2Sky_TAN()
lon_pole = _compute_lon_pole(slit_center_sky, tan)
skyrot = RotateNative2Celestial(slit_center_sky.ra.value,
slit_center_sky.dec.value,
lon_pole.value)
spatial_trans = offset | rot | scale | tan | skyrot
spectral_trans = Scale(spectral_scale) | Shift(min_lam)
mapping = Mapping((1, 1, 0))
mapping.inverse = Mapping((2, 1))
transform = mapping | spatial_trans & spectral_trans
transform.outputs = ('ra', 'dec', 'lamda')
# Build the output wcs
input_frame = refwcs.input_frame
output_frame = refwcs.output_frame
wnew = WCS(output_frame=output_frame, forward_transform=transform)
# Build the bounding_box in the output frame wcs object
bounding_box_grid = wnew.backward_transform(ra, dec, lam)
bounding_box = []
for axis in input_frame.axes_order:
axis_min = np.nanmin(bounding_box_grid[axis])
axis_max = np.nanmax(bounding_box_grid[axis])
bounding_box.append((axis_min, axis_max))
wnew.bounding_box = tuple(bounding_box)
# Update class properties
self.output_spatial_scale = spatial_scale
self.output_spectral_scale = spectral_scale
self.output_wcs = wnew
def test_drop_axes_3():
mapping = Mapping((1, ), n_inputs=2)
assert mapping.n_inputs == 2
rotation = Rotation2D(60)
model = rotation | mapping
assert_allclose(model(1, 2), 1.86602540378)
# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
import numpy as np
from astropy.modeling.models import Shift, Rotation2D
from asdf.tests import helpers
from ...import jwextension
from ...models import (AngleFromGratingEquation, WavelengthFromGratingEquation,
Unitless2DirCos, DirCos2Unitless, Rotation3DToGWA, Gwa2Slit,
Snell, Logical, V23ToSky, Slit)
import pytest
m1 = Shift(1) & Shift(2) | Rotation2D(3.1)
m2 = Shift(2) & Shift(2) | Rotation2D(23.1)
test_models = [DirCos2Unitless(), Unitless2DirCos(),
Rotation3DToGWA(angles=[12.1, 1.3, 0.5, 3.4], axes_order='xyzx'),
AngleFromGratingEquation(20000, -1), WavelengthFromGratingEquation(25000, 2),
Logical('GT', 5, 10), Logical('LT', np.ones((10,))* 5, np.arange(10)),
V23ToSky(angles=[0.1259, -0.1037, -146.03468, 69.503032, 80.46448], axes_order="zyxyz"),
Snell(angle=-16.5, kcoef=[0.583, 0.462, 3.891], lcoef=[0.002526, 0.01, 1200.556],
tcoef=[-2.66e-05, 0.0, 0.0, 0.0, 0.0, 0.0], tref=35, pref=0,
temperature=35, pressure=0),
]
@pytest.mark.parametrize(('model'), test_models)
def test_model(tmpdir, model):
tree = {'model': model}
helpers.assert_roundtrip_tree(tree, tmpdir, extensions=jwextension.JWSTExtension())
def evaluate(self, x, y, x0, y0, order):
"""Return the valid pixel(s) and wavelengths given x0, y0, x, y, order
Parameters
----------
x0: int,float,list
Source object x-center
y0: int,float,list
Source object y-center
x : int,float,list
Input x location
y : int,float,list
Input y location
order : int
Spectral order to use
Returns
-------
x0, y0, lambda, order in the direct image for the pixel that was
specified as input using the wavelength l and spectral order
Notes
-----
I kept the possibility of having a rotation like NIRISS, although I
don't know if there is a use case for it for WFC3.
The two `flatten` lines may need to be uncommented if we want to use
this for array input.
"""
try:
iorder = self._order_mapping[int(order.flatten()[0])]
except AttributeError:
iorder = self._order_mapping[order]
except KeyError:
raise ValueError("Specified order is not available")
# The next two lines are to get around the fact that
# modeling.standard_broadcasting=False does not work.
#x00 = x0.flatten()[0]
#y00 = y0.flatten()[0]
t = np.linspace(0, 1, 10) #sample t
xmodel = self.xmodels[iorder]
ymodel = self.ymodels[iorder]
lmodel = self.lmodels[iorder]
dx = xmodel.evaluate(x0, y0, t)
dy = ymodel.evaluate(x0, y0, t)
if self.theta != 0.0:
rotate = Rotation2D(self.theta)
dx, dy = rotate(dx, dy)
so = np.argsort(dx)
tab = Tabular1D(dx[so], t[so], bounds_error=False, fill_value=None)
dxr = astmath.SubtractUfunc()
wavelength = dxr | tab | lmodel
model = Mapping(
(2, 3, 0, 2,
4)) | Const1D(x0) & Const1D(y0) & wavelength & Const1D(order)
return model(x, y, x0, y0, order)
