solar_tom.py

# -*- coding: utf-8 -*-
"""
Created on Fri Jun 13 10:50:03 2014

@author: stuart
"""
import glob
import numpy as np

import matplotlib
import matplotlib.pyplot as plt

import astropy.units as u
from astropy.wcs import WCS

import sunpy.map

from wcsaxes import WCSAxes

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()

def get_wcs(smap):

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

    return w2

def draw_wcs(ax):
    """
    Make the WCAxes look pretty
    """
    x = ax.coords['hpln']
    y = ax.coords['hplt']

    x.coord_wrap = 180

    x.set_ticklabel(size='x-small', color='white')
    y.set_ticklabel(size='x-small', color='white')

    x.set_axislabel('Solar-x [arcsec]', weight='bold', color='white')
    y.set_axislabel('Solar-y [arcsec]', weight='bold', color='white')

    x.set_ticks(spacing=60. * u.arcsec, color='white')
    y.set_ticks(spacing=60. * u.arcsec, color='white')
    
    x.set_ticks_position('bl')
    y.set_ticks_position('bl')

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

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

    # ax.coords.grid(color='white', alpha=0.2, ls='solid')

    # Show overlay with solar lon/lat
    coord_meta = {'type': ['longitude', 'latitude'], 'wrap':[180, None], 'unit':['arcsec', 'arcsec'], 'name':['lon', 'lat']}
    overlay = ax.get_coords_overlay(SolarHGtoHPC().inverted(), coord_meta=coord_meta)

    lon = overlay[0]
    lat = overlay[1]

    lon.coord_wrap = 180
    lon.set_major_formatter('dd')

    lon.set_ticklabel(size='x-small', color='white')
    lat.set_ticklabel(size='x-small', color='white')
 
    lon.set_axislabel('Solar Longitude', weight='bold', color='white')
    lat.set_axislabel('Solar Latitude', weight='bold', color='white')

    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='black', alpha=0.8, linestyle='solid')


Writer = matplotlib.animation.writers['ffmpeg']
writer = Writer(fps=30, metadata=dict(artist='The SunPy Project'), bitrate=10000)

files = glob.glob('./cropped/*fits')
files.sort()
files = files[::10]

mapc = sunpy.map.Map(files, cube=True)

w2 = get_wcs(mapc.maps[0])

fig2 = plt.figure(figsize=(8,8))
ax = WCSAxes(fig2, [0.1, 0.1, 0.8, 0.8], wcs=w2)
fig2.patch.set_facecolor('black')
fig2.add_axes(ax)

def updatefig(i, im, ani_data, ax):
    print("Animating Frame: {}/{}".format(i, len(files)))
    amap = mapc[i]
    im.set_array(ani_data[i].data)
    # im.set_cmap(plt.cm.RdBu)
    from sunpy.cm import cm
    im.set_cmap(cm.cmlist.get('hmimag'))
    im.set_norm(mapc[0].mpl_color_normalizer)
    im.set_clim(vmin=-50, vmax=50)
    ax.set_title("%s %s" % (amap.name, amap.date), color='white')
    ax.reset_wcs(get_wcs(ani_data[i]))
    draw_wcs(ax)

im = mapc[0].plot(axes=ax, extent=None)
ax.title.set_position([0.5,1.07])

draw_wcs(ax)
ani = matplotlib.animation.FuncAnimation(fig2, updatefig,
                                        frames=range(0,len(mapc.maps)),
                                        fargs=[im,mapc, ax],
                                        interval=300,
                                        blit=False)

ani.save('HMI_cropped_wcs_axes_rot_hmimag.mp4', writer=writer, savefig_kwargs={'facecolor':'black'})
fig2.show()