#!/usr/bin/env python

#filename=channel_maps_v2

#---------------------------------------by Gao John

#this file is to generate the integretion of the cube in CO observation from velocity argv[1] to some velocity by 1km/s

from __future__ import division

from pylab import *

import pyfits

#import pywcs

from kapteyn import wcs

from sys import argv

from string import atoi,atof

from matplotlib.axes import Axes#

from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable#

import pywcsgrid2

import mpl_toolkits.axes_grid1.axes_grid as axes_grid

from mpl_toolkits.axes_grid1.inset_locator import inset_axes

import matplotlib.colors as mcolors

use_path_effect = True

try:

from matplotlib.patheffects import withStroke

except ImportError:

use_path_effect = False

import gaoyuan

#define class Velo

#class Velo(object):

# def __init__(self, header):

# wcs = pywcs.WCS(header)

# self.wcs_vel = wcs.sub([3])

#

# def to_vel(self, p):

# v = self.wcs_vel.wcs_pix2sky([[p]], 0)

# return v[0][0]

#function setup_axes()

def setup_axes(fig, header):

return g, cax

#================================================

#h1 file: kes41

file_h1 = pyfits.open('..//kes41_13coall_MEANeff_based.fits')

h1_data = file_h1[0].data

h1_header = file_h1[0].header

h1_cube = gaoyuan.Cube(h1_data, h1_header).cube

(nn, ny, nx) = h1_cube.shape

#=================================================

#radio file: kes41

file_radio = pyfits.open('..//most//MOS336p5_smaller.fits')

radio_data = file_radio[0].data

radio_header = file_radio[0].header

#==================================================

vcut1 = atof(argv[1])

vcut_for_name = vcut1

#vcut2 = atof(argv[2])

#if (vcut1>vcut2):

# vcut1 = atof(argv[2])

# vcut2 = atof(argv[1])

#

#vn = ((arange(nn)-co_header['CRPIX3'])*co_header['CDELT3']+co_header['CRVAL3'])/1000.

#============================================

#prepare figure & axes

fig = figure(1, figsize=(12,16))

g, cax = setup_axes(fig, h1_header)

#draw image

i = 0

#dxy = kcut1

#nxy = 3*3

#cmap = cm.ocean_r

#cmap = cm.gist_ncar

#cmap = cm.Greys

#cmap = cm.jet

#cmap = cm.gist_yarg

cmap = cm.binary

#cmap = cm.YlOrRd

#norm = mcolors.Normalize()

images = []

#start_channel = i*nxy+dxy

proj = wcs.Projection(h1_header)

line_ax = proj.specaxnum

line = proj.sub(line_ax)

del_v = 2

#unit change

vcut1 = vcut1*1e3

del_v = del_v*1e3

for i, ax in enumerate(g):

# channel_number = start_channel + i

# channel = co_data[channel_number]

# k = 0

# kcut1 = 0

# kcut2 = int(co_header['CRPIX3'])

# while (k<nn-1):

# if (vn[k]<=vcut1 and vn[k+1]>=vcut1):# or (vn[k]>=-200 and vn[k+1]<=-200)):

# kcut1 = k

# break

# k+=1

# k = 0

# while (k<nn):

# if (vn[k]<=vcut1+del_v and vn[k+1]>=vcut1+del_v):

# kcut2 = k

# break

# k+=1

#

kcut1 = int(line.topixel1d(vcut1))

vcut2 = vcut1+del_v

kcut2 = int(line.topixel1d(vcut2))

co_per1 = zeros((ny,nx))

for ii in range(nx):

for jj in range(ny):

# #ns = 0.

for kk in range(nn)[kcut1:kcut2+1]:

if (isnan(h1_cube[kk,jj,ii])):

print 'nan found'

continue

co_per1[jj,ii] += h1_cube[kk,jj,ii]

im = ax.imshow(co_per1, origin="lower", cmap=cmap,

vmin=12, vmax=58,

interpolation="none")

#draw contour

cont = ax[radio_header].contour(radio_data, linewidth=0.5,

levels=arange(0.07,0.88,0.12), colors="c", alpha=0.8)

v = vcut1/1e3

t = ax.add_inner_title(r"$%4.1f$" % (v), loc=3, frameon=False,

prop={'size':16})

#maser point

m_alpha = 15.*(16.+(38.+52.15/60.)/60)

m_delta = -(46.+(56.+16.1/60)/60)

plt = ax["fk5"].plot([m_alpha], [m_delta], color="#fa626e", marker="+", ms=20, mew=2, zorder=3)

if use_path_effect:

t.txt._text.set_path_effects([withStroke(foreground='w',

linewidth=1)])

ax.axis["bottom"].major_ticklabels.set(fontsize=18)

ax.axis["bottom"].label.set(fontsize=18)

ax.axis["left"].major_ticklabels.set(fontsize=18)

ax.axis["left"].label.set(fontsize=18)

images.append(im)

vcut1 += del_v

#label with velocities

#for i, ax in enumerate(g):

# channel_number = start_channel+i

# v = vn[channel_number]

# t = ax.add_inner_title(r"$v=%4.1f$ km s$^{-1}$" % (v), loc=2, frameon=False)

# if use_path_effect:

# t.txt._text.set_path_effects([withStroke(foreground='w',

# linewidth=3)])

#make colorbar

cb = colorbar(im, cax=cax)

cb.set_label("K",fontsize=18)

for t in cb.ax.get_yticklabels():

t.set_fontsize(18)

#cb.set_ticks([28,36,44,52])

#adjust norm

#norm.vmin = 125

#norm.vmax = 300

#for im in images:

# im.changed()

savefig("co13_channel_maps_%d"%vcut_for_name+"_v2.eps",

bbox_inches="tight", dpi=300, format="eps")