How do I trim a .fits image and keep world coordinates for plotting in astropy Python?

Question

This issue has been plaguing me for some time. I'm trying to handle some large amount of data that is in the form of a .fits file (on the order of 11000x9000 pixels). What I need to do is create a 'zoomed in' RA/Dec coordinate plot (ideally using astropy.wcs) for many objects on the sky with contours from one fits file, and greyscale (or heatmap from another).

My problem is that whenever I slice the data from the image (to my region of interest) I lose the association with the sky coordinates. This means that the sliced image isn't in the correct location.

I've adapted an example from the astropy docs to save you the pain of my data. (Note: I want the contours to cover more area than the image, whatever the solution is for this should work on both data)

Here is the code I am having trouble with:

from matplotlib import pyplot as plt
from astropy.io import fits
from astropy.wcs import WCS
from astropy.utils.data import download_file
import numpy as np

fits_file = 'http://data.astropy.org/tutorials/FITS-images/HorseHead.fits'
image_file = download_file(fits_file, cache=True)
hdu = fits.open(image_file)[0]
wmap = WCS(hdu.header)
data = hdu.data

fig = plt.figure()
ax1 = fig.add_subplot(121, projection=wmap)
ax2 = fig.add_subplot(122, projection=wmap)
# Scale input image
bottom, top = 0., 12000.
data = (((top - bottom) * (data - data.min())) / (data.max() - data.min())) + bottom


'''First plot'''
ax1.imshow(data, origin='lower', cmap='gist_heat_r')

# Now plot contours
xcont = np.arange(np.size(data, axis=1))
ycont = np.arange(np.size(data, axis=0))
colors = ['forestgreen','green', 'limegreen']
levels = [2000., 7000., 11800.]

ax1.contour(xcont, ycont, data, colors=colors, levels=levels, linewidths=0.5, smooth=16)
ax1.set_xlabel('RA')
ax1.set_ylabel('Dec')
ax1.set_title('Full image')

''' Second plot ''' 
datacut = data[250:650, 250:650]
ax2.imshow(datacut, origin='lower', cmap=cmap)
ax2.contour(xcont, ycont, data, colors=colors, levels=levels, linewidths=0.5, smooth=16)

ax2.set_xlabel('RA')
ax2.set_ylabel('')
ax2.set_title('Sliced image')
plt.show()

I tried using the WCS coords of my sliced chunk to fix this, but I'm not sure if I can pass it in anywhere!

pixcoords = wcs.wcs_pix2world(zip(*[range(250,650),range(250,650)]),1)

Cobaltous

Answer 1

The good news is: You can simply slice your astropy.WCS as well which makes your task relativly trivial:

...

wmapcut = wmap[250:650, 250:650] # sliced here
datacut = data[250:650, 250:650]
ax2 = fig.add_subplot(122, projection=wmapcut) # use sliced wcs as projection
ax2.imshow(datacut, origin='lower', cmap='gist_heat_r')
# contour has to be sliced as well
ax2.contour(np.arange(datacut.shape[0]), np.arange(datacut.shape[1]), datacut, 
            colors=colors, levels=levels, linewidths=0.5, smooth=16)
...

If your files have different WCS you might need to do some reprojection (see for example reproject)

Lurie

This is an excellent concise answer (using .shape[0] is super tidy!). Thank you for pointing out reproject, I didn't know it existed and I will definitely need it! One more thing; I'm trying to scatter a point on top of this and the axes coordinates don't appear to be mapped in degrees. For example plotting ax1.scatter(85.33, -2.5) should appear at (85 20' , -2 30') but it's a fair way off. – Cobaltous

pywcsgrid2 also has lots of good reprojection capabilities, and I find it a little more flexible than astropy.WCS, at the cost of a steeper learning curve. It can also overplot reasonably well. – Urbana

@Cobaltous To overlay scatter plots check out this example: wcsaxes.readthedocs.io/en/latest/…. If you have any trouble afterwards just ask another question (it's easier to deal with one problem per question). In short you need to apply a transform or specify the coordinates in pixel-coordinates. – Lurie