wcsaxes_coords.py

# -*- coding: utf-8 -*-
"""
Created on Mon Mar 24 17:37:51 2014

@author: stuart
"""
import numpy as np
import matplotlib.pyplot as plt

import astropy.units as u
from astropy.wcs import WCS
from astropy.io import fits
import wcsaxes
import wcsaxes.utils
import sunpy.map

import sunpy.coordinates

hmimap = sunpy.map.Map(sunpy.AIA_171_IMAGE)
#hmimap = sunpy.map.Map('hmi_m_45s_2014_03_01_00_01_30_tai_magnetogram_fits.fits')
#hmimap = sunpy.map.Map('/home/stuart/sunpy/data/aia_lev1_171a_2014_01_03t10_00_11_34z_image_lev1_fits')
#hmimap = sunpy.map.Map('/storage2/Lth_AR_Project/HMI/AR11148/cropped/HMI_2011-01-16T02:00:25.40_centred_131.3x-382.8.fits')

w2 = WCS(naxis=2)
w2.wcs.crpix = [hmimap.reference_pixel['x'], hmimap.reference_pixel['y']]
w2.wcs.cdelt = [hmimap.scale['x'], hmimap.scale['y']]
w2.wcs.crval = [hmimap.reference_coordinate['x'], hmimap.reference_coordinate['y']]
w2.wcs.ctype = [hmimap.coordinate_system['x'], hmimap.coordinate_system['y']]
w2.wcs.pc = hmimap.rotation_matrix
w2.wcs.cunit = [hmimap.units['x'], hmimap.units['y']]


def hpln_hplt_frame(wcs):
    xcoord = wcs.wcs.ctype[0][0:4]
    ycoord = wcs.wcs.ctype[1][0:4]

    if xcoord == 'HPLN' and ycoord == 'HPLT':
        return sunpy.coordinates.HelioProjective

wcsaxes.utils.register_frame_identifier(hpln_hplt_frame)


# Make colorscale figure
fig = plt.figure(figsize=(6,6))
 
# Create axes
ax = wcsaxes.WCSAxes(fig, [0.1, 0.1, 0.8, 0.8], wcs=w2)
fig.add_axes(ax)
 
# Show colormap
hmimap.plot(axes=ax, extent=None, vmin=np.nanmin(hmimap), vmax=np.nanmax(hmimap)*0.3)
 
# Esthetics 
x = ax.coords[0]#['hpln']
y = ax.coords[1]#['hplt']

#x.coord_wrap = 180

x.set_ticks(color='white')
y.set_ticks(color='white')
# 
x.set_ticklabel(size='x-small')
y.set_ticklabel(size='x-small')
 
x.set_axislabel('Solar-x [arcsec]', weight='bold')
y.set_axislabel('Solar-y [arcsec]', weight='bold')

 
x.set_ticks_position('bl')
y.set_ticks_position('bl')

#x.set_ticks(spacing=5 * u.deg)
#y.set_ticks(spacing=5 * u.deg)

x.set_major_formatter('s.s')
y.set_major_formatter('s.s')

 
ax.coords.grid(color='white', alpha=0.1)
 
# Show overlay with solar lon/lat
 
overlay = ax.get_coords_overlay('heliographicstonyhurst')
print overlay
 
lon = overlay[0]
lat = overlay[1]
 
lon.coord_wrap = 180
lon.set_major_formatter('dd')
 
lon.set_ticklabel(size='x-small')
lat.set_ticklabel(size='x-small')
 
lon.set_axislabel('Solar Longitude', weight='bold')
lat.set_axislabel('Solar Latitude', weight='bold')
 
lon.set_ticks_position('tr')
lat.set_ticks_position('tr')

lon.set_ticks(spacing=5. * u.deg, color='white')
lat.set_ticks(spacing=5. * u.deg, color='white')
 
overlay.grid(color='white', alpha=0.3)
 
# Set title
ax.set_title("{} {}".format(hmimap.name, sunpy.time.parse_time(hmimap.date).strftime("%Y-%m-%d %H:%M:%S.%f")),
             y=1.05, size='large', weight='bold')
 
plt.show()