from copy import deepcopy
from astropy.modeling.rotations import RotationSequence3D
from astropy import units
from gwcs.wcs import WCS
import numpy as np
from tweakwcs.correctors import JWSTWCSCorrector
from tweakwcs.linearfit import build_fit_matrix
from stdatamodels.jwst.datamodels import ImageModel
from ..assign_wcs.util import update_fits_wcsinfo
from ..assign_wcs.pointing import _v23tosky
_RAD2ARCSEC = 3600.0 * np.rad2deg(1.0)
__all__ = ["adjust_wcs", "transfer_wcs_correction"]
def _wcsinfo_from_wcs_transform(wcs):
frames = wcs.available_frames
if 'v2v3' not in frames or 'world' not in frames or frames[-1] != 'world':
raise ValueError(
"Unsupported WCS structure."
)
# Initially get v2_ref, v3_ref, and roll_ref from
# the v2v3 to world transform. Also get ra_ref, dec_ref
t = wcs.get_transform(frames[-2], 'world')
for m in t:
if isinstance(m, RotationSequence3D) and m.parameters.size == 5:
v2_ref, nv3_ref, roll_ref, dec_ref, nra_ref = m.angles.value
break
else:
raise ValueError(
"Unsupported WCS structure."
)
# overwrite v2_ref, v3_ref, and roll_ref with
# values from the tangent plane when available:
if 'v2v3corr' in frames:
# get v2_ref, v3_ref, and roll_ref from
# the v2v3 to v2v3corr transform:
frm1 = 'v2v3vacorr' if 'v2v3vacorr' in frames else 'v2v3'
tpcorr = wcs.get_transform(frm1, 'v2v3corr')
v2_ref, nv3_ref, roll_ref = tpcorr['det_to_optic_axis'].angles.value
wcsinfo = {
'v2_ref': 3600 * v2_ref,
'v3_ref': -3600 * nv3_ref,
'roll_ref': roll_ref,
'ra_ref': -nra_ref,
'dec_ref': dec_ref
}
return wcsinfo
[docs]
def adjust_wcs(wcs, delta_ra=0.0, delta_dec=0.0, delta_roll=0.0,
scale_factor=1.0):
"""
Apply corrections to an imaging WCS of 'cal' data models.
.. warning::
This function is not designed to handle neither FITS WCS nor
GWCS of resampled images. It is designed specifically for GWCS of
calibrated imaging data models that can be used as input to Stage 3
of the JWST pipeline (with suffixes '_cal', '_tweakreg', '_skymatch').
.. warning::
This function modifies the WCS of calibrated imaging data models in a
way that is NOT compatible with ``tweakreg``: once a WCS
was modified using ``adjust_wcs()``, the corresponding imaging data
model (whose WCS was modified) no longer be aligned using the
``tweakreg`` step.
Parameters
----------
wcs : `gwcs.WCS`
WCS object to be adjusted. Must be an imaging JWST WCS of a calibrated
data model.
delta_ra : float, astropy.units.Quantity, optional
Additional rotation (in degrees if units not provided) to be applied
along the longitude direction.
delta_dec : float, astropy.units.Quantity, optional
Additional rotation (in degrees if units not provided) to be applied
along the latitude direction.
delta_roll : float, astropy.units.Quantity, optional
Additional rotation (in degrees if units not provided) to be applied
to the telescope roll angle (rotation about V1 axis).
scale_factor : float, optional
A multiplicative scale factor to be applied to the current scale
(if any) in the WCS. If input ``wcs`` does not have a scale factor
applied, it is assumed to be 1. The scale factor is applied in
a tangent plane perpendicular to the V1 axis of the telescope.
Returns
-------
wcs : `gwcs.WCS`
Adjusted WCS object.
"""
# convert input angles to degrees:
u_deg = units.Unit('deg')
if isinstance(delta_ra, units.Quantity):
delta_ra = delta_ra.to(u_deg).value
if isinstance(delta_dec, units.Quantity):
delta_dec = delta_dec.to(u_deg).value
if isinstance(delta_roll, units.Quantity):
delta_roll = delta_roll.to(u_deg).value
# find the last frame in the pipeline that starts with 'v2v3':
pipeline = deepcopy(wcs.pipeline)
for step in pipeline[::-1]:
if (
step.frame.name and
step.frame.name.startswith('v2v3') and
step.transform.name == 'v23tosky'
):
s = step
break
else:
raise ValueError("Unknown WCS structure")
v2, v3, roll, dec, ra = s.transform.parameters[-5:]
v3 = -v3
ra = -ra + delta_ra
dec += delta_dec
roll += delta_roll
s.transform = _v23tosky(v2, v3, roll, ra, dec)
wcs = WCS(pipeline)
if scale_factor != 1.0:
# apply scale factor in the tangent plane:
corr = JWSTWCSCorrector(
wcs,
{
'v2_ref': 3600.0 * v2,
'v3_ref': 3600.0 * v3,
'roll_ref': 0.0
}
)
corr.set_correction(matrix=build_fit_matrix(0.0, scale_factor))
wcs = corr.wcs
return wcs
[docs]
def transfer_wcs_correction(to_image, from_image, matrix=None, shift=None):
"""
Applies the same *total* WCS correction that was applied by ``tweakreg`` (!)
to the WCS in the ``from_image`` data model to the WCS of the ``to_image``
data model. In some ways this function is analogous function to the
``tweakback`` function for HST available in the
`drizzlepac package <https://github.com/spacetelescope/drizzlepac>`_.
One fundamental difference between this function and ``tweakback``
is that JWST data models do not keep a history
of data's WCS via alternative WCS as it is done in HST data and so it is
impossible to select and apply only one particular WCS correction if
there were multiple corrections previously applied to a WCS. The
tangent-plane correction in JWST WCS is cumulative/total correction.
If you would like to apply a specific/custom correction, you can do that
via ``matrix`` and ``shift`` arguments which is defined in the
reference tangent plane provided by the ``from_image``'s WCS.
When providing your own corrections via ``matrix`` and ``shift`` arguments,
this function is similar to the :py:func:`adjust_wcs` function
but provides an alternative way of specifying corrections via affine
transformations in a reference tangent plane.
.. warning::
Upon return, if the ``to_image`` argument is an `ImageModel` it will be
modified with an updated ``ImageModel.meta.wcs`` WCS model.
If ``to_image`` argument is a file name of an ``ImageModel``, that
model will be read in, its WCS will be updated, and the updated model
will be written out to the same file. BACKUP the file in ``to_image``
argument before calling this function.
.. warning::
This function does not support input data models whose WCS were
modified by :py:func:`adjust_wcs`. Only WCS corrections computed by
either the ``tweakreg`` step or by ``tweakwcs`` package are supported.
Parameters
----------
to_image : str, ImageModel
Image model to which the correction should be applied/transferred to.
.. warning::
If it is a string file name then, upon return, this file
will be **overwritten** with a data model with an updated WCS.
from_image : str, ImageModel, gwcs.wcs.WCS
A data model whose WCS was previously corrected.
This data model plays two roles: 1) it is the reference WCS which
provides a tangent plane in which corrections have been defined, and
2) it provides WCS corrections to be applied to ``to_image``'s
WCS.
If the WCS of the ``from_image`` data model does not contain
corrections, then *both* ``matrix`` *and* ``shift`` arguments *must be
supplied*.
matrix : 2D list, 2D numpy.ndarray, None, optional
A 2D matrix part of an affine transformation defined in the tangent
plane derived from the ``from_image``'s WCS.
.. note::
When provided, ``shift`` argument *must also be provided* in
which case ``matrix`` and ``shift`` arguments override the
correction (if present) from the ``from_file``'s WCS.
shift : list, numpy.ndarray, None, optional
A list of length 2 representing the translational part of an affine
transformation (in arcsec) defined in the tangent plane derived
from the ``from_image``'s WCS.
.. note::
When provided, ``matrix`` argument *must also be provided* in
which case ``matrix`` and ``shift`` arguments override the
correction (if present) from the ``from_file``'s WCS.
Returns
-------
Upon return, if the ``to_image`` argument is an `ImageModel` it will be
modified with an updated ``ImageModel.meta.wcs`` WCS model.
If ``to_image`` argument is a file name of an ``ImageModel``, that
model will be read in, its WCS will be updated, and the updated model
will be written to the same file. BACKUP the file in ``to_image``
argument before calling this function.
"""
if isinstance(to_image, str):
to_file = to_image
to_image = ImageModel(to_file)
else:
to_file = None
if isinstance(from_image, str):
from_image = ImageModel(from_image)
refwcs = from_image.meta.wcs
ref_wcsinfo = from_image.meta.wcsinfo.instance
elif isinstance(from_image, WCS):
refwcs = from_image
ref_wcsinfo = _wcsinfo_from_wcs_transform(refwcs)
else: # an ImageModel object
refwcs = from_image.meta.wcs
ref_wcsinfo = from_image.meta.wcsinfo.instance
if ((matrix is None and shift is not None) or
(matrix is not None and shift is None)):
raise ValueError(
"Both 'matrix' and 'shift' must be either None or not None."
)
elif matrix is None and shift is None:
# use the correction that was applied to the "reference" WCS - the
# WCS from the "from_image":
if 'v2v3corr' not in refwcs.available_frames:
raise ValueError(
"The WCS of the 'from_image' data model has not been "
"previously corrected. A corrected WCS is needed in order to "
"\"transfer\" it to the 'to_file' when both 'matrix' and "
"'shift' are None."
)
t = refwcs.get_transform('v2v3', 'v2v3corr')
affine = t['tp_affine']
matrix = affine.matrix.value
shift = _RAD2ARCSEC * affine.translation.value # convert to arcsec
from_corr = JWSTWCSCorrector(
refwcs,
ref_wcsinfo
)
to_corr = JWSTWCSCorrector(
to_image.meta.wcs,
to_image.meta.wcsinfo.instance
)
to_corr.set_correction(matrix=matrix, shift=shift, ref_tpwcs=from_corr)
to_image.meta.wcs = to_corr.wcs
update_fits_wcsinfo(to_image)
if to_file:
to_image.write(to_file)