Python JLA_library类代码示例

def compute_model(options):

    import numpy
    import astropy.io.fits as fits
    import JLA_library as JLA
    from astropy.table import Table
    from astropy.cosmology import FlatwCDM
    from scipy.interpolate import interp1d


    # -----------  Read in the configuration file ------------
    params=JLA.build_dictionary(options.config)

    # -----------  Read in the SN ordering ------------------------
    SNeList = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc'])
    nSNe = len(SNeList)

    for i, SN in enumerate(SNeList):
        SNeList['id'][i] = SNeList['id'][i].replace('lc-', '').replace('.list', '').replace('_smp', '')

    lcfile = JLA.get_full_path(params[options.lcfits])
    SNe = Table.read(lcfile, format='fits')

    print 'There are %d SNe' % (nSNe)

    indices = JLA.reindex_SNe(SNeList['id'], SNe)
    SNe = SNe[indices]

    redshift = SNe['zcmb']
    replace=(redshift < 0)

    # For SNe that do not have the CMB redshift
    redshift[replace]=SNe[replace]['zhel']
    print len(redshift)

    if options.raw:
        # Data from the bottom left hand figure of Mosher et al. 2014.
        # This is option ii) that is descibed above
        offsets=Table.read(JLA.get_full_path(params['modelOffset']),format='ascii.csv')
        Delta_M=interp1d(offsets['z'], offsets['offset'], kind='linear',bounds_error=False,fill_value='extrapolate')(redshift)
    else:
        Om_0=0.303 # JLA value in the wCDM model
        cosmo1 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.0)
        cosmo2 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.024)
        Delta_M=5*numpy.log10(cosmo1.luminosity_distance(redshift)/cosmo2.luminosity_distance(redshift))
    
    # Build the covariance matrix. Note that only magnitudes are affected
    Zero=numpy.zeros(nSNe)
    H=numpy.concatenate((Delta_M,Zero,Zero)).reshape(3,nSNe).ravel(order='F')
    C_model=numpy.matrix(H).T * numpy.matrix(H)

    date = JLA.get_date()
    fits.writeto('C_model_%s.fits' % (date),numpy.array(C_model),clobber=True) 

    return None

def compute_Cstat(options):
    """Python program to compute C_stat
    """

    import numpy
    import astropy.io.fits as fits
    from astropy.table import Table
    import JLA_library as JLA

    # -----------  Read in the configuration file ------------

    params=JLA.build_dictionary(options.config)

    # -----------  Read in the SN ordering ------------------------
    SNeList = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc'])
    nSNe = len(SNeList)

    for i, SN in enumerate(SNeList):
        SNeList['id'][i] = SNeList['id'][i].replace('lc-', '').replace('.list', '')

    lcfile = JLA.get_full_path(params[options.lcfits])
    SNe = Table.read(lcfile, format='fits')


    # -----------  Read in the data --------------------------

    print 'There are %d SNe in the sample' % (nSNe)

    indices = JLA.reindex_SNe(SNeList['id'], SNe)
    SNe=SNe[indices]

    C_stat=numpy.zeros(9*nSNe*nSNe).reshape(3*nSNe,3*nSNe)

    for i,SN in enumerate(SNe):
        cov=numpy.zeros(9).reshape(3,3)
        cov[0,0]=SN['dmb']**2.
        cov[1,1]=SN['dx1']**2.
        cov[2,2]=SN['dcolor']**2.
        cov[0,1]=SN['cov_m_s']
        cov[0,2]=SN['cov_m_c']
        cov[1,2]=SN['cov_s_c']
        # symmetrise
        cov=cov+cov.T-numpy.diag(cov.diagonal())
        C_stat[i*3:i*3+3,i*3:i*3+3]=cov

    # -----------  Read in the base matrix computed using salt2_stat.cc ------------

    if options.base!=None:
        C_stat+=fits.getdata(options.base)

    date = JLA.get_date()
    fits.writeto('C_stat_%s.fits' % date,C_stat,clobber=True) 

    return

def add_covar_matrices(covmatrices,diag):
    """
    Python program that adds the individual covariance matrices into a single matrix
    """

    import numpy
    import astropy.io.fits as fits
    import JLA_library as JLA

    # Read in the terms that account for uncertainties in perculiar velocities, 
    # instrinsic dispersion and, lensing

    # Read in the covariance matrices
    matrices = []
    for matrix in covmatrices:
        matrices.append(fits.getdata(JLA.get_full_path(covmatrices[matrix]), 0))
        # Test for NaNs and replace them with zero
        if numpy.isnan(matrices[-1]).any():
            print 'Found a NaN in %s ... replacing them with zero' % (covmatrices[matrix])
            print numpy.isnan(matrices[-1]).sum()
            matrices[-1][numpy.isnan(matrices[-1])]=0.0

    # Add the matrices
    size = matrices[0].shape
    add = numpy.zeros(size[0]**2.).reshape(size[0], size[0])
    for matrix in matrices:
        add += matrix

    # Compute A

    nSNe = size[0]/3

    jla_results = {'Om':0.303, 'w':-1.027, 'alpha':0.141, 'beta':3.102}

    arr = numpy.zeros(nSNe*3*nSNe).reshape(nSNe, 3*nSNe)

    for i in range(nSNe):
        arr[i, 3*i] = 1.0
        arr[i, 3*i+1] = jla_results['alpha']
        arr[i, 3*i+2] = -jla_results['beta']

    cov = numpy.matrix(arr) * numpy.matrix(add) * numpy.matrix(arr).T

    # Add the diagonal terms

    sigma = numpy.genfromtxt(JLA.get_full_path(diag),
                             comments='#',
                             usecols=(0, 1, 2),
                             dtype='f8,f8,f8',
                             names=['sigma_coh', 'sigma_lens', 'sigma_pecvel'])

    for i in range(nSNe):
        cov[i, i] += sigma['sigma_coh'][i]**2 + \
        sigma['sigma_lens'][i]**2 + \
        sigma['sigma_pecvel'][i]**2

    return cov

def compute_dust(options):
    """Python program to compute C_dust
    """

    import numpy
    import astropy.io.fits as fits
    import os
    import JLA_library as JLA

    # ---------- Read in the SNe list -------------------------

    SNelist = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S110',
                               names=['id', 'lc'])

    for i, SN in enumerate(SNelist):
        SNelist['id'][i] = SNelist['id'][i].replace('lc-','').replace('.list','')

    # -----------  Read in the configuration file ------------

    params=JLA.build_dictionary(options.config)
    try:
        salt_path = JLA.get_full_path(params['defsaltModel'])
    except KeyError:
        salt_path = ''
        
    # -----------   The lightcurve fitting -------------------

    # Compute the offset between the nominal value of the extinciton 
    # and the adjusted value
    # We first compute the difference in light curve fit parameters for E(B-V) * (1+offset)
    offset = 0.1

    j = []

    for SN in SNelist:
        inputFile = SN['lc']
        print 'Fitting %s ' % (SN['lc'])
        workArea = JLA.get_full_path(options.workArea)
        dm, dx1, dc = JLA.compute_extinction_offset(SN['id'], inputFile, offset, workArea, salt_path)
        j.extend([dm, dx1, dc])
    
    # But we want to compute the impact of an offset that is twice as large, hence the factor of 4 in the expression
    # 2017/10/13
    # But we want to compute the impact of an offset that is half as large, hence the factor of 4 in the denominator
    # cdust = numpy.matrix(j).T * numpy.matrix(j) * 4.0
    cdust = numpy.matrix(j).T * numpy.matrix(j) / 4.0

    date = JLA.get_date()

    fits.writeto('C_dust_%s.fits' % date, cdust, clobber=True) 

    return

def compute_model(options):

    import numpy
    import astropy.io.fits as fits
    import JLA_library as JLA
    from astropy.table import Table
    from astropy.cosmology import FlatwCDM



    # -----------  Read in the configuration file ------------

    params=JLA.build_dictionary(options.config)

    # -----------  Read in the SN ordering ------------------------
    SNeList = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc'])
    nSNe = len(SNeList)

    for i, SN in enumerate(SNeList):
        SNeList['id'][i] = SNeList['id'][i].replace('lc-', '').replace('.list', '')

    lcfile = JLA.get_full_path(params[options.lcfits])
    SNe = Table.read(lcfile, format='fits')

    print 'There are %d SNe' % (nSNe)

    #z=numpy.array([])
    #offset=numpy.array([])
    Om_0=0.303 # JLA value in the wCDM model

    cosmo1 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.0)
    cosmo2 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.024)
    
    # For the JLA SNe
    redshift = SNe['zcmb']
    replace=(redshift < 0)
    # For the non JLA SNe
    redshift[replace]=SNe[replace]['zhel']

    Delta_M=5*numpy.log10(cosmo1.luminosity_distance(redshift)/cosmo2.luminosity_distance(redshift))

    # Build the covariance matrix. Note that only magnitudes are affected
    Zero=numpy.zeros(nSNe)
    H=numpy.concatenate((Delta_M,Zero,Zero)).reshape(3,nSNe).ravel(order='F')
    C_model=numpy.matrix(H).T * numpy.matrix(H)

    date = JLA.get_date()
    fits.writeto('C_model_%s.fits' % (date),numpy.array(C_model),clobber=True) 

    return None

def compute_filterTransVar(options):
    import JLA_library as JLA
    import numpy
    from astropy.table import Table

    eff=[]
    # Read in the filter curves
    filt=Table.read(JLA.get_full_path(options.filter),format='ascii.csv')
    for col in filt.colnames:
        if 'ccd' in col and 'amp' not in col:
            f=JLA.filterCurve(filt['wavelength'],filt[col])
            eff.append(f.eff())

    print 'Examined the transmission curves for %d CCDs' % (len(eff))
    print 'Mean effective wavelength is %6.1f' % (numpy.mean(eff))
    print 'Range of effective wavelength is %6.1f-%6.1f' % (numpy.min(eff),numpy.max(eff))
    print 'RMS effective wavelength %6.1f' % (numpy.std(eff))

    return

def compute_dust(options):
    """Python program to compute C_dust
    """

    import numpy
    import astropy.io.fits as fits
    import os
    import JLA_library as JLA

    # ---------- Read in the SNe list -------------------------

    SNelist = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S100',
                               names=['id', 'lc'])

    for i, SN in enumerate(SNelist):
        SNelist['id'][i] = SNelist['id'][i].replace('lc-','').replace('.list','')

    # -----------   The lightcurve fitting -------------------

    # Compute the offset between the nominal value of the extinciton 
    # and the adjusted value
    offset = 0.1

    j = []

    for SN in SNelist:
        inputFile = SN['lc']
        print 'Fitting %s' % (SN['id'])
        dm, dx1, dc = JLA.compute_extinction_offset(SN['id'], inputFile, offset)
        j.extend([dm, dx1, dc])
    
    cdust = numpy.matrix(j).T * numpy.matrix(j) * 4.0

    date = JLA.get_date()

    fits.writeto('C_dust_%s.fits' % date, cdust, clobber=True) 

    return

def compute_date_of_max(options):
    import numpy
    from astropy.table import Table
    import JLA_library as JLA

    params=JLA.build_dictionary(options.config)

    # -----------  Read in the configuration file ------------


    lightCurveFits=JLA.get_full_path(params['lightCurveFits'])
    lightCurves=JLA.get_full_path(params['lightCurves'])
    adjlightCurves=JLA.get_full_path(params['adjLightCurves'])


    # ---------  Read in the list of SNe ---------------------
    SNe = Table.read(lightCurveFits, format='fits')

    nSNe=len(SNe)
    print 'There are %d SNe' % (nSNe)

    # -----------   The lightcurve fitting -------------------

    J=[]

    for SN in SNe:
        SNfile='lc-'+SN['name']+'.list'
        #print 'Examining %s' % SN['name']
        inputFile=lightCurves+SNfile
        outputFile=adjlightCurves+SNfile
        # If needed refit the lightcurve and insert the date of maximum into the input file
        JLA.insertDateOfMax(SN['name'].strip(),inputFile,outputFile,options.force)

    return

def compute_ZP(options):

    import JLA_library as JLA
    import numpy as np

    params=JLA.build_dictionary(options.config)

    # Read in the standard star

    standard=JLA.spectrum(JLA.get_full_path(params['magSys'])+options.standard)

    # Read in the filter

    filt=JLA.filterCurve(JLA.get_full_path(params['filterDir'])+options.filter)

    # Compute the ZP
    if options.system=='AB':
        print '%s in %s %s %5.3f' % (options.standard,options.filter,options.system,filt.AB(standard))
    else:
        pass
#        print '%s in %s %s %5.3f' % (options.standard,options.filter,options.system,filt.Vega(standard))


    return

def compute_date_of_max(options):
    import numpy
    from astropy.table import Table
    import JLA_library as JLA

    params=JLA.build_dictionary(options.config)

    # ----------- Correction factor for extinction -----------
    # See ApJ 737 103
    extinctionFactor=0.86

    # -----------  Read in the configuration file ------------

    lightCurveFits=JLA.get_full_path(params['lightCurveFits'])
    lightCurves=JLA.get_full_path(params['lightCurves'])
    adjlightCurves=JLA.get_full_path(params['adjLightCurves'])

    # ---------  Read in the list of SNe ---------------------
    # One can either use an ASCII file with the SN list or a fits file
    if options.SNlist == None:
        SNe = Table.read(lightCurveFits, format='fits')
    else:
    # We use the ascii file, which gives the full path name
        SNe = Table.read(options.SNlist, format='ascii',names=['name','type','lc'],data_start=0)

    nSNe=len(SNe)
    print 'There are %d SNe' % (nSNe)

    # -----------   The lightcurve fitting -------------------

    for SN in SNe:
        if options.SNlist == None:
            SNfile='lc-'+SN['name']+'.list'
            inputFile=lightCurves+SNfile
            outputFile=adjlightCurves+SNfile
        else:
            inputFile=SN['lc']
            outputFile=SN['lc'].replace(lightCurves,adjlightCurves)

        print 'Examining %s' % SN['name']
        # If needed refit the lightcurve and insert the date of maximum into the input file
        JLA.insertDateOfMax(SN['name'].strip(),inputFile,outputFile,options.force,params)
        # Shouldn't we adjust the extinction first
        if options.adjustExtinction:
            adjustExtinction(outputFile,extinctionFactor)
    return

def reorderSNe(options):
    # The ordering of the SNe produced by salt2_stat does not reflect the order they were input
    # The ordering in the output file is written to the file sne_mu.list
    # SDSS SNe in the JLA sample get called @SN 12856.0, which means that the output of salt2_stat has these names.
    # Some nearby SNe have sn in front of their names. Others do not 
    # In one case a lower case v is used
    # DES SNe in the DES sample are listed as 01248677 

    import numpy
    import astropy.io.fits as fits
    import JLA_library as JLA
    from astropy.table import Table

    # -----------  Read in the SN ordering ------------------------
    SNeList = Table(numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc']))


    # -----------  Read in the file that specifies the ordering of the matrix produced by Cstat ------------
    statList = Table(numpy.genfromtxt(options.input,
                                      usecols=(0,1),
                                      dtype='S30,float',
                                      names=['id','z'],skip_header=13))


    # We use the -9 as a way to catch errors.
    reindex=numpy.zeros(len(SNeList),int)-9

    for i,SNname in enumerate(SNeList['id']):
        name=SNname.replace("lc-","").replace(".list","").replace('_smp','')
        print i,name
        for j,SNname2 in enumerate(statList['id']):
            if SNname2 == name or ("SDSS"+SNname2.replace(".0","") == name and "SDSS" in name) or \
                    (SNname2.replace("sn","")==name.replace("sn","")) or ("DES" in name and "DES_0"+SNname2==name):
                if reindex[i]!=-9:
                    print SNname,SNname2
                reindex[i]=j
                print SNname,SNname2,i,j
                
    for index,value in enumerate(reindex):
        if value==-9:
            print "Error"
            print index,SNeList[index]['id']
            exit()

    print "The numbers should be the same"
    print len(SNeList), len(statList), len(numpy.unique(reindex))

    # -----------  Read in Cstat and re-order --------------------------------------

    Cstat=fits.getdata(options.file)

    # We use brute force to reorder the elements
    # Recall that for each SNe, there is an error associated with the peak mag, colour and stretch

    Cstat_new=numpy.copy(Cstat) * 0.0

    nSNe=len(SNeList)
    for i in  range(nSNe):
        for j in range(nSNe):
            Cstat_new[3*reindex[i]:3*reindex[i]+3,3*reindex[j]:3*reindex[j]+3]=Cstat[3*i:3*i+3,3*j:3*j+3]


    date = JLA.get_date()
    fits.writeto('%s_Cstat_%s.fits' % (options.prefix,date),Cstat_new,clobber=True) 

    return None

def updateDES(options,params,model):
    try:
        shutil.rmtree(options.output+'/'+model['modelNumber']+'/snfit_data/Instruments/DECam')
    except:
        pass
    shutil.copytree(JLA.get_full_path(params['DES_instrument']),options.output+'/'+model['modelNumber']+'/snfit_data/Instruments/DECam')
    
    # Update the DES magnitude system
    shutil.copy(JLA.get_full_path(params['DES_magsys']),options.output+'/'+model['modelNumber']+'/snfit_data/MagSys/')

    return

def runSALT(SALTpath, SALTmodel, salt_prefix, inputFile, SN):
    import os
    
    # Set up the path to the SALT model and the name of the outputFile
    os.environ['SALTPATH']=SALTpath+SALTmodel['directory']+'/snfit_data/'
    outputFile=JLA.get_full_path(options.workArea)+'/'+SN+'/'+SN+'_'+SALTmodel['directory']+'.dat'
    if os.path.isfile(outputFile):
        pass
        #print "Skipping, fit with SALT model %s for %s already done" % (SALTmodel['directory'],os.path.split(inputFile)[1])
    else:
        # Otherwise, do the fit with the date of Max set to the value in the lightcurve file
        JLA.fitLC(inputFile, outputFile, salt_prefix, forceDayMax=True)
    return outputFile

def convert_lightcurves(options):

    # Read in the configuration file
    # The configuraiton file contains the location of various files
    params=JLA.build_dictionary(options.config)

    # Read in the extra variance
    # This depends on the photometric method. It is lower for SMP 
    extraVariance=get_extra_variance(JLA.get_full_path(params['extraVariance']),options)

    # Read in the extinction values
    # A temporary fix as the lightcurves do not currently have it
    # Still needed
    extinction=get_extinction(JLA.get_full_path(params['extinction']),options)

    snanaDir=JLA.get_full_path(params['snanaLightCurves'])
    saltDir=JLA.get_full_path(params['adjLightCurves'])

    try:
        os.mkdir(saltDir)
    except:
        pass
        
    saltDir=saltDir+'DES/'

    try:
        os.mkdir(saltDir)
    except:
        pass

    for lightcurve in os.listdir(snanaDir):
        if '.dat' in lightcurve:
            # Read in the snana file
            lc=snanaLightCurve(snanaDir+lightcurve)
            lightCurveFile=saltDir+lightcurve.replace('des_real','lc-DES').replace('.dat','.list')
            if lc.parameters['TYPE'].split()[0] in ['1','101']:   # Is a SN Ia or a SN Ia?
                print lightcurve, lightCurveFile
                lc.clean()                                # Remove bad photometry
                lc.addNoise(extraVariance)                # Add additional variance to the lightcurve points
                # It is not clear if we need to compute a rough date of max before doing the more precise fit
                lc.estimateDateOfMax(options)             # Sets an approximate date of max for the light curve fitting done below.
                # Apply cuts
                # lc.applySNCuts()
                # lc.applySamplingCuts()
                lc.write(lightCurveFile,options.format)   # Write out the resutlt
                lc.fitDateOfMax(lightCurveFile,params)    # Get a more precise estimate of the data of peak brightness
                lc.updateExtinction(lightCurveFile,extinction) # Temporary code
    return

def compute_diag(SNe):
    lens = 0.055
    #c = 3e5 #km/s
    sigma_lens = 0.055 * SNe['zcmb']
    sigma_pecvel = 5*8.33e-4/(np.log(10) *SNe['zcmb']) # See eq. 13 in B14
    sigma_coh = np.array([coh_dict[JLA.survey(sn)] for sn in SNe])
    return np.column_stack((sigma_coh, sigma_lens, sigma_pecvel))