from __future__ import absolute_import, division, print_function
import copy
import numpy as np
from astropy.io import fits
import stwcs
from . import altwcs
from ..updatewcs import utils
from stsci.tools import fileutil
DEFAULT_WCS_KEYS = ['CRVAL1', 'CRVAL2', 'CRPIX1', 'CRPIX2',
'CD1_1', 'CD1_2', 'CD2_1', 'CD2_2',
'CTYPE1', 'CTYPE2', 'ORIENTAT']
DEFAULT_PRI_KEYS = ['HDRNAME', 'SIPNAME', 'NPOLNAME', 'D2IMNAME', 'DESCRIP']
COL_FITSKW_DICT = {'RMS_RA': 'sci.crder1', 'RMS_DEC': 'sci.crder2',
'NMatch': 'sci.nmatch', 'Catalog': 'sci.catalog'}
###
### WCSEXT table related keyword archive functions
###
[docs]def init_wcscorr(input, force=False):
"""
This function will initialize the WCSCORR table if it is not already present,
and look for WCS keywords with a prefix of 'O' as the original OPUS
generated WCS as the initial row for the table or use the current WCS
keywords as initial row if no 'O' prefix keywords are found.
This function will NOT overwrite any rows already present.
This function works on all SCI extensions at one time.
"""
# TODO: Create some sort of decorator or (for Python2.5) context for
# opening a FITS file and closing it when done, if necessary
if not isinstance(input, fits.HDUList):
# input must be a filename, so open as `astropy.io.fits.HDUList` object
fimg = fits.open(input, mode='update')
need_to_close = True
else:
fimg = input
need_to_close = False
# Do not try to generate a WCSCORR table for a simple FITS file
numsci = fileutil.countExtn(fimg)
if len(fimg) == 1 or numsci == 0:
return
enames = []
for e in fimg: enames.append(e.name)
if 'WCSCORR' in enames:
if not force:
return
else:
del fimg['wcscorr']
print('Initializing new WCSCORR table for ', fimg.filename())
used_wcskeys = altwcs.wcskeys(fimg['SCI', 1].header)
# define the primary columns of the WCSEXT table with initial rows for each
# SCI extension for the original OPUS solution
numwcs = len(used_wcskeys)
if numwcs == 0: numwcs = 1
# create new table with more rows than needed initially to make it easier to
# add new rows later
wcsext = create_wcscorr(descrip=True, numrows=numsci, padding=(numsci * numwcs) + numsci * 4)
# Assign the correct EXTNAME value to this table extension
wcsext.header['TROWS'] = (numsci * 2, 'Number of updated rows in table')
wcsext.header['EXTNAME'] = ('WCSCORR', 'Table with WCS Update history')
wcsext.header['EXTVER'] = 1
# define set of WCS keywords which need to be managed and copied to the table
wcs1 = stwcs.wcsutil.HSTWCS(fimg, ext=('SCI', 1))
idc2header = True
if wcs1.idcscale is None:
idc2header = False
wcs_keywords = list(wcs1.wcs2header(idc2hdr=idc2header).keys())
prihdr = fimg[0].header
prihdr_keys = DEFAULT_PRI_KEYS
pri_funcs = {'SIPNAME': stwcs.updatewcs.utils.build_sipname,
'NPOLNAME': stwcs.updatewcs.utils.build_npolname,
'D2IMNAME': stwcs.updatewcs.utils.build_d2imname}
# Now copy original OPUS values into table
for extver in range(1, numsci + 1):
rowind = find_wcscorr_row(wcsext.data,
{'WCS_ID': 'OPUS', 'EXTVER': extver,
'WCS_key': 'O'})
# There should only EVER be a single row for each extension with OPUS values
rownum = np.where(rowind)[0][0]
# print 'Archiving OPUS WCS in row number ',rownum,' in WCSCORR table for SCI,',extver
hdr = fimg['SCI', extver].header
# define set of WCS keywords which need to be managed and copied to the table
if used_wcskeys is None:
used_wcskeys = altwcs.wcskeys(hdr)
# Check to see whether or not there is an OPUS alternate WCS present,
# if so, get its values directly, otherwise, archive the PRIMARY WCS
# as the OPUS values in the WCSCORR table
if 'O' not in used_wcskeys:
altwcs.archiveWCS(fimg, ('SCI', extver), wcskey='O', wcsname='OPUS')
wkey = 'O'
wcs = stwcs.wcsutil.HSTWCS(fimg, ext=('SCI', extver), wcskey=wkey)
wcshdr = wcs.wcs2header(idc2hdr=idc2header)
if wcsext.data.field('CRVAL1')[rownum] != 0:
# If we find values for these keywords already in the table, do not
# overwrite them again
print('WCS keywords already updated...')
break
for key in wcs_keywords:
if key in wcsext.data.names:
wcsext.data.field(key)[rownum] = wcshdr[(key + wkey)[:8]]
# Now get any keywords from PRIMARY header needed for WCS updates
for key in prihdr_keys:
if key in prihdr:
val = prihdr[key]
else:
val = ''
wcsext.data.field(key)[rownum] = val
# Now that we have archived the OPUS alternate WCS, remove it from the list
# of used_wcskeys
if 'O' in used_wcskeys:
used_wcskeys.remove('O')
# Now copy remaining alternate WCSs into table
# TODO: Much of this appears to be redundant with update_wcscorr; consider
# merging them...
for uwkey in used_wcskeys:
for extver in range(1, numsci + 1):
hdr = fimg['SCI', extver].header
wcs = stwcs.wcsutil.HSTWCS(fimg, ext=('SCI', extver),
wcskey=uwkey)
wcshdr = wcs.wcs2header()
if 'WCSNAME' + uwkey not in wcshdr:
wcsid = utils.build_default_wcsname(fimg[0].header['idctab'])
else:
wcsid = wcshdr['WCSNAME' + uwkey]
# identify next empty row
rowind = find_wcscorr_row(wcsext.data,
selections={'wcs_id': ['', '0.0']})
rows = np.where(rowind)
if len(rows[0]) > 0:
rownum = np.where(rowind)[0][0]
else:
print('No available rows found for updating. ')
# Update selection columns for this row with relevant values
wcsext.data.field('WCS_ID')[rownum] = wcsid
wcsext.data.field('EXTVER')[rownum] = extver
wcsext.data.field('WCS_key')[rownum] = uwkey
# Look for standard WCS keyword values
for key in wcs_keywords:
if key in wcsext.data.names:
wcsext.data.field(key)[rownum] = wcshdr[key + uwkey]
# Now get any keywords from PRIMARY header needed for WCS updates
for key in prihdr_keys:
if key in pri_funcs:
val = pri_funcs[key](fimg)[0]
else:
if key in prihdr:
val = prihdr[key]
else:
val = ''
wcsext.data.field(key)[rownum] = val
# Append this table to the image FITS file
fimg.append(wcsext)
# force an update now
# set the verify flag to 'warn' so that it will always succeed, but still
# tell the user if PyFITS detects any problems with the file as a whole
utils.updateNEXTENDKw(fimg)
fimg.flush('warn')
if need_to_close:
fimg.close()
[docs]def find_wcscorr_row(wcstab, selections):
"""
Return an array of indices from the table (NOT HDU) 'wcstab' that matches the
selections specified by the user.
The row selection criteria must be specified as a dictionary with
column name as key and value(s) representing the valid desired row values.
For example, {'wcs_id':'OPUS','extver':2}.
"""
mask = None
for i in selections:
icol = wcstab.field(i)
if isinstance(icol, np.chararray): icol = icol.rstrip()
selecti = selections[i]
if not isinstance(selecti, list):
if isinstance(selecti, str):
selecti = selecti.rstrip()
bmask = (icol == selecti)
if mask is None:
mask = bmask.copy()
else:
mask = np.logical_and(mask, bmask)
del bmask
else:
for si in selecti:
if isinstance(si, str):
si = si.rstrip()
bmask = (icol == si)
if mask is None:
mask = bmask.copy()
else:
mask = np.logical_or(mask, bmask)
del bmask
return mask
[docs]def archive_wcs_file(image, wcs_id=None):
"""
Update WCSCORR table with rows for each SCI extension to record the
newly updated WCS keyword values.
"""
if not isinstance(image, fits.HDUList):
fimg = fits.open(image, mode='update')
close_image = True
else:
fimg = image
close_image = False
update_wcscorr(fimg, wcs_id=wcs_id)
if close_image:
fimg.close()
[docs]def update_wcscorr(dest, source=None, extname='SCI', wcs_id=None, active=True):
"""
Update WCSCORR table with a new row or rows for this extension header. It
copies the current set of WCS keywords as a new row of the table based on
keyed WCSs as per Paper I Multiple WCS standard).
Parameters
----------
dest : HDUList
The HDU list whose WCSCORR table should be appended to (the WCSCORR HDU
must already exist)
source : HDUList, optional
The HDU list containing the extension from which to extract the WCS
keywords to add to the WCSCORR table. If None, the dest is also used
as the source.
extname : str, optional
The extension name from which to take new WCS keywords. If there are
multiple extensions with that name, rows are added for each extension
version.
wcs_id : str, optional
The name of the WCS to add, as in the WCSNAMEa keyword. If
unspecified, all the WCSs in the specified extensions are added.
active: bool, optional
When True, indicates that the update should reflect an update of the
active WCS information, not just appending the WCS to the file as a
headerlet
"""
if not isinstance(dest, fits.HDUList):
dest = fits.open(dest, mode='update')
fname = dest.filename()
if source is None:
source = dest
if extname == 'PRIMARY':
return
numext = fileutil.countExtn(source, extname)
if numext == 0:
raise ValueError('No %s extensions found in the source HDU list.'
% extname)
# Initialize the WCSCORR table extension in dest if not already present
init_wcscorr(dest)
try:
dest.index_of('WCSCORR')
except KeyError:
return
# check to see whether or not this is an up-to-date table
# replace with newly initialized table with current format
old_table = dest['WCSCORR']
wcscorr_cols = ['WCS_ID', 'EXTVER', 'SIPNAME',
'HDRNAME', 'NPOLNAME', 'D2IMNAME']
for colname in wcscorr_cols:
if colname not in old_table.data.columns.names:
print("WARNING: Replacing outdated WCSCORR table...")
outdated_table = old_table.copy()
del dest['WCSCORR']
init_wcscorr(dest)
old_table = dest['WCSCORR']
break
# Current implementation assumes the same WCS keywords are in each
# extension version; if this should not be assumed then this can be
# modified...
wcs_keys = altwcs.wcskeys(source[(extname, 1)].header)
wcs_keys = [kk for kk in wcs_keys if kk]
if ' ' not in wcs_keys: wcs_keys.append(' ') # Insure that primary WCS gets used
# apply logic for only updating WCSCORR table with specified keywords
# corresponding to the WCS with WCSNAME=wcs_id
if wcs_id is not None:
wnames = altwcs.wcsnames(source[(extname, 1)].header)
wkeys = []
for letter in wnames:
if wnames[letter] == wcs_id:
wkeys.append(letter)
if len(wkeys) > 1 and ' ' in wkeys:
wkeys.remove(' ')
wcs_keys = wkeys
wcshdr = stwcs.wcsutil.HSTWCS(source, ext=(extname, 1)).wcs2header()
wcs_keywords = list(wcshdr.keys())
if 'O' in wcs_keys:
wcs_keys.remove('O') # 'O' is reserved for original OPUS WCS
# create new table for hdr and populate it with the newly updated values
new_table = create_wcscorr(descrip=True, numrows=0, padding=len(wcs_keys) * numext)
prihdr = source[0].header
# Get headerlet related keywords here
sipname, idctab = utils.build_sipname(source, fname, "None")
npolname, npolfile = utils.build_npolname(source, None)
d2imname, d2imfile = utils.build_d2imname(source, None)
if 'hdrname' in prihdr:
hdrname = prihdr['hdrname']
else:
hdrname = ''
idx = -1
for wcs_key in wcs_keys:
for extver in range(1, numext + 1):
extn = (extname, extver)
if 'SIPWCS' in extname and not active:
tab_extver = 0 # Since it has not been added to the SCI header yet
else:
tab_extver = extver
hdr = source[extn].header
if 'WCSNAME' + wcs_key in hdr:
wcsname = hdr['WCSNAME' + wcs_key]
else:
wcsname = utils.build_default_wcsname(hdr['idctab'])
selection = {'WCS_ID': wcsname, 'EXTVER': tab_extver,
'SIPNAME': sipname, 'HDRNAME': hdrname,
'NPOLNAME': npolname, 'D2IMNAME': d2imname
}
# Ensure that an entry for this WCS is not already in the dest
# table; if so just skip it
rowind = find_wcscorr_row(old_table.data, selection)
if np.any(rowind):
continue
idx += 1
wcs = stwcs.wcsutil.HSTWCS(source, ext=extn, wcskey=wcs_key)
wcshdr = wcs.wcs2header()
# Update selection column values
for key, val in selection.items():
if key in new_table.data.names:
new_table.data.field(key)[idx] = val
for key in wcs_keywords:
if key in new_table.data.names:
new_table.data.field(key)[idx] = wcshdr[key + wcs_key]
for key in DEFAULT_PRI_KEYS:
if key in new_table.data.names and key in prihdr:
new_table.data.field(key)[idx] = prihdr[key]
# Now look for additional, non-WCS-keyword table column data
for key in COL_FITSKW_DICT:
fitkw = COL_FITSKW_DICT[key]
# Interpret any 'pri.hdrname' or
# 'sci.crpix1' formatted keyword names
if '.' in fitkw:
srchdr, fitkw = fitkw.split('.')
if 'pri' in srchdr.lower(): srchdr = prihdr
else: srchdr = source[extn].header
else:
srchdr = source[extn].header
if fitkw + wcs_key in srchdr:
new_table.data.field(key)[idx] = srchdr[fitkw + wcs_key]
# If idx was never incremented, no rows were added, so there's nothing else
# to do...
if idx < 0:
return
# Now, we need to merge this into the existing table
rowind = find_wcscorr_row(old_table.data, {'wcs_id': ['', '0.0']})
old_nrows = np.where(rowind)[0][0]
new_nrows = new_table.data.shape[0]
# check to see if there is room for the new row
if (old_nrows + new_nrows) > old_table.data.shape[0] - 1:
pad_rows = 2 * new_nrows
# if not, create a new table with 'pad_rows' new empty rows
upd_table = fits.BinTableHDU.from_columns(old_table.columns, header=old_table.header,
nrows=old_table.data.shape[0] + pad_rows)
else:
upd_table = old_table
pad_rows = 0
# Now, add
for name in old_table.columns.names:
if name in new_table.data.names:
# reset the default values to ones specific to the row definitions
for i in range(pad_rows):
upd_table.data.field(name)[old_nrows + i] = old_table.data.field(name)[-1]
# Now populate with values from new table
upd_table.data.field(name)[old_nrows:old_nrows + new_nrows] = \
new_table.data.field(name)
upd_table.header['TROWS'] = old_nrows + new_nrows
# replace old extension with newly updated table extension
dest['WCSCORR'] = upd_table
[docs]def restore_file_from_wcscorr(image, id='OPUS', wcskey=''):
""" Copies the values of the WCS from the WCSCORR based on ID specified by user.
The default will be to restore the original OPUS-derived values to the Primary WCS.
If wcskey is specified, the WCS with that key will be updated instead.
"""
if not isinstance(image, fits.HDUList):
fimg = fits.open(image, mode='update')
close_image = True
else:
fimg = image
close_image = False
numsci = fileutil.countExtn(fimg)
wcs_table = fimg['WCSCORR']
orig_rows = (wcs_table.data.field('WCS_ID') == 'OPUS')
# create an HSTWCS object to figure out what WCS keywords need to be updated
wcsobj = stwcs.wcsutil.HSTWCS(fimg, ext=('sci', 1))
wcshdr = wcsobj.wcs2header()
for extn in range(1, numsci + 1):
# find corresponding row from table
ext_rows = (wcs_table.data.field('EXTVER') == extn)
erow = np.where(np.logical_and(ext_rows, orig_rows))[0][0]
for key in wcshdr:
if key in wcs_table.data.names: # insure that keyword is column in table
tkey = key
if 'orient' in key.lower():
key = 'ORIENT'
if wcskey == '':
skey = key
else:
skey = key[:7] + wcskey
fimg['sci', extn].header[skey] = wcs_table.data.field(tkey)[erow]
for key in DEFAULT_PRI_KEYS:
if key in wcs_table.data.names:
if wcskey == '':
pkey = key
else:
pkey = key[: 7] + wcskey
fimg[0].header[pkey] = wcs_table.data.field(key)[erow]
utils.updateNEXTENDKw(fimg)
# close the image now that the update has been completed.
if close_image:
fimg.close()
[docs]def create_wcscorr(descrip=False, numrows=1, padding=0):
"""
Return the basic definitions for a WCSCORR table.
The dtype definitions for the string columns are set to the maximum allowed so
that all new elements will have the same max size which will be automatically
truncated to this limit upon updating (if needed).
The table is initialized with rows corresponding to the OPUS solution
for all the 'SCI' extensions.
"""
trows = numrows + padding
# define initialized arrays as placeholders for column data
# TODO: I'm certain there's an easier way to do this... for example, simply
# define the column names and formats, then create an empty array using
# them as a dtype, then create the new table from that array.
def_float64_zeros = np.array([0.0] * trows, dtype=np.float64)
def_float64_ones = def_float64_zeros + 1.0
def_float_col = {'format': 'D', 'array': def_float64_zeros.copy()}
def_float1_col = {'format': 'D', 'array': def_float64_ones.copy()}
def_str40_col = {'format': '40A',
'array': np.array([''] * trows, dtype='S40')}
def_str24_col = {'format': '24A',
'array': np.array([''] * trows, dtype='S24')}
def_int32_col = {'format': 'J',
'array': np.array([0] * trows, dtype=np.int32)}
# If more columns are needed, simply add their definitions to this list
col_names = [('HDRNAME', def_str24_col), ('SIPNAME', def_str24_col),
('NPOLNAME', def_str24_col), ('D2IMNAME', def_str24_col),
('CRVAL1', def_float_col), ('CRVAL2', def_float_col),
('CRPIX1', def_float_col), ('CRPIX2', def_float_col),
('CD1_1', def_float_col), ('CD1_2', def_float_col),
('CD2_1', def_float_col), ('CD2_2', def_float_col),
('CTYPE1', def_str24_col), ('CTYPE2', def_str24_col),
('ORIENTAT', def_float_col), ('PA_V3', def_float_col),
('RMS_RA', def_float_col), ('RMS_Dec', def_float_col),
('NMatch', def_int32_col), ('Catalog', def_str40_col)]
# Define selector columns
id_col = fits.Column(name='WCS_ID', format='40A',
array=np.array(['OPUS'] * numrows + [''] * padding,
dtype='S24'))
extver_col = fits.Column(name='EXTVER', format='I',
array=np.array(list(range(1, numrows + 1)),
dtype=np.int16))
wcskey_col = fits.Column(name='WCS_key', format='A',
array=np.array(['O'] * numrows + [''] * padding,
dtype='S'))
# create list of remaining columns to be added to table
col_list = [id_col, extver_col, wcskey_col] # start with selector columns
for c in col_names:
cdef = copy.deepcopy(c[1])
col_list.append(fits.Column(name=c[0], format=cdef['format'],
array=cdef['array']))
if descrip:
col_list.append(
fits.Column(name='DESCRIP', format='128A',
array=np.array(
['Original WCS computed by OPUS'] * numrows,
dtype='S128')))
# Now create the new table from the column definitions
newtab = fits.BinTableHDU.from_columns(fits.ColDefs(col_list), nrows=trows)
# The fact that setting .name is necessary should be considered a bug in
# pyfits.
# TODO: Make sure this is fixed in pyfits, then remove this
newtab.name = 'WCSCORR'
return newtab
[docs]def delete_wcscorr_row(wcstab, selections=None, rows=None):
"""
Sets all values in a specified row or set of rows to default values
This function will essentially erase the specified row from the table
without actually removing the row from the table. This avoids the problems
with trying to resize the number of rows in the table while preserving the
ability to update the table with new rows again without resizing the table.
Parameters
----------
wcstab: object
PyFITS binTable object for WCSCORR table
selections: dict
Dictionary of wcscorr column names and values to be used to select
the row or set of rows to erase
rows: int, list
If specified, will specify what rows from the table to erase regardless
of the value of 'selections'
"""
if selections is None and rows is None:
print('ERROR: Some row selection information must be provided!')
print(' Either a row numbers or "selections" must be provided.')
raise ValueError
delete_rows = None
if rows is None:
if 'wcs_id' in selections and selections['wcs_id'] == 'OPUS':
delete_rows = None
print('WARNING: OPUS WCS information can not be deleted from WCSCORR table.')
print(' This row will not be deleted!')
else:
rowind = find_wcscorr_row(wcstab, selections=selections)
delete_rows = np.where(rowind)[0].tolist()
else:
if not isinstance(rows, list):
rows = [rows]
delete_rows = rows
# Insure that rows pointing to OPUS WCS do not get deleted, even by accident
for row in delete_rows:
if wcstab['WCS_key'][row] == 'O' or wcstab['WCS_ID'][row] == 'OPUS':
del delete_rows[delete_rows.index(row)]
if delete_rows is None:
return
# identify next empty row
rowind = find_wcscorr_row(wcstab, selections={'wcs_id': ['', '0.0']})
last_blank_row = np.where(rowind)[0][-1]
# copy values from blank row into user-specified rows
for colname in wcstab.names:
wcstab[colname][delete_rows] = wcstab[colname][last_blank_row]
[docs]def update_wcscorr_column(wcstab, column, values, selections=None, rows=None):
"""
Update the values in 'column' with 'values' for selected rows
Parameters
----------
wcstab: object
PyFITS binTable object for WCSCORR table
column: string
Name of table column with values that need to be updated
values: string, int, or list
Value or set of values to copy into the selected rows for the column
selections: dict
Dictionary of wcscorr column names and values to be used to select
the row or set of rows to erase
rows: int, list
If specified, will specify what rows from the table to erase regardless
of the value of 'selections'
"""
if selections is None and rows is None:
print('ERROR: Some row selection information must be provided!')
print(' Either a row numbers or "selections" must be provided.')
raise ValueError
if not isinstance(values, list):
values = [values]
update_rows = None
if rows is None:
if 'wcs_id' in selections and selections['wcs_id'] == 'OPUS':
update_rows = None
print('WARNING: OPUS WCS information can not be deleted from WCSCORR table.')
print(' This row will not be deleted!')
else:
rowind = find_wcscorr_row(wcstab, selections=selections)
update_rows = np.where(rowind)[0].tolist()
else:
if not isinstance(rows, list):
rows = [rows]
update_rows = rows
if update_rows is None:
return
# Expand single input value to apply to all selected rows
if len(values) > 1 and len(values) < len(update_rows):
print('ERROR: Number of new values', len(values))
print(' does not match number of rows', len(update_rows), ' to be updated!')
print(' Please enter either 1 value or the same number of values')
print(' as there are rows to be updated.')
print(' Table will not be updated...')
raise ValueError
if len(values) == 1 and len(values) < len(update_rows):
values = values * len(update_rows)
# copy values from blank row into user-specified rows
for row in update_rows:
wcstab[column][row] = values[row]