viewtfmap.py 10.4 KB
Newer Older
OP's avatar
OP committed
1 2 3 4 5 6 7 8 9 10 11 12
import math as math
import sys, platform, os
import numpy as np
from astropy.io import fits

from matplotlib import pyplot as plt
import itertools as itert
from skimage import data, img_as_float
from skimage import exposure
import mynormalize as mynorm
from matplotlib.widgets import Slider, Button, RadioButtons
import time
perdereau's avatar
perdereau committed
13
from colorsys import hsv_to_rgb, hls_to_rgb
OP's avatar
OP committed
14 15
#    Choices: ['auto', 'gtk', 'gtk3', 'inline', 'nbagg', 'notebook', 'osx', 'qt', 'qt4', 'qt5', 'tk', 'wx']

16
rfreq = [1381.,1390.,1420.5,1430.]
perdereau's avatar
perdereau committed
17
normf=[1325.,1465.]
OP's avatar
OP committed
18 19 20 21
names = ['1H','2H','3H','4H','1V','2V','3V','4V','1Hx2H','1Hx3H','1Hx4H','2Hx3H','2Hx4H','3Hx4H',
         '1Vx2V','1Vx3V','1Vx4V','2Vx3V','2Vx4V','3Vx4V','1Hx1V','1Hx2V','1Hx3V','1Hx4V','2Hx1V','2Hx2V','2Hx3V','2Hx4V',
         '3Hx1V','3Hx2V','3Hx3V','3Hx4V','4Hx1V','4Hx2V','4Hx3V','4Hx4V']

perdereau's avatar
perdereau committed
22 23 24 25 26
def viewtfmap(folder,num,mvmin=0.,mvmax=0.,save=False,   raplot=False,timeplot=False,mtitle="",clean=True,mod=False,
              arg=False,cmplx=False,noplo=False,verb=False,yrcm=False,cor=[],cimag=False,creal=False,hasvar=True,
              ramod=False,plras=[],nmras=[],mask=False):
    if verb :
        print len(cor)
OP's avatar
OP committed
27 28 29
    filename = folder
    #print   i
    #print num
perdereau's avatar
perdereau committed
30 31
    if noplo :
        clean=False
OP's avatar
OP committed
32 33
    if clean:
        plt.close("all")
perdereau's avatar
perdereau committed
34 35


OP's avatar
OP committed
36 37
    hdulist = fits.open(folder+'.fits')

38 39 40 41
    winwid = (18,8)

    times=hdulist[len(hdulist)-1].data
    ras=hdulist[len(hdulist)-2].data #41 in 43
OP's avatar
OP committed
42
    freqs= hdulist[len(hdulist)-3].data
43 44 45 46
    if ramod :
        ras=hdulist[len(hdulist)-1].data #41 in 43
        freqs= hdulist[len(hdulist)-2].data
        
47 48 49 50 51 52 53 54 55
    i=2*num
    if num>7 :
        i= 16 + 4*(num-8)

    if (not hasvar) :
        i=num     
        if num>7 :
            i= 8 + 2*(num-8)

56
    if len(hdulist) <8 :
57 58 59 60
        i=num # devrait etre juste la pour les sorties des voies th et rfi

    if (verb) :
        print hdulist[i].header
61
        
OP's avatar
OP committed
62
    bfreq = [ (np.abs(freqs-rf)).argmin() for rf in rfreq ]
perdereau's avatar
perdereau committed
63 64 65
    bfnor = [(np.abs(freqs-rf)).argmin() for rf in normf ]
    if (verb) :
        print ' number of hds ' + str(len(hdulist))
66

OP's avatar
OP committed
67 68 69
    nfre = np.size(freqs)
    ntim = np.size(times)
    timin = min(times)/3600.
70 71
    if timin ==0. :
        timin=times[0]/3600.
OP's avatar
OP committed
72
    timax = max(times)/3600.
73
    print '  timin/max = '+str(timin)+', ' +str(timax)
74 75 76
    if ramod :
        timin = min(ras)
        timax= max(ras)
OP's avatar
OP committed
77 78 79
    frmin = min(freqs)
    frmax=max(freqs)
    if (not noplo):
80
        fig,ax1=plt.subplots(figsize=winwid)
OP's avatar
OP committed
81 82 83 84
        
    img = hdulist[  i].data
    if(cimag):
        img = hdulist[  i+1].data
perdereau's avatar
perdereau committed
85 86 87
    if (verb) :
        print 'Image shape :'
        print np.shape(img)
OP's avatar
OP committed
88 89 90 91 92
        
    if mod :
        if num>7 :
            
            imgi= hdulist[  i+1].data
OP's avatar
OP committed
93
            img = np.sqrt(img**2+imgi**2)
OP's avatar
OP committed
94 95 96
            if len(cor) ==2 :
                imca = hdulist[cor[0]].data
                imcb = hdulist[4+cor[1]].data
97
                img=img/np.sqrt(imca*imcb)
OP's avatar
OP committed
98 99 100 101 102 103 104

    if arg :
        if num>7 :
            imgi= hdulist[  i+1].data
            img = np.arctan2(imgi,img) * 180. / np.pi
            mvmin=-180.
            mvmax=180.
perdereau's avatar
perdereau committed
105
            mtitle= names[num]+" arg.(deg)" 
OP's avatar
OP committed
106 107 108
        else :
            print "arg demande mais auto selectionne : ERREUR "
            exit
perdereau's avatar
perdereau committed
109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
    mskdat=img*0.
    if cmplx :
        imgi= hdulist[  i+1].data
        zmod = np.sqrt(img**2+imgi**2)
        
        arg =  np.arctan2(imgi,img)
        aa = (arg+np.pi) /(2.*np.pi)
        aa = (aa+0.5)%1
        bb = 1./(abs(zmod)**0.3 +1.)
        bb = -bb+1.
        print np.shape(aa)
        cc = np.ndarray(np.shape(aa))
        print  np.shape(cc)
        cc.fill(1.)
        iimg =[]
        for ii in np.arange(np.shape(aa)[0]):
            for jj in np.arange(np.shape(aa)[1]):
                iimg.append( hls_to_rgb(aa[ii,jj], bb[ii,jj],cc[ii,jj] ))
        img = np.reshape( iimg,(np.shape(aa)[0],np.shape(aa)[1],3))
    
    if mask :
        if cmplx :
            print "mask & cmplx not yet compatible "
            exit
        msknam = filename +'_masks.fits'
        print msknam+ "   " + str(num)
        hdmskl = fits.open(msknam)
        
        mskdat = np.fabs(hdmskl[num+1].data-1.)
        img = img*mskdat
        filename=filename+"_masked"
    if not cmplx :   
        vmin = np.asarray(img[bfnor[0]:bfnor[1],]).min() 
        vmax = np.asarray(img[bfnor[0]:bfnor[1],]).max()
    else :
        vmin=0
        vmax=1 
146
    
OP's avatar
OP committed
147 148
    if mvmin !=0. :
        vmin = mvmin
149
        
OP's avatar
OP committed
150 151 152 153 154 155 156
    if mvmax !=0. :
        vmax = mvmax
    if verb :
        print vmin, np.asarray(img[300:900,]).min()
        print vmax,np.asarray(img[300:900,]).max()

    if (not noplo) :
157 158
        im = ax1.imshow(img,aspect='auto' , interpolation='none',vmin=vmin,vmax=vmax)
        #m = ax1.imshow(img,aspect='auto' ,extent=(timin,timax,frmax,frmin), interpolation='none',vmin=vmin,vmax=vmax)
OP's avatar
OP committed
159
        if (yrcm) :
OP's avatar
OP committed
160
            im.set_cmap('YlOrBr_r')
OP's avatar
OP committed
161 162 163 164 165 166 167 168 169 170 171 172
        cbar = plt.colorbar( im ) #,fraction=0.021,pad=.05)
        cbar.set_norm(mynorm.MyNormalize(vmin=vmin,vmax=vmax,stretch='linear'))
        cm_cycle = sorted([i for i in dir(plt.cm) if hasattr(getattr(plt.cm,i),'N')])
        cm_index = cm_cycle.index(cbar.get_cmap().name)

        plt.show()
        if mtitle =="":
            plt.title(filename + ' ' + names[num],y=1.08)
        else :
            plt.title(filename + ' ' +mtitle,y=1.08)

        plt.ylabel("Frequency (Mhz)")
173 174 175 176
        if ramod:
            plt.xlabel("RA (h)")
        else:
            plt.xlabel("Temps (h)")
OP's avatar
OP committed
177

178 179 180
        if ramod :
            for rap in plras:
                ax1.plot([rap,rap],[frmax,frmin],color='red',linestyle='--')
perdereau's avatar
perdereau committed
181 182
                #ax1.annotate(nmras[plras.index(rap)],(rap,frmin),horizontalalignment='center',verticalalignment='bottom')
                ax1.annotate(nmras[plras.index(rap)],(rap,frmin),horizontalalignment='left',verticalalignment='bottom',rotation=30)
OP's avatar
OP committed
183 184
        ax1.yaxis.grid(which="major", color='gray', linestyle='-', linewidth=.5)
        ax1.xaxis.grid(which="major", color='gray', linestyle='-', linewidth=.5)
185
        
OP's avatar
OP committed
186
        ax2=ax1.twinx()
187
        ax2.yaxis.set_view_interval(frmax,frmin  ,ignore=True)
OP's avatar
OP committed
188
        ax2.set_ylabel("Freq. bin",rotation=270.,y=0.9)
189 190
        if not ramod :
            ax3=ax1.twiny()
191
            ax3.xaxis.set_view_interval( timin,timax  ,ignore=True)            
192
            ax3.set_xlabel("time bin",x=0.85)
OP's avatar
OP committed
193 194 195 196 197

        if save:
            ext=""
            if creal :
                ext = ext+"_real"
perdereau's avatar
perdereau committed
198 199
            if mask :
                ext=ext+"_mask"
200 201 202 203 204 205 206
            if ramod :
                if len(plras) >0 :
                    plt.savefig(filename+"_"+ names[num]+"_RAM"+ext+"_srcs.png")
                else :
                    plt.savefig(filename+"_"+ names[num]+"_RAM"+ext+".png")
            else :
                plt.savefig(filename+"_"+ names[num]+"_TFM"+ext+".png")
OP's avatar
OP committed
207
            
OP's avatar
OP committed
208 209 210 211 212 213 214

        axnext = fig.add_axes([0.6, 0.04, 0.05, 0.05])
        bnext = Button(axnext,cm_cycle[cm_index+1] )
        axprev = fig.add_axes([0.7, 0.04, 0.05, 0.05])
        bprev = Button(axprev,cm_cycle[cm_index-1] )
        axcur = fig.add_axes([0.65, 0.04, 0.05, 0.05], )
        bcur = Button(axcur,cm_cycle[cm_index])
OP's avatar
OP committed
215 216 217 218 219 220 221 222 223
        def cm_next(val):
            cm_cycle = sorted([i for i in dir(plt.cm) if hasattr(getattr(plt.cm,i),'N')])
            cm_index = cm_cycle.index(cbar.get_cmap().name)
            cm_index = cm_index+1

            if cm_index>= len(cm_cycle):
                cm_index=0

            change_cm(cm_index)
OP's avatar
OP committed
224
            #print cm_index
OP's avatar
OP committed
225 226 227 228 229 230 231
            
        def cm_prev(val):
            cm_cycle = sorted([i for i in dir(plt.cm) if hasattr(getattr(plt.cm,i),'N')])
            cm_index = cm_cycle.index(cbar.get_cmap().name)
            cm_index = cm_index-1
            
            if cm_index < 0 :
OP's avatar
OP committed
232
                cm_index=len(cm_cycle)-1
OP's avatar
OP committed
233 234

            change_cm(cm_index)
OP's avatar
OP committed
235
            #print cm_index
OP's avatar
OP committed
236 237
            
        def change_cm(index):
OP's avatar
OP committed
238
            
OP's avatar
OP committed
239 240
            cm_cycle = sorted([i for i in dir(plt.cm) if hasattr(getattr(plt.cm,i),'N')])
            cmap = cm_cycle[index]
OP's avatar
OP committed
241
            #print cmap
OP's avatar
OP committed
242 243 244 245
            cbar.set_cmap(cmap)
            cbar.draw_all()        
            im.set_cmap(cmap)
            cbar.patch.figure.canvas.draw()
OP's avatar
OP committed
246 247 248 249 250 251 252 253 254 255
            bcur.label.set_text(cmap)
            indprev=index-1
            if indprev<0 :
                indprev=len(cm_cycle)
            indnext=index+1
            if indnext>=len(cm_cycle):
                indnext=0
            bnext.label.set_text(cm_cycle[indnext])
            bprev.label.set_text(cm_cycle[indprev])
             
OP's avatar
OP committed
256 257 258
        bnext.on_clicked(cm_next)
        bprev.on_clicked(cm_prev)
        
259
        #plt.tight_layout()
OP's avatar
OP committed
260 261 262 263 264
        plt.subplots_adjust(bottom=0.15)
        plt.show()


    if raplot : 
265
        fig,ax1=plt.subplots(figsize=winwid)
OP's avatar
OP committed
266 267 268 269 270 271 272 273
    
        [ plt.plot(ras, img[b,],label="F=%6.2f MHz" %(freqs[b]) ) for b in bfreq ]
        plt.xlabel("RA (hour)")
        plt.ylabel("Signal ")
        plt.ylim(np.asarray(img[[bfreq],]).min(),np.asarray(img[[bfreq],]).max())
        leg = plt.legend()
        for legobj in leg.legendHandles:
            legobj.set_linewidth(2.0)
274 275
        plt.ylim((mvmin,mvmax))
        plt.xlim((0.,24.))
OP's avatar
OP committed
276 277 278 279 280 281
        if mtitle =="":
            plt.title(filename + ' ' + names[num],y=1.08)
        else :
            plt.title(filename + ' ' +mtitle,y=1.08)

        if save:
perdereau's avatar
perdereau committed
282 283 284 285
            if mask :
                plt.savefig(filename+"_"+ names[num]+"_mask_TFM_ra.png")
            else:
                plt.savefig(filename+"_"+ names[num]+"_TFM_ra.png")
OP's avatar
OP committed
286 287 288
    

    if timeplot : 
289
        fig,ax1=plt.subplots(figsize=winwid)
OP's avatar
OP committed
290 291 292 293 294 295 296
        [ plt.plot(times/3600., img[b,],label="F=%6.2f MHz" %(freqs[b]) ) for b in bfreq ]
        plt.xlabel('Time (TU, hour)')
        plt.ylabel("Signal ")
        plt.ylim(np.asarray(img[[bfreq],]).min(),np.asarray(img[[bfreq],]).max())
        leg = plt.legend()
        for legobj in leg.legendHandles:
            legobj.set_linewidth(2.0)
297
        plt.ylim((mvmin,mvmax))
OP's avatar
OP committed
298 299 300 301 302 303 304

        
        if mtitle =="":
            plt.title(filename + ' ' + names[num],y=1.08)
        else :
            plt.title(filename + ' ' +mtitle,y=1.08)
        if save:
perdereau's avatar
perdereau committed
305 306 307 308
            if mask :
                plt.savefig(filename+"_"+ names[num]+"_TFM_mask_freq.png")
            else:
                plt.savefig(filename+"_"+ names[num]+"_TFM_freq.png")
OP's avatar
OP committed
309 310 311 312 313
    if (verb) :
        print times.size
    if (verb) :
        print np.shape(img)
    if noplo :
perdereau's avatar
perdereau committed
314 315 316
        return times/3600.,img,freqs,ras,[ img[b,] for b in bfreq],mskdat
    elif    cmplx:
        return times/3600.,img,freqs,ras
OP's avatar
OP committed
317
    else :
318
        return times/3600.,img,freqs,ras,[ img[b,] for b in bfreq],ax1,bnext,bprev,mskdat  
OP's avatar
OP committed
319 320 321 322