solarwcsaxes.py

import numpy as np
import matplotlib.pyplot as plt

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

from sunpy.net import vso
from sunpy.wcs import (convert_hcc_hpc, convert_hg_hcc,
convert_hcc_hg, convert_hpc_hcc)

from wcsaxes.transforms import CurvedTransform

class SolarHGtoHPC(CurvedTransform):

    def transform(self, input_coords):
        ilon, ilat = input_coords[:,0], input_coords[:,1]
        olon, olat = convert_hcc_hpc(*convert_hg_hcc(ilon, ilat))
        olon /= 3600.
        olat /= 3600
        return np.vstack([olon, olat]).transpose()

    transform_non_affine = transform

    def inverted(self):
        return SolarHPCtoHG()

class SolarHPCtoHG(CurvedTransform):

    def transform(self, input_coords):
        ilon, ilat = input_coords[:,0], input_coords[:,1]
        olon, olat = convert_hcc_hg(*convert_hpc_hcc(ilon * 3600., ilat * 3600., z=True))
        return np.vstack([olon, olat]).transpose()

    transform_non_affine = transform

    def inverted(self):
        return SolarHGtoHPC()

#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.crota = [hmimap.rotation_angle['x'], hmimap.rotation_angle['y']]
w2.wcs.cunit = [hmimap.units['x'], hmimap.units['y']]

# Make colorscale figure
fig = plt.figure(figsize=(6,6))
 
# Create axes
ax = 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(SolarHGtoHPC())
 
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()