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)

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)