sunpy__plot.py

"""
routines for plotting images from the SUNRISE data output


Example usage:

import sunpy.sunpy__plot as sunplot

sunplot.plot_synthetic_sdss_gri('broadband_1234.fits')
sunplot.plot_sdss_gri('broadband_1234.fits')

Dependencies:
  numpy
  matplotlib
"""

import numpy as np
import os
import sys
import sunpy.sunpy__load as sunpy__load			# used for noiseless images, for which we can return the image input directly
import sunpy__synthetic_image		# used for images with noise, pixel scaling, etc.
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import gc

__author__ = "Paul Torrey and Greg Snyder"
__copyright__ = "Copyright 2014, The Authors"
__credits__ = ["Paul Torrey", "Greg Snyder"]
__license__ = "GPL"
__version__ = "1.0"
__maintainer__ = "Paul Torrey"
__email__ = "ptorrey@mit.harvard.edu"
__status__ = "Production"
if __name__ == '__main__':    #code to execute if called from command-line
    pass    #do nothing 



def plot_spectrum(filelist, savefile='spectrum.pdf', fontsize=14, ymin=None, ymax=None, **kwargs):
    """ routine for plotting the spectra for a list of files """
    fig = plt.figure(figsize=(5,5))
    ax = fig.add_subplot(111)
    print filelist
    for file in filelist:
	print file
	wavelength = sunpy__load.load_sed_lambda(file)
        sed        = sunpy__load.load_sed_l_lambda(file)
        ax.plot(wavelength,wavelength * sed, label=file)

    ax.set_xlabel(r'$\lambda$ (m)',fontsize=fontsize)
    ax.set_ylabel(r'$\lambda$L${}_{\lambda}$ (W)', fontsize=fontsize)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize)

    ax.legend
    ax.set_xscale('log')
    ax.set_yscale('log')
    if (ymin != None) and (ymax != None):
	ax.set_ylim([ymin, ymax])

    fig.subplots_adjust(left=0.18, right=0.95, top=0.95, bottom=0.12, wspace=0.0, hspace=0.0)
    fig.savefig(savefile)





def plot_synthetic_sdss_gri(filename, savefile='syn_sdss_gri.png', **kwargs):
    """ routine for plotting synthetic sdss gri images from Illustris idealized images including appropriate pixel scaling, noise, etc.  """

    rp, img = return_synthetic_sdss_gri_img(filename, **kwargs)
    my_save_image(img, savefile)
    del img
    gc.collect()

def plot_synthetic_hst(filename, savefile='syn_hst.png', **kwargs):	#mass_string=None, full_fov=None, cut_bad_pixels=False, **kwargs):
    """ routine for plotting synthetic hst gri images from Illustris idealized images including appropriate pixel scaling, noise, etc.  """

    rp, img = return_synthetic_hst_img(filename, **kwargs)	#return_synthetic_sdss_gri_img(filename, **kwargs)
    my_save_image(img, savefile, **kwargs)	#top_opt_text=mass_string, full_fov=full_fov, cut_bad_pixels=cut_bad_pixels, **kwargs)
    del img
    gc.collect()


def plot_sdss_gri(filename, savefile='./sdss_gri.png', **kwargs):
    """ routine for plotting synthetic sdss gri images from Illustris idealized images *without* additional image effects """

    img = return_sdss_gri_img(filename, **kwargs)
    my_save_image(img, savefile)
    del img
    gc.collect()


def return_synthetic_sdss_gri_img(filename, 
				lupton_alpha=0.5, lupton_Q=0.5, scale_min=1e-4, 
                                b_fac=0.7, g_fac=1.0, r_fac=1.3,
				seed_boost=1.0,
			 	r_petro_kpc=None,
				**kwargs):

    fail_flag=True		# looks for "bad" backgrounds, and tells us to try again
    n_iter = 1
    while(fail_flag and (n_iter < 2)):
        fail_flag=False
	try:
            seed=int(filename[filename.index('broadband_')+10:filename.index('.fits')])*(n_iter)*seed_boost
        except:
	    try:
                seed=int(filename[filename.index('.fits')-3:filename.index('.fits')])*(n_iter)*seed_boost
	    except:
		seed=1234

        n_iter+=1

        # The calls to build_synthetic_image had to be updated to account for changes in the number
        # of return values (previously 4, now 6).  We don't need the other 2 values here, but we
        # have to catch them so the code won't complain about too many return values to unpack.
        b_image, rp, the_used_seed,this_fail_flag, _, _    = \
            sunpy__synthetic_image.build_synthetic_image(filename, 'g_SDSS.res', 
				seed=seed,
				r_petro_kpc=r_petro_kpc, 
				fix_seed=False,
				**kwargs)
        if(this_fail_flag):
	  fail_flag=True

        g_image, dummy, the_used_seed,this_fail_flag, _, _ = \
            sunpy__synthetic_image.build_synthetic_image(filename, 'r_SDSS.res', 
				seed=the_used_seed,
				r_petro_kpc=rp,
				fix_seed=True, 
				**kwargs)
        if(this_fail_flag):
          fail_flag=True

        r_image, dummy, the_used_seed, this_fail_flag, _, _ = \
            sunpy__synthetic_image.build_synthetic_image(filename, 'i_SDSS.res', 
                                seed=the_used_seed,
				r_petro_kpc=rp,
				fix_seed=True, 
				**kwargs)
        if(this_fail_flag):
            fail_flag=True

    n_pixels = r_image.shape[0]
    img = np.zeros((n_pixels, n_pixels, 3), dtype=float)
 
    b_image *= b_fac
    g_image *= g_fac
    r_image *= r_fac
 
    I = (r_image + g_image + b_image)/3
    val = np.arcsinh( lupton_alpha * lupton_Q * (I - scale_min))/lupton_Q
    I[ I < 1e-6 ] = 1e100		# from below, this effectively sets the pixel to 0

    img[:,:,0] = r_image * val / I
    img[:,:,1] = g_image * val / I
    img[:,:,2] = b_image * val / I

    maxrgbval = np.amax(img, axis=2)
 
    changeind = maxrgbval > 1.0
    img[changeind,0] = img[changeind,0]/maxrgbval[changeind]
    img[changeind,1] = img[changeind,1]/maxrgbval[changeind]
    img[changeind,2] = img[changeind,2]/maxrgbval[changeind]

    minrgbval = np.amin(img, axis=2)
    changeind = minrgbval < 0.0
    img[changeind,0] = 0
    img[changeind,1] = 0
    img[changeind,2] = 0

    changind = I < 0
    img[changind,0] = 0
    img[changind,1] = 0
    img[changind,2] = 0
    img[img<0] = 0


    del b_image, g_image, r_image, I, val
    gc.collect()

    img[img<0] = 0
    return rp, img


def return_synthetic_hst_img(filename,
                                lupton_alpha=0.5, lupton_Q=0.5, scale_min=1e-4,
                                b_fac=1.0, g_fac=1.0, r_fac=1.0, max=1.0, dynrng=1e3,
                                **kwargs):

    try:
       seed=int(filename[filename.index('broadband_')+10:filename.index('.fits')])
    except:
       try:
          seed=int(filename[filename.index('broadband_red_')+14:filename.index('.fits')])
       except:
          seed=int(filename[filename.index('broadband_rest_')+15:filename.index('.fits')])
    

    b_image, rp, dummy, dummy    = sunpy__synthetic_image.build_synthetic_image(filename, 22,		#25,
                                seed=seed, fix_seed=True,
                                #r_petro_kpc=None,
                                **kwargs)
    dummy, dummy, dummy, dummy  = sunpy__synthetic_image.build_synthetic_image(filename, 25,
                                seed=seed, fix_seed=True,
                                #r_petro_kpc=rp,
                                **kwargs)

    g_image, dummy, dummy, dummy = sunpy__synthetic_image.build_synthetic_image(filename, 26,
                                seed=seed, fix_seed=True,
                                #r_petro_kpc=rp,
                                **kwargs)
    r_image, dummy, dummy, dummy = sunpy__synthetic_image.build_synthetic_image(filename, 27,
                                seed=seed, fix_seed=True,
                                #r_petro_kpc=rp,
                                **kwargs)

    n_pixels = r_image.shape[0]
    img = np.zeros((n_pixels, n_pixels, 3), dtype=float)

    b_image *= b_fac
    g_image *= g_fac
    r_image *= r_fac

    if 0:
       I = (r_image + g_image + b_image)/3
       val = np.arcsinh( lupton_alpha * lupton_Q * (I - scale_min))/lupton_Q
       I[ I < 1e-8 ] = 1e20               # from below, this effectively sets the pixel to 0

       img[:,:,0] = r_image * val / I
       img[:,:,1] = g_image * val / I
       img[:,:,2] = b_image * val / I

       maxrgbval = np.amax(img, axis=2)

       changeind = maxrgbval > 1.0
       img[changeind,0] = img[changeind,0]/maxrgbval[changeind]
       img[changeind,1] = img[changeind,1]/maxrgbval[changeind]
       img[changeind,2] = img[changeind,2]/maxrgbval[changeind]

       minrgbval = np.amin(img, axis=2)
       changeind = minrgbval < 0.0
       img[changeind,0] = 0
       img[changeind,1] = 0
       img[changeind,2] = 0

       changind = I < 0
       img[changind,0] = 0
       img[changind,1] = 0
       img[changind,2] = 0
       img[img<0] = 0

    else:
        img[:,:,0] = np.log10(r_image )  # / max=1.0, dynrng=1e3,
	img[:,:,1] = np.log10(g_image )
	img[:,:,2] = np.log10(b_image )

        img -= np.log10( max/dynrng ) 
        img /= np.log10( dynrng ) 

        img[img>1] = 1
	I = img
	val = img

    img[img<0] = 0
    print "img min/max/mean "+str(img.min())+"  "+str(img.max())+"  "+str(img.mean())
    print " "

    del b_image, g_image, r_image, I, val
    gc.collect()

    img[img<0] = 0
    return rp, img


def return_sdss_gri_img(filename,camera=0,scale_min=0.1,scale_max=50,size_scale=1.0, non_linear=0.5):
    if (not os.path.exists(filename)):
        print "file not found:", filename
        sys.exit()

    b_image = sunpy__load.load_broadband_image(filename,band='g_SDSS.res',camera=camera) * 0.7
    g_image = sunpy__load.load_broadband_image(filename,band='r_SDSS.res',camera=camera) * 1.0
    r_image = sunpy__load.load_broadband_image(filename,band='i_SDSS.res',camera=camera) * 1.4
    n_pixels = r_image.shape[0]
    img = np.zeros((n_pixels, n_pixels, 3), dtype=float)

    img[:,:,0] = asinh(r_image, scale_min=scale_min, scale_max=scale_max,non_linear=non_linear)
    img[:,:,1] = asinh(g_image, scale_min=scale_min, scale_max=scale_max,non_linear=non_linear)
    img[:,:,2] = asinh(b_image, scale_min=scale_min, scale_max=scale_max,non_linear=non_linear)
    img[img<0] = 0

    del b_image, g_image, r_image
    gc.collect()

    return img


def return_h_band_img(filename,camera=0,scale_min=0.1,scale_max=50,size_scale=1.0):
    if (not os.path.exists(filename)):
        print "file not found:", filename
        sys.exit()

    image = sunpy__load.load_broadband_image(filename,band='H_Johnson.res', camera=camera) 
    n_pixels = image.shape[0]
    img = np.zeros((n_pixels, n_pixels), dtype=float)
    img[:,:] = asinh(image, scale_min=scale_min, scale_max=scale_max,non_linear=0.5)
    img[img<0] = 0
    return img


def return_johnson_uvk_img(filename,camera=0,scale_min=0.1,scale_max=50,size_scale=1.0):
    if (not os.path.exists(filename)):
        print "file not found:", filename
        sys.exit()

    b_effective_wavelength = sunpy__load.load_broadband_effective_wavelengths(filename,band="U_Johnson.res")
    g_effective_wavelength = sunpy__load.load_broadband_effective_wavelengths(filename,band="V_Johnson.res")
    r_effective_wavelength = sunpy__load.load_broadband_effective_wavelengths(filename,band="K_Johnson.res")

    b_image = sunpy__load.load_broadband_image(filename,band='U_Johnson.res',camera=camera) * b_effective_wavelength / g_effective_wavelength * 2.5
    g_image = sunpy__load.load_broadband_image(filename,band='V_Johnson.res',camera=camera) * g_effective_wavelength / g_effective_wavelength 
    r_image = sunpy__load.load_broadband_image(filename,band='K_Johnson.res',camera=camera) * r_effective_wavelength / g_effective_wavelength * 1.5 

    n_pixels
    img = np.zeros((n_pixels, n_pixels, 3), dtype=float)
    img[:,:,0] = asinh(r_image, scale_min=scale_min, scale_max=scale_max,non_linear=0.5)
    img[:,:,1] = asinh(g_image, scale_min=scale_min, scale_max=scale_max,non_linear=0.5)
    img[:,:,2] = asinh(b_image, scale_min=scale_min, scale_max=scale_max,non_linear=0.5)
    img[img<0] = 0
    return img


def return_stellar_mass_img(filename, camera=0, scale_min=1e8, scale_max=1e10, size_scale=1.0, non_linear=1e8):
    image = sunpy__load.load_stellar_mass_map(filename, camera=camera)
    n_pixels = image.shape[0]
    img = np.zeros((n_pixels, n_pixels), dtype=float)
    img[:,:] = asinh(image, scale_min=scale_min, scale_max=scale_max, non_linear=non_linear)
    img[img<0] = 0
    return img

def return_mass_weighted_age_img(filename, camera=0, scale_min=None, scale_max=None, size_scale=1.0):
    image = sunpy__load.load_mass_weighted_stellar_age_map(filename, camera=camera)
    image += 1e5
    return image

def return_stellar_metal_img(filename, camera=0, scale_min=None, scale_max=None, size_scale=1.0, non_linear=None):
    image1 = sunpy__load.load_stellar_mass_map(filename, camera=camera)
    image2 = sunpy__load.load_stellar_metal_map(filename, camera=camera)
    image = image2/image1
    image[image<0]      = 0     #image.min()
    image[image*0 != 0] = 0     #image.min()
    return image

def my_save_image(img, savefile, opt_text=None, top_opt_text=None, full_fov=None, cut_bad_pixels=False, zoom=None, save_indiv_bands=True, **kwargs):
    if img.shape[0] >1:
	n_pixels_save = img.shape[0]


	if zoom!=None:
	    n_pixels_current = img.shape[0]
	    center = np.floor(img.shape[0]/(2.0))
	    n_pixels_from_center  = np.floor(img.shape[0]/(2.0*zoom))
	    img = img[center - n_pixels_from_center:center+n_pixels_from_center, center - n_pixels_from_center:center+n_pixels_from_center]
	    print "resized image from "+str(n_pixels_current)+" to "+str(img.shape[0])+" pixels"
	    if full_fov!=None:
		full_fov /= (1.0*zoom)

	if cut_bad_pixels:
    	    cut_index=-1
	    print np.mean(img[cut_index:, cut_index:])
	    while( np.mean(img[cut_index:, cut_index:]) == 0 ):
	        print cut_index
	        cut_index -= 1
	    img = img[:cut_index,:cut_index]


        fig = plt.figure(figsize=(1,1))
        ax = fig.add_subplot(111)
        imgplot = ax.imshow(img,origin='lower', interpolation='nearest')
        plt.axis('off')

        if not opt_text==None:
            ax.text(img.shape[0]/2.0, img.shape[0]/2.0, opt_text, ha='center',va='center', color='white', fontsize=4)

        if not top_opt_text==None:
            ax.text(img.shape[0]/2.0, 0.9*img.shape[0], top_opt_text, ha='center',va='center', color='white', fontsize=4)

	if not full_fov==None:
	    bar_size_in_kpc    = np.round(full_fov/5.0)
	    pixel_size_in_kpc  = full_fov / (1.0 * img.shape[0])
	    bar_size_in_pixels = bar_size_in_kpc / pixel_size_in_kpc
	    center = img.shape[0]/2.0
	    
	    ax.text( center, 0.15 * img.shape[0], str(bar_size_in_kpc)+" kpc", color='w', fontsize=4, ha='center', va='center')
	    ax.plot( [center-bar_size_in_pixels/2.0, center+bar_size_in_pixels/2.0] , [0.1*img.shape[0], 0.1*img.shape[0]], lw=2,color='w')
	    ax.set_xlim([0,img.shape[0]-1])
	    ax.set_ylim([0,img.shape[0]-1])


        fig.subplots_adjust(left=0.0, right=1.0, top=1.0, bottom=0.0)
        fig.savefig(savefile, dpi=n_pixels_save)
        fig.clf()
        plt.close()


	if save_indiv_bands:
	    print img.shape
	    for iii in np.arange(3):
		fig = plt.figure(figsize=(1,1))
                ax = fig.add_subplot(111)
		print img[:,:,iii].min(), img[:,:,iii].max()

                imgplot = ax.imshow(img[:,:,iii],origin='lower', interpolation='nearest', cmap = cm.Greys, vmin=0, vmax=1)
                plt.axis('off')
                fig.subplots_adjust(left=0.0, right=1.0, top=1.0, bottom=0.0)
                fig.savefig(savefile+"_band_"+str(iii)+'.png', dpi=n_pixels_save)
                fig.clf()
                plt.close()
        del img
        gc.collect()





def asinh(inputArray, scale_min=None, scale_max=None, non_linear=2.0):
        imageData=np.array(inputArray, copy=True)

        if scale_min == None:
                scale_min = imageData.min()
        if scale_max == None:
                scale_max = imageData.max()
        factor = np.arcsinh((scale_max - scale_min)/non_linear)
        indices0 = np.where(imageData < scale_min)
        indices1 = np.where((imageData >= scale_min) & (imageData <= scale_max))
        indices2 = np.where(imageData > scale_max)
        imageData[indices0] = 0.0
        imageData[indices2] = 1.0
        imageData[indices1] = np.arcsinh((imageData[indices1] - scale_min)/non_linear)/factor

        return imageData