Python astropy.wcs.WCS Examples

def update_wcs_fits_log(file, wcs_ref, xyscale=[0, 0, 0, 1], initialize=True, replace=('.fits', '.wcslog.fits'), wcsname='SHIFT'):
    """
    Make FITS log when updating WCS
    """
    new_hdu = wcs_ref.to_fits(relax=True)[0]
    new_hdu.header['XSHIFT'] = xyscale[0]
    new_hdu.header['YSHIFT'] = xyscale[1]
    new_hdu.header['ROT'] = xyscale[2], 'WCS fit rotation, degrees'
    new_hdu.header['SCALE'] = xyscale[3], 'WCS fit scale'
    new_hdu.header['WCSNAME'] = wcsname

    wcs_logfile = file.replace(replace[0], replace[1])

    if os.path.exists(wcs_logfile):
        if initialize:
            os.remove(wcs_logfile)
            hdu = pyfits.HDUList([pyfits.PrimaryHDU()])
        else:
            hdu = pyfits.open(wcs_logfile)
    else:
        hdu = pyfits.HDUList([pyfits.PrimaryHDU()])

    hdu.append(new_hdu)
    hdu.writeto(wcs_logfile, overwrite=True, output_verify='fix') 

def _load_cube_from_file(self, data=None):
        """Initialises a cube from a file."""

        if data is not None:
            assert isinstance(data, fits.HDUList), 'data is not an HDUList object'
        else:
            try:
                with gunzip(self.filename) as gg:
                    self.data = fits.open(gg.name)
            except (IOError, OSError) as err:
                raise OSError('filename {0} cannot be found: {1}'.format(self.filename, err))

        self.header = self.data[1].header
        self.wcs = WCS(self.header)

        self._check_file(self.data[0].header, self.data, 'Cube')

        self._wavelength = self.data['WAVE'].data
        self._shape = (self.header['NAXIS2'], self.header['NAXIS1'])

        self._do_file_checks(self) 

def _get_fits_mbr(fin, row_ignore_factor=10):
    hdulist = pyfits.open(fin)
    data = hdulist[0].data
    wcs = pywcs.WCS(hdulist[0].header)
    width = data.shape[1]
    height = data.shape[0]

    bottom_left = [0, 0, 0, 0]
    top_left = [0, height - 1, 0, 0]
    top_right = [width - 1, height - 1, 0, 0]
    bottom_right = [width - 1, 0, 0, 0]

    def pix2sky(pix_coord):
        return wcs.wcs_pix2world([pix_coord], 0)[0][0:2]

    ret = np.zeros([4, 2])
    ret[0, :] = pix2sky(bottom_left)
    ret[1, :] = pix2sky(top_left)
    ret[2, :] = pix2sky(top_right)
    ret[3, :] = pix2sky(bottom_right)
    RA_min, DEC_min, RA_max, DEC_max = np.min(ret[:, 0]),   np.min(ret[:, 1]),  np.max(ret[:, 0]),  np.max(ret[:, 1])
    
    # http://pgsphere.projects.pgfoundry.org/types.html
    sqlStr = "SELECT sbox '((%10fd, %10fd), (%10fd, %10fd))'" % (RA_min, DEC_min, RA_max, DEC_max)
    return sqlStr 

def pixel_to_world_values(self, *pixel_arrays):
        """
        Convert pixel coordinates to world coordinates.

        This method takes `~astropy.wcs.wcsapi.BaseLowLevelWCS.pixel_n_dim` scalars or arrays as
        input, and pixel coordinates should be zero-based. Returns
        `~astropy.wcs.wcsapi.BaseLowLevelWCS.world_n_dim` scalars or arrays in units given by
        `~astropy.wcs.wcsapi.BaseLowLevelWCS.world_axis_units`. Note that pixel coordinates are
        assumed to be 0 at the center of the first pixel in each dimension. If a
        pixel is in a region where the WCS is not defined, NaN can be returned.
        The coordinates should be specified in the ``(x, y)`` order, where for
        an image, ``x`` is the horizontal coordinate and ``y`` is the vertical
        coordinate.

        If `~astropy.wcs.wcsapi.BaseLowLevelWCS.world_n_dim` is ``1``, this
        method returns a single scalar or array, otherwise a tuple of scalars or
        arrays is returned.
        """ 

def __init__(self, input, output, origin=1):
        """Sample class to demonstrate how to define a coordinate transformation

        Modified from `drizzlepac.wcs_functions.WCSMap` to use full SIP header
        in the `forward` and `backward` methods. Use this class to drizzle to
        an output distorted WCS, e.g.,

            >>> drizzlepac.astrodrizzle.do_driz(..., wcsmap=SIP_WCSMap)
        """

        # Verify that we have valid WCS input objects
        self.checkWCS(input, 'Input')
        self.checkWCS(output, 'Output')

        self.input = input
        self.output = copy.deepcopy(output)
        #self.output = output

        self.origin = origin
        self.shift = None
        self.rot = None
        self.scale = None 

def _load_cube_from_api(self):
        """Calls the API and retrieves the necessary information to instantiate the cube."""

        url = marvin.config.urlmap['api']['getCube']['url']

        try:
            response = self._toolInteraction(url.format(name=self.plateifu))
        except Exception as ee:
            raise MarvinError('found a problem when checking if remote cube '
                              'exists: {0}'.format(str(ee)))

        data = response.getData()

        self.header = fits.Header.fromstring(data['header'])
        self.wcs = WCS(fits.Header.fromstring(data['wcs_header']))
        self._wavelength = data['wavelength']
        self._shape = data['shape']

        if self.plateifu != data['plateifu']:
            raise MarvinError('remote cube has a different plateifu!')

        return 

def _load_modelcube_from_api(self):
        """Initialises a model cube from the API."""

        url = marvin.config.urlmap['api']['getModelCube']['url']
        url_full = url.format(name=self.plateifu, bintype=self.bintype.name,
                              template=self.template.name)

        try:
            response = self._toolInteraction(url_full)
        except Exception as ee:
            raise MarvinError('found a problem when checking if remote model cube '
                              'exists: {0}'.format(str(ee)))

        data = response.getData()

        self.header = fits.Header.fromstring(data['header'])
        self.wcs = WCS(fits.Header.fromstring(data['wcs_header']))
        self._wavelength = np.array(data['wavelength'])
        self._redcorr = np.array(data['redcorr'])
        self._shape = tuple(data['shape'])

        self.plateifu = str(self.header['PLATEIFU'].strip())
        self.mangaid = str(self.header['MANGAID'].strip()) 

def pixel_shape(self):
        """
        The shape of the data that the WCS applies to as a tuple of length
        `~astropy.wcs.wcsapi.BaseLowLevelWCS.pixel_n_dim` in ``(x, y)``
        order (where for an image, ``x`` is the horizontal coordinate and ``y``
        is the vertical coordinate).

        If the WCS is valid in the context of a dataset with a particular
        shape, then this property can be used to store the shape of the
        data. This can be used for example if implementing slicing of WCS
        objects. This is an optional property, and it should return `None`
        if a shape is not known or relevant.

        If you are interested in getting a shape that is comparable to that of
        a Numpy array, you should use
        `~astropy.wcs.wcsapi.BaseLowLevelWCS.array_shape` instead.
        """
        return None 

def get_data(self, filter='f140w', catfile=None, segfile=None):
        import glob

        sci_file = glob.glob('{0}-{1}_dr?_sci.fits'.format(self.root, filter))

        sci = pyfits.open(sci_file[0])
        wht = pyfits.open(sci_file[0].replace('sci', 'wht'))

        if segfile is None:
            segfile = sci_file[0].split('_dr')[0]+'_seg.fits'

        seg = pyfits.open(segfile)

        if catfile is None:
            catfile = '{0}-{1}.cat.fits'.format(self.root, filter)

        cat = utils.GTable.gread(catfile)

        wcs = pywcs.WCS(sci[0].header, relax=True)
        return sci, wht, seg, cat, wcs 

def get_header(self, ext=0):
        """Returns the metadata embedded in the file.

        Target Pixel Files contain embedded metadata headers spread across three
        different FITS extensions:

        1. The "PRIMARY" extension (``ext=0``) provides a metadata header
           providing details on the target and its CCD position.
        2. The "PIXELS" extension (``ext=1``) provides details on the
           data column and their coordinate system (WCS).
        3. The "APERTURE" extension (``ext=2``) provides details on the
           aperture pixel mask and the expected coordinate system (WCS).

        Parameters
        ----------
        ext : int or str
            FITS extension name or number.

        Returns
        -------
        header : `~astropy.io.fits.header.Header`
            Header object containing metadata keywords.
        """
        return self.hdu[ext].header 

def make_SEP_FLT_catalog(flt_file, ext=1, column_case=str.upper, **kwargs):
    import astropy.io.fits as pyfits
    import astropy.wcs as pywcs

    im = pyfits.open(flt_file)
    sci = im['SCI', ext].data - im['SCI', ext].header['MDRIZSKY']
    err = im['ERR', ext].data
    mask = im['DQ', ext].data > 0

    ZP = utils.calc_header_zeropoint(im, ext=('SCI', ext))

    try:
        wcs = pywcs.WCS(im['SCI', ext].header, fobj=im)
    except:
        wcs = None

    tab, seg = make_SEP_catalog_from_arrays(sci, err, mask, wcs=wcs, ZP=ZP, **kwargs)
    tab.meta['ABZP'] = ZP
    tab.meta['FILTER'] = utils.get_hst_filter(im[0].header)
    tab['mag_auto'] = ZP - 2.5*np.log10(tab['flux'])

    for c in tab.colnames:
        tab.rename_column(c, column_case(c))

    return tab, seg 

def test_dist():
    with get_pkg_data_fileobj(
            os.path.join("data", "dist.fits"), encoding='binary') as test_file:
        hdulist = fits.open(test_file)
        # The use of ``AXISCORR`` for D2IM correction has been deprecated
        with pytest.warns(AstropyDeprecationWarning):
            wcs1 = wcs.WCS(hdulist[0].header, hdulist)
        assert wcs1.det2im2 is not None
        with pytest.warns(VerifyWarning):
            s = pickle.dumps(wcs1)
        with pytest.warns(FITSFixedWarning):
            wcs2 = pickle.loads(s)

        with NumpyRNGContext(123456789):
            x = np.random.rand(2 ** 16, wcs1.wcs.naxis)
            world1 = wcs1.all_pix2world(x, 1)
            world2 = wcs2.all_pix2world(x, 1)

        assert_array_almost_equal(world1, world2) 

def test_sip():
    with get_pkg_data_fileobj(
            os.path.join("data", "sip.fits"), encoding='binary') as test_file:
        hdulist = fits.open(test_file, ignore_missing_end=True)
        with pytest.warns(FITSFixedWarning):
            wcs1 = wcs.WCS(hdulist[0].header)
        assert wcs1.sip is not None
        s = pickle.dumps(wcs1)
        with pytest.warns(FITSFixedWarning):
            wcs2 = pickle.loads(s)

        with NumpyRNGContext(123456789):
            x = np.random.rand(2 ** 16, wcs1.wcs.naxis)
            world1 = wcs1.all_pix2world(x, 1)
            world2 = wcs2.all_pix2world(x, 1)

        assert_array_almost_equal(world1, world2) 

def __init__(self, base_fname):
        """
        Args:
            base_fname (str): The map should be stored in two FITS files, named
                ``base_fname + '_' + X + '.fits'``, where ``X`` is ``'ngp'`` and
                ``'sgp'``.
        """
        self._data = {}

        for pole in self.poles:
            fname = '{}_{}.fits'.format(base_fname, pole)
            try:
                with fits.open(fname) as hdulist:
                    self._data[pole] = [hdulist[0].data, wcs.WCS(hdulist[0].header)]
            except IOError as error:
                print(dustexceptions.data_missing_message(self.map_name,
                                                          self.map_name_long))
                raise error 

def _test_init(self, model_cube, galaxy, bintype=None, template=None):
        assert model_cube.release == galaxy.release
        assert model_cube._drpver == galaxy.drpver
        assert model_cube._dapver == galaxy.dapver
        assert model_cube.bintype.name == bintype if bintype else galaxy.bintype.name
        assert model_cube.template == template if template else galaxy.template
        assert model_cube.plateifu == galaxy.plateifu
        assert model_cube.mangaid == galaxy.mangaid
        assert isinstance(model_cube.header, fits.Header)
        assert isinstance(model_cube.wcs, WCS)
        assert model_cube._wavelength is not None
        assert model_cube._redcorr is not None 

def time_1d_wcs(header_time_1d):
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', FITSFixedWarning)
        return WCS(header_time_1d) 

def test_map(filename):
    header = get_pkg_data_contents(os.path.join("data/maps", filename))
    wcsobj = wcs.WCS(header)

    with NumpyRNGContext(123456789):
        x = np.random.rand(2 ** 12, wcsobj.wcs.naxis)
        wcsobj.wcs_pix2world(x, 1)
        wcsobj.wcs_world2pix(x, 1) 

def test_wcs_info_str():

    # The tests in test_sliced_low_level_wcs.py excercise wcs_info_str
    # extensively. This test is to ensure that the function exists and the
    # API of the function works as expected.

    wcs_empty = WCS(naxis=1)

    assert wcs_info_str(wcs_empty).strip() == DEFAULT_1D_STR.strip() 

def test_basic():
    wcs1 = wcs.WCS()
    s = pickle.dumps(wcs1)
    with pytest.warns(FITSFixedWarning):
        pickle.loads(s) 

def test_too_much_slicing():
    wcs = WCS_SPECTRAL_CUBE
    with pytest.raises(ValueError, match='Cannot slice WCS: the resulting WCS '
                                         'should have at least one pixel and '
                                         'one world dimension'):
        wcs[0, 1, 2] 

def wcs(self):
        return WCS(self.header) 

def test_wcs(self, cube):
        assert cube.data_origin == cube.exporigin
        assert isinstance(cube.wcs, wcs.WCS)
        comp = cube.wcs.wcs.pc if cube.data_origin == 'api' else cube.wcs.wcs.cd
        assert comp[1, 1] == pytest.approx(0.000138889) 

def test_image_has_wcs(self, galaxy):
        w = getWCSFromPng(galaxy.imgpath)
        assert isinstance(w, WCS) is True 

def wcs(galaxy):
    return WCS(fits.getheader(galaxy.cubepath, 1)) 

def _make_metadata(self, filetype=None):
        ''' Make the meta data for the image '''

        if 'png' in filetype:
            meta = PIL.PngImagePlugin.PngInfo()
        else:
            meta = None
            warnings.warn('Can only save WCS metadata with PNG filetype', MarvinUserWarning)

        if meta:
            info = {key: str(val) for key, val in self.image.info.items()}
            for row in info:
                meta.add_text(row, info[row])

        return meta 

def _wcs_to_dict(self):
        ''' Convert and return the WCS as a dictionary'''
        wcshdr = None
        if self.wcs:
            wcshdr = self.wcs.to_header()
            wcshdr = dict(wcshdr)
            wcshdr = {key: str(val) for key, val in wcshdr.items()}
        return wcshdr 

def _define_wcs(self):
        """
        Given what we know about the scale of the image,
        define a nearly-correct world coordinate system to use with it.
        """
        w = WCS(naxis=2)
        w.wcs.crpix = self.size_pix / 2
        w.wcs.crval = self.center
        w.wcs.cd = np.array([[-1, 0], [0, 1]]) * self.scale / 3600.
        w.wcs.ctype = ['RA---TAN', 'DEC--TAN']
        w.wcs.cunit = ['deg', 'deg']
        w.wcs.radesys = 'ICRS'
        w.wcs.equinox = 2000.0
        self.wcs = w 

def convertImgCoords(coords, image, to_pix=None, to_radec=None):
    """Transform the WCS info in an image.

    Convert image pixel coordinates to RA/Dec based on
    PNG image metadata or vice_versa

    Parameters:
        coords (tuple):
            The input coordindates to transform
        image (str):
            The full path to the image
        to_pix (bool):
            Set to convert to pixel coordinates
        to_radec (bool):
            Set to convert to RA/Dec coordinates

    Returns:
        newcoords (tuple):
            Tuple of either (x, y) pixel coordinates
            or (RA, Dec) coordinates
    """

    try:
        wcs = getWCSFromPng(image)
    except Exception as e:
        raise MarvinError('Cannot get wcs info from image {0}: {1}'.format(image, e))

    if to_radec:
        try:
            newcoords = wcs.all_pix2world([coords], 1)[0]
        except AttributeError as e:
            raise MarvinError('Cannot convert coords to RA/Dec.  No wcs! {0}'.format(e))
    if to_pix:
        try:
            newcoords = wcs.all_world2pix([coords], 1)[0]
        except AttributeError as e:
            raise MarvinError('Cannot convert coords to image pixels.  No wcs! {0}'.format(e))
    return newcoords 

def populate_datamodel(self, array):
        """Place the image and accopanying WCS information in an
        ImageModel instance. This makes passing the information to
        the blotting function easier.

        Parameters
        ----------
        array : numpy.ndarray
            2D array containing the cropped image
        """

        # now we need to update the WCS of the mosaic piece
        cropped = datamodels.ImageModel()
        cropped.meta.wcsinfo.cdelt1 = self.mosaic_scale_x / 3600.
        cropped.meta.wcsinfo.cdelt2 = self.mosaic_scale_y / 3600.
        cropped.meta.wcsinfo.crpix1 = self.crpix1
        cropped.meta.wcsinfo.crpix2 = self.crpix2
        cropped.meta.wcsinfo.crval1 = self.center_ra
        cropped.meta.wcsinfo.crval2 = self.center_dec
        cropped.meta.wcsinfo.ctype1 = self.mosaic_x_type
        cropped.meta.wcsinfo.ctype2 = self.mosaic_y_type
        cropped.meta.wcsinfo.cunit1 = 'deg'
        cropped.meta.wcsinfo.cunit2 = 'deg'
        cropped.meta.wcsinfo.ra_ref = self.center_ra
        cropped.meta.wcsinfo.dec_ref = self.center_dec
        cropped.meta.wcsinfo.roll_ref = self.mosaic_roll * 180./np.pi
        cropped.meta.wcsinfo.wcsaxes = 2
        cropped.meta.wcsinfo.pc1_1 = self.cd11 / (self.mosaic_scale_x / 3600.)
        cropped.meta.wcsinfo.pc1_2 = self.cd12 / (self.mosaic_scale_x / 3600.)
        cropped.meta.wcsinfo.pc2_1 = self.cd21 / (self.mosaic_scale_y / 3600.)
        cropped.meta.wcsinfo.pc2_2 = self.cd22 / (self.mosaic_scale_y / 3600.)

        cropped.data = array
        return cropped 

def add_ra_dec_to_catalog(image):
    image_wcs = WCS(image.meta)
    image_catalog = image.catalog
    ras, decs = image_wcs.all_pix2world(image_catalog['x'], image_catalog['y'], 1)

    image_catalog['ra'] = ras
    image_catalog['dec'] = decs
    image_catalog['ra'].unit = 'degree'
    image_catalog['dec'].unit = 'degree'
    image_catalog['ra'].description = 'Right Ascension'
    image_catalog['dec'].description = 'Declination'