WCS as matplotlib projection for a data-cube slice loaded using astropy?

Question

I have a FITS file named 'my_cube.fits' with WCS. The file has spatial information on axis 1 and 2 (X and Y) and spectral information on axis 3 (Z). When I load it using astropy.io.fits, the spectral axis is 0 and spatial axes are 1 and 2. The file is loaded like this:

    import astropy.io.fits as pyfits

    filename = 'my_cube.fits'
    my_data = pyfits.getdata(filename)
    my_header = pyfits.getheader(filename)

I have been using matplotlib to display data and I want to know how to display a single spectral frame of my data cube using its WCS. Let's say:

    from astropy.wcs import WCS
    from matplotlib import pyplot as plt

    my_wcs = WCS(my_header)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection=my_wcs)
    ax.imshow(my_data[5, :, :])

    plt.show()

If I just do that, I have:

    ...
    File "/usr/local/lib/python3.4/dist-packages/matplotlib/figure.py", line 1005, in add_subplot
    a = subplot_class_factory(projection_class)(self, *args, **kwargs)
    File "/usr/local/lib/python3.4/dist-packages/matplotlib/axes/_subplots.py", line 73, in __init__
    self._axes_class.__init__(self, fig, self.figbox, **kwargs)
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/core.py", line 49, in __init__
    self.patch = self.coords.frame.patch
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/frame.py", line 129, in patch
    self._update_patch_path()
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/frame.py", line 115, in _update_patch_path
    self.update_spines()
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/frame.py", line 192, in update_spines
    self['b'].data = np.array(([xmin, ymin], [xmax, ymin]))
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/frame.py", line 40, in data
    self._world = self.transform.transform(self._data)
    File "/usr/local/lib/python3.4/dist-packages/wcsaxes/transforms.py", line 166, in transform
    world = self.wcs.wcs_pix2world(pixel_full, 1)
    File "/usr/local/lib/python3.4/dist-packages/astropy/wcs/wcs.py", line 1329, in wcs_pix2world
'output', *args, **kwargs)
    File "/usr/local/lib/python3.4/dist-packages/astropy/wcs/wcs.py", line 1231, in _array_converter
    return _return_single_array(xy, origin)
    File "/usr/local/lib/python3.4/dist-packages/astropy/wcs/wcs.py", line 1212, in _return_single_array
"of shape (N, {0})".format(self.naxis))
    ValueError: When providing two arguments, the array must be of shape (N, 3)

Citreous

Answer 1

Your WCS object contains the two spatial axes and the spectral axis, since you initialized it with the full header (I assume your header also contains a spectral WCS solution). Thus, creating a 2D plot with just the spatial won't work, since your projection given to the subplot is 3D.

I haven't done this myself, nor do I have an example file, but the WCS documentation mentions that you can use the sub function, or even the celestial functions to get the individual axes or the celestial (spatial) axes; these functions will return a WCS object with just those axes.

Thus, you can probably use the following:

my_wcs = WCS(my_header).celestial
fig = plt.figure()
ax = fig.add_subplot(111, projection=my_wcs)

Batton

celestial is a property, not a function, so it doesn't need to be called – Cupro

That will require a fix in the docs then; the API documentation is correct, but the part I link to calls it a function. – Batton

Answer 2

This is a good use-case for spectral-cube, which effectively wraps astropy.io.fits for cube uses.

Evert's solution, once corrected, will work, assuming you have wcsaxes installed.

To use spectral_cube for this, a simple example is:

cube = SpectralCube.read(filename)
cube[5,:,:].quicklook() # uses aplpy

# using wcsaxes
fig = plt.figure()
ax = fig.add_axes([0.1,0.1,0.8,0.8], projection=cube[5,:,:].wcs)
ax.imshow(cube[5,:,:]) # you may need cube[5,:,:].value depending on mpl version

Cupro