2016年3月31日木曜日

Arf_BG_corrector.py

お察し下さい (pyfits練習) 第2弾。前よりは多少ファイルI/Oの書き方が賢くなっている...かも。あくまで個人的に作ったスクリプト。
SCRIPT_SHELL="python" 
SCRIPT_NAME="Arf_BG_corrector.py"
SCRIPT_HELP_ARG="1) Input Rate Bg PI file made by mathpha 2) Src Arf 3) Bg Arf 4) Output Bg PI file"
SCRIPT_HELP_SENTENCE="Modiefy Bg PI Channel with Src / Bg  arf ratio."
SCRIPT_NUM_ARG=4

###IMPORT###
import sys
import pyfits
import os
###HELP###

if len(sys.argv[1:])!=SCRIPT_NUM_ARG:
    print 'Name: '+SCRIPT_NAME
    print 'Arguments: '+SCRIPT_HELP_ARG
    print 'Explanation: '+SCRIPT_HELP_SENTENCE
    sys.exit()


###MAIN###
ev_per_ch=3.65

bg_pif=str(sys.argv[1])
src_arf=str(sys.argv[2])
bg_arf=str(sys.argv[3])
out_pif=str(sys.argv[4])

def main() :
    
    
    if os.path.isfile(out_pif) :
        os.remove(out_pif)
    
    data_src_arf = (pyfits.open(src_arf))['SPECRESP'].data
    data_bg_arf = (pyfits.open(bg_arf))['SPECRESP'].data
   
    ratio=[] 
    for n in range(len(data_src_arf)) :
        ratio.append([data_src_arf['ENERG_LO'][n],data_src_arf['ENERG_HI'][n],data_src_arf['SPECRESP'][n] / data_bg_arf['SPECRESP'][n]])
    
    hdulist=pyfits.open(bg_pif)
    
    sp=hdulist['SPECTRUM'].data
    ch_num=len(hdulist['SPECTRUM'].data)
    
    hdulist['SPECTRUM'].header['COMMENT']='Modefied with Arf_BG_corrector.py'
    hdulist['SPECTRUM'].header['COMMENT']='Src arf : '+src_arf
    hdulist['SPECTRUM'].header['COMMENT']='Bg arf : '+bg_arf
    hdulist['SPECTRUM'].header['COMMENT']='Bg pi : '+bg_pif

        
    for n in xrange(ch_num) :
        enegy=ev_per_ch*n/1e3
        for erange in ratio :
            if erange[0] <= enegy and enegy < erange[1] :
                sp['RATE'][n]=sp['RATE'][n]*erange[2]
                sp['STAT_ERR'][n]=sp['STAT_ERR'][n]*erange[2]
                break

    hdulist.writeto(out_pif)

main()