camel.py 59.2 KB
Newer Older
Matthieu Tristram's avatar
Matthieu Tristram committed
1 2 3 4
#
#CAMEL tools library
#
#Sep 2015   - M. Tristram -
5 6 7
import numpy as np
import matplotlib.pyplot as plt
import scipy.ndimage as nd
8
import sys
9
import re
perdereau's avatar
perdereau committed
10
import glob,os
11
import gzip
12
import mmap
13
import pdb
Matthieu Tristram's avatar
Matthieu Tristram committed
14 15


16
def readCosmoMC( filename, nchain=4):
17
    cosmopars = list(np.loadtxt( "%s.paramnames" % (filename), dtype='str', unpack=True, usecols=[0]))
Matthieu Tristram's avatar
Matthieu Tristram committed
18

19
    parlist = ['iter','chi2']+[from_cosmomcName.get(p,p) for p in cosmopars]
Matthieu Tristram's avatar
Matthieu Tristram committed
20

21
    data = extract_chains( filename+"_", nchain)
Matthieu Tristram's avatar
Matthieu Tristram committed
22

23 24
    chain = {}
    for p in range(len(parlist)):
25
        chain[parlist[p]] = np.array([s for tmp in data for s in tmp[p]])
Matthieu Tristram's avatar
Matthieu Tristram committed
26
#        chain[parlist[p]] = np.concatenate( [tmp[p] for tmp in data])
Matthieu Tristram's avatar
Matthieu Tristram committed
27

28
    return( chain)
Matthieu Tristram's avatar
Matthieu Tristram committed
29

Matthieu Tristram's avatar
Matthieu Tristram committed
30
def mergeMC( filename, burnin=0.9, num=None, ext="txt", par=None, nelts=None):
Matthieu Tristram's avatar
Matthieu Tristram committed
31 32 33 34
    import subprocess
    import os
    import sys

Matthieu Tristram's avatar
Matthieu Tristram committed
35
    if num == None:
Matthieu Tristram's avatar
Matthieu Tristram committed
36
        num=[n for n in range(100) if os.path.isfile("%s%d.%s" % (filename,n,ext))]
Matthieu Tristram's avatar
Matthieu Tristram committed
37

Matthieu Tristram's avatar
Matthieu Tristram committed
38
    if par == None:
Matthieu Tristram's avatar
Matthieu Tristram committed
39
        par = MCparnames( filename, num[0], ext)
Matthieu Tristram's avatar
Matthieu Tristram committed
40 41

    print(num)
Matthieu Tristram's avatar
Matthieu Tristram committed
42 43 44 45 46 47 48
    chain = []
    for name in ["%s%d.%s" % (filename,n,ext) for n in num]:
        f = open( name, 'rb')
        f.readline() #labels

        line_size=np.mean([len(f.readline()) for i in range(100)])
#        print( line_size)
49

Matthieu Tristram's avatar
Matthieu Tristram committed
50 51 52 53 54 55 56 57
        nsamp = os.stat(name).st_size / line_size
#        print(nsamp)
        if 0 < burnin < 1:
            ist = int((1.-burnin)*nsamp)
        elif burnin < 0:
            ist = int(-burnin)
        else:
            ist = int(nsamp-burnin)
58

Matthieu Tristram's avatar
Matthieu Tristram committed
59 60
        if int(line_size*ist*1.2) < os.path.getsize(name):
            f.seek(-int(line_size*ist*1.2),os.SEEK_END)
Matthieu Tristram's avatar
Matthieu Tristram committed
61 62
        lines_found=f.readlines()
        f.close()
63

Matthieu Tristram's avatar
Matthieu Tristram committed
64 65
#        data=[np.array(x.split(),float) for x in lines_found[-ist:]]
        data=[[float(c) for c in x.split()] for x in lines_found[-ist:]]
66

Matthieu Tristram's avatar
Matthieu Tristram committed
67 68 69 70 71
        if len(chain)==0:
            chain = data
        else:
            chain = np.concatenate((chain,data),axis=0)

72 73 74
    if nelts:
        if nelts < len(chain):
            chain = chain[np.array(np.random.random(nelts)*len(chain),int),:]
75 76

    dchain = dict( list(zip(par,np.transpose(chain))))
Matthieu Tristram's avatar
Matthieu Tristram committed
77 78 79 80

    return(dchain)


Matthieu Tristram's avatar
Matthieu Tristram committed
81 82
#Read parameter names in MCMC sample file
def MCparnames( filename, num=1, ext="txt"):
83 84 85

    f = open( "%s%d.%s" % (filename,num,ext))
    par = [v for v in f.readline().replace("\n","").split("\t") if v != ""]
Matthieu Tristram's avatar
Matthieu Tristram committed
86
    f.close()
87 88 89
    return(par)


Matthieu Tristram's avatar
Matthieu Tristram committed
90 91
def Parnames( parfile, par_type=None):
    param_list = []
92

Matthieu Tristram's avatar
Matthieu Tristram committed
93 94
    if not os.path.isfile(parfile):
        raise ValueError( "Filename do not exists: %s" % parfile)
95

Matthieu Tristram's avatar
Matthieu Tristram committed
96 97 98 99 100 101 102 103 104 105 106 107 108 109
    sf = open(parfile)

    par_type=[par_type] if par_type else ["par","fix"]

    for l in sf.readlines():
        fval = re.findall(r'\S+',l)
#        print(fval)
        if len(fval) >1 :
            if fval[0] in par_type:
                if not fval[1] in param_list:
                    param_list.append(fval[1])
    sf.close()
    return( param_list)

Matthieu Tristram's avatar
Matthieu Tristram committed
110 111
def Interval( chain, par, CL=68, symmetrical=False):
    return( getci( np.percentile( chain[par], [(100-CL)/2,50,100-(100-CL)/2]), symmetrical=symmetrical))
112

Matthieu Tristram's avatar
Matthieu Tristram committed
113 114 115
def Upper( chain, par, CL=95):
    return( np.percentile( chain[par], CL))

Matthieu Tristram's avatar
Matthieu Tristram committed
116

Matthieu Tristram's avatar
Matthieu Tristram committed
117 118 119 120 121 122 123 124 125
def scan(filename):
    f = open( filename, 'r')
    l=f.readlines()
    f.close()

    index=0
    while not l[index].startswith("SCANNING PAR"):
        index = index+1

126 127
    print((l[index].replace("\n","")))

Matthieu Tristram's avatar
Matthieu Tristram committed
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
    while not l[index].startswith("it "):
        index = index+1

    par = [v.strip() for v in l[index].replace("\n","").split("|") if v != ""]
    index = index + 1

    scan = []
    while not l[index].startswith("===="):
        scan.append( [v for v in l[index].replace("\n","").split("\t") if v != ""])
        index = index + 1

    return( scan,par)


#expand given mat with the given parameters on the diag
def expand_cov( mat, par, newpar, newvar):
144

Matthieu Tristram's avatar
Matthieu Tristram committed
145 146
    match    = [newpar.index(i) for i in newpar if i in par]
    notmatch = [newpar.index(i) for i in newpar if i not in par]
147

Matthieu Tristram's avatar
Matthieu Tristram committed
148
    sz = len(newpar)
149 150
    newmat = np.zeros( (sz, sz))

Matthieu Tristram's avatar
Matthieu Tristram committed
151 152 153
    for i in match:
        for j in match:
            newmat[i,j] = mat[par.index(newpar[i]),par.index(newpar[j])]
154

Matthieu Tristram's avatar
Matthieu Tristram committed
155 156
    for i in notmatch:
        newmat[i,i] = newvar[i]
157

Matthieu Tristram's avatar
Matthieu Tristram committed
158 159 160 161
    return( newmat)



Matthieu Tristram's avatar
Matthieu Tristram committed
162
def extract_chains( chainame, nchain):
163

164
    chain = []
Matthieu Tristram's avatar
Matthieu Tristram committed
165
    for c in range(0,nchain):
166
        print(( "%s%d.txt" % (chainame,c+1)))
Matthieu Tristram's avatar
Matthieu Tristram committed
167
        data = []
168

Matthieu Tristram's avatar
Matthieu Tristram committed
169 170 171 172
        with open("%s%d.txt" % (chainame,c+1)) as f:
            f.readline()
            for line in f:
                data.append( [float(a) for a in line.split()])
Matthieu Tristram's avatar
Matthieu Tristram committed
173

174
        chain.append( np.array(np.transpose(data)))
175

176
    return(chain)
Matthieu Tristram's avatar
Matthieu Tristram committed
177 178 179



180 181 182 183 184 185
def getci( c, symmetrical=False):
    if symmetrical:
        return(c[1], np.mean([abs(c[0]-c[1]), abs(c[2]-c[1])]))
    else:
        return(c[1], c[0]-c[1], c[2]-c[1])

Matthieu Tristram's avatar
Matthieu Tristram committed
186 187 188 189 190 191 192 193 194 195 196 197 198
def ci2sig( val, ci, symmetrical=False):
    if symmetrical:
        return(val, np.mean([abs(ci[0]-val), abs(ci[1]-val)]))
    else:
        return(val, ci[0]-val, ci[1]-val)

def sig2ci( val, sig):
    if len(sig) == 1:
        return( val, (val-sig,val+sig))
    else:
        return( val, (val-np.abs(sig[0]),val+sig[1]))


199
def cov2cor( mat):
200

201 202 203 204
    cor = mat.copy()
    for i in range( len(mat)):
        for j in range( len(mat)):
            cor[i,j] = mat[i,j]/np.sqrt(mat[i,i]*mat[j,j])
205

206 207 208
    return(cor)


Matthieu Tristram's avatar
Matthieu Tristram committed
209

210 211 212
########################################################################
#CONVERGENCE
########################################################################
Matthieu Tristram's avatar
Matthieu Tristram committed
213 214
def GelmanRubin( chains, gap=10000, length_min=1000, length_max=None, new=False):
    #plot GelmanRubin test skipping first parameter=chi2
215
    nchain = len(chains)
216

217
    nsamples = min( [len(c[0]) for c in chains])
218 219
    if length_max==None:
        length_max = nsamples
220

221 222
    if length_max > nsamples:
        print( "Not enough samples...")
223 224 225
        exit()

    it = list(range(length_min,length_max,gap))[1:-1]
226
    R = []
227
    for isamp in it:
228
        print(( "%d%%" % int(it.index(isamp)*100./len(it))))
Matthieu Tristram's avatar
Matthieu Tristram committed
229 230 231
        if new:
            #Gelman-Rubin stat on the last "gap" samples for each iteration
            n = gap
Matthieu Tristram's avatar
Matthieu Tristram committed
232
            tmp = [data[1:,isamp:isamp+gap] for data in chains]
Matthieu Tristram's avatar
Matthieu Tristram committed
233 234 235
        else:
            #Gelman-Runbin stat on the last half samples for each iteration
            n = (isamp - length_min)/2
Matthieu Tristram's avatar
Matthieu Tristram committed
236
            tmp = [data[1:,length_min+n:isamp] for data in chains]
237 238

        #within Chain
239 240
        mchain = np.mean( tmp,2)
        schain =  np.var( tmp,2, ddof=1)
241

242
        #between chains
243
        W = np.mean( schain,0)
244
        B =  np.var( mchain,0, ddof=1)*n
245

246
        #sqrt ?
247
        R.append( ( (n-1.)/n*W + B/n )/W )
248

249
    return(it,np.array(R))
Matthieu Tristram's avatar
Matthieu Tristram committed
250 251


252 253 254
def Geweke( chains, gap=10000, length_max=None):
    length_min = 1000
    nchain = len(chains)
255

Matthieu Tristram's avatar
Matthieu Tristram committed
256
    nsamples = min( [len(c[0]) for c in chains])
257 258
    if length_max==None:
        length_max = nsamples
Matthieu Tristram's avatar
Matthieu Tristram committed
259

260 261 262
    if length_max > nsamples:
        print( "Not enough samples...")
        exit()
263 264

    it = list(range(length_min,length_max,gap))[1:]
265
    T = []
266
    for isamp in it:
267
        print(( "%d%%" % int(it.index(isamp)*100./len(it))))
268
        n = length_max-isamp
Matthieu Tristram's avatar
Matthieu Tristram committed
269 270
        tmp1 = [data[:,isamp:isamp+n*0.1] for data in chains]
        tmp2 = [data[:,isamp+n*0.5:length_max] for data in chains]
271

272
        #Tscore
Matthieu Tristram's avatar
Matthieu Tristram committed
273 274 275
        Z1 = np.mean( tmp1,2)
        Z2 = np.mean( tmp2,2)
        s = np.sqrt( np.var(tmp1,2)/(n*0.1) + np.var(tmp2,2)/(n*0.5) )
276

277
        T.append( (Z1-Z2)/s)
278

279 280
    return(it,T)
########################################################################
Matthieu Tristram's avatar
Matthieu Tristram committed
281 282 283 284 285





286 287
########################################################################
#PARAMETER NAMES
Matthieu Tristram's avatar
Matthieu Tristram committed
288 289
########################################################################
parname = {
290
    'omega_b'              : '$\\Omega_\mathrm{b}h^2$',
291
    'omega_cdm'            : '$\\Omega_\mathrm{c}h^2$',
292
    'Omega_b'              : '$\\Omega_\mathrm{b}$',
293
    'Omega_cdm'            : '$\\Omega_\mathrm{c}$',
Matthieu Tristram's avatar
Matthieu Tristram committed
294 295 296 297 298
    '100*theta_s'          : '$100\\theta_\mathrm{s}$',
    'tau_reio'             : '$\\tau$',
    'z_reio'               : '$\mathrm{z}_\mathrm{reio}$',
    'n_s'                  : '$\mathrm{n}_\mathrm{s}$',
    'log(10^10A_s)'        : '$log(10^{10}A_\mathrm{s})$',
299 300
    'H0'                   : '$H_0$',
    'sigma8'               : '$\\sigma_8$',
301
    'Alens'                : '$\mathrm{A}_\mathrm{L}$',
302
    'm_ncdm'               : '$\\Sigma m_\\nu$',
303 304
    'N_eff'                : '$N_\mathrm{eff}$',
    'r'                    : '$r$',
Matthieu Tristram's avatar
Matthieu Tristram committed
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336
    'A_planck'             : '$A_\mathrm{planck}$',
    'A_cib_217'            : '$A_\mathrm{CIB}^{217}$',
    'xi_sz_cib'            : '$\\xi_\mathrm{SZxCIB}$',
    'A_sz'                 : '$A_\mathrm{SZ}$',
    'ps_A_100_100'         : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(100)}$',
    'ps_A_143_143'         : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(143)}$',
    'ps_A_217_217'         : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(217)}$',
    'ps_A_143_217'         : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(143x217)}$',
    'ksz_norm'             : '$A_\mathrm{kSZ}$',
    'gal545_A_100'         : '$A_\mathrm{gal545}^\mathrm{100}$',
    'gal545_A_143'         : '$A_\mathrm{gal545}^\mathrm{143}$',
    'gal545_A_143_217'     : '$A_\mathrm{gal545}^\mathrm{143x217}$',
    'gal545_A_217'         : '$A_\mathrm{gal545}^\mathrm{217}$',
    'calib_100T'           : '$c^\mathrm{Planck}_{100}$',
    'calib_217T'           : '$c^\mathrm{Planck}_{217}$',
    'c0'                   : '$c^\mathrm{Planck}_\mathrm{100hm1}$',
    'c1'                   : '$c^\mathrm{Planck}_\mathrm{100hm2}$',
    'c2'                   : '$c^\mathrm{Planck}_\mathrm{143hm1}$',
    'c3'                   : '$c^\mathrm{Planck}_\mathrm{143hm2}$',
    'c4'                   : '$c^\mathrm{Planck}_\mathrm{217hm1}$',
    'c5'                   : '$c^\mathrm{Planck}_\mathrm{217hm2}$',
    'Aps100x100'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(100x100)}$',
    'Aps100x143'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(100x143)}$',
    'Aps100x217'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(100x217)}$',
    'Aps143x143'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(143x143)}$',
    'Aps217x217'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(217x217)}$',
    'Aps143x217'           : '$A_\mathrm{PS}^\mathrm{Planck}\mathrm{(143x217)}$',
    'Asz'                  : '$A_\mathrm{SZ}$',
    'Acib'                 : '$A_\mathrm{CIB}$',
    'Aksz'                 : '$A_\mathrm{kSZ}$',
    'Aszxcib'              : '$A_\mathrm{SZxCIB}$',
    'AdustTT'              : '$A_\mathrm{dust}^\mathrm{TT}$',
337 338
    'AdustTP'              : '$A_\mathrm{dust}^\mathrm{TE}$',
    'AdustPP'              : '$A_\mathrm{dust}^\mathrm{EE}$',
perdereau's avatar
perdereau committed
339
    'SPT_ADust'            : '$A^\mathrm{dust}_\mathrm{SPT}$',
Matthieu Tristram's avatar
Matthieu Tristram committed
340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364
    'SPT_high_95_cal'      : '$c^\mathrm{SPT\ high}_{95}$',
    'SPT_high_150_cal'     : '$c^\mathrm{SPT\ high}_{150}$',
    'SPT_high_220_cal'     : '$c^\mathrm{SPT\ high}_{220}$',
    'SPT_high_Aps_95x95'   : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(95x95)}$',
    'SPT_high_Aps_95x150'  : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(95x150)}$',
    'SPT_high_Aps_95x220'  : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(95x220)}$',
    'SPT_high_Aps_150x150' : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(150x150)}$',
    'SPT_high_Aps_150x220' : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(150x220)}$',
    'SPT_high_Aps_220x220' : '$A_{PS}^\mathrm{SPT\ high}\mathrm{(220x220)}$',
    'SPT_low_Aps'          : '$A_{PS}^\mathrm{SPT\ low}$',
    'SPT_low_cal'          : '$c^\mathrm{SPT\ low}$',
    'ACT_equat_148_cal'    : '$c^\mathrm{ACT\ equat}_{148}$',
    'ACT_equat_220_cal'    : '$c^\mathrm{ACT\ equat}_{220}$',
    'ACT_equat_ADust'      : '$A_\mathrm{dust}^\mathrm{ACT\ equat}$',
    'ACT_equat_Aps_148x148' : '$A_{PS}^\mathrm{ACT\ equat}\mathrm{(148x148)}$',
    'ACT_equat_Aps_148x220' : '$A_{PS}^\mathrm{ACT\ equat}\mathrm{(148x220)}$',
    'ACT_equat_Aps_220x220' : '$A_{PS}^\mathrm{ACT\ equat}\mathrm{(220x220)}$',
    'ACT_south_148_cal'     : '$c^\mathrm{ACT\ south}_{148}$',
    'ACT_south_220_cal'     : '$c^\mathrm{ACT\ south}_{220}$',
    'ACT_south_ADust'       : '$A_\mathrm{dust}^\mathrm{ACT\ south}$',
    'ACT_south_Aps_148x148' : '$A_{PS}^\mathrm{ACT\ south}\mathrm{(148x148)}$',
    'ACT_south_Aps_148x220' : '$A_{PS}^\mathrm{ACT\ south}\mathrm{(148x220)}$',
    'ACT_south_Aps_220x220' : '$A_{PS}^\mathrm{ACT\ south}\mathrm{(220x220)}$',
    'chi2'                  : '$\\chi^2$'
    }
Matthieu Tristram's avatar
Matthieu Tristram committed
365 366


367
to_cosmomcName = {
Matthieu Tristram's avatar
Matthieu Tristram committed
368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402
    'omega_b'              : 'omegabh2',
    'omega_cdm'            : 'omegach2',
    '100*theta_s'          : 'theta',
    'tau_reio'             : 'tau',
    'z_reio'               : 'zrei',
    'n_s'                  : 'ns',
    'log(10^10A_s)'        : 'logA',
    'A_planck'             : 'calPlanck',
    'A_cib_217'            : 'acib217',
    'xi_sz_cib'            : 'xi',
    'A_sz'                 : 'asz143',
    'ps_A_100_100'         : 'aps100',
    'ps_A_143_143'         : 'aps143',
    'ps_A_217_217'         : 'aps217',
    'ps_A_143_217'         : 'aps143217',
    'ksz_norm'             : 'aksz',
    'gal545_A_100'         : 'kgal100',
    'gal545_A_143'         : 'kgal143',
    'gal545_A_143_217'     : 'kgal143217',
    'gal545_A_217'         : 'kgal217',
    'calib_100T'           : 'cal0',
    'calib_217T'           : 'cal2',
    'c0'                   : 'c0',
    'c1'                   : 'c1',
    'c2'                   : 'c2',
    'c3'                   : 'c3',
    'c4'                   : 'c4',
    'c5'                   : 'c5',
    'Aps100x100'           : 'Aps100x100',
    'Aps100x143'           : 'Aps100x143',
    'Aps100x217'           : 'Aps100x217',
    'Aps143x143'           : 'Aps143x143',
    'Aps217x217'           : 'Aps217x217',
    'Aps143x217'           : 'Aps143x217',
    }
Matthieu Tristram's avatar
Matthieu Tristram committed
403

404
from_cosmomcName = dict(list(zip(list(to_cosmomcName.values()),list(to_cosmomcName.keys()))))
405 406


Matthieu Tristram's avatar
Matthieu Tristram committed
407 408

########################################################################
409 410
#PLOTS
########################################################################
411
def posterior1d( chains, params, nbin=50, smooth=1,
412
                 names=[], colors=[], linestyles=[], lw=[], extent=[], subplot=[1,1], figsize=None):
413 414 415
    nchain = len(chains)
    npar   = len(params)

Matthieu Tristram's avatar
Matthieu Tristram committed
416
    fig=plt.figure(figsize=figsize)
417
    if len(colors) < nchain:
418 419 420
#        ax = plt.gca()
#        color_cycle = ax._get_lines.prop_cycler #ax._get_lines.color_cycle
#        colors = [next(color_cycle) for i in range(nchain)]
Matthieu Tristram's avatar
Matthieu Tristram committed
421
        colors = plt.rcParams['axes.color_cycle']
422 423 424 425 426 427 428

    if len(linestyles) < nchain:
        linestyles = ['-']*nchain

    if np.size(lw) == 1:
        lw = [lw]*nchain
    elif np.size(lw) < nchain:
429
        lw = [1]*nchain
430 431 432 433

    if np.prod(subplot) < npar:
        subplot[1] = np.ceil(np.double(npar)/subplot[0])

434 435 436 437 438
    if type(extent) == tuple:
        extent = [extent]*npar
    if len(extent) != npar:
        extent=None

439 440 441 442 443
    for par in params:
        ax=plt.subplot( subplot[0], subplot[1], params.index(par)+1)

        for c in range(nchain):
            chain = chains[c]
444
            if par in chain:
445
                hist1d( chain[par], smooth=smooth, color=colors[c], linestyle=linestyles[c], lw=lw[c], bins=nbin)
446

447 448
        ax.tick_params(axis='both', which='major', labelsize=16)
        ax.locator_params(tight=True, nbins=6)
449
        plt.xlabel( parname.get(par,par))
450 451 452

        if extent:
            ax.set_xlim( extent[params.index(par)])
453

454 455
    #legend
    if names != []:
Matthieu Tristram's avatar
Matthieu Tristram committed
456
#        if len(params) >= 2:
Matthieu Tristram's avatar
Matthieu Tristram committed
457 458
        if len(params) ==1:
            ax.legend(names, loc='upper center')
459
        else:
Matthieu Tristram's avatar
Matthieu Tristram committed
460 461 462 463 464 465
            if len(params) == subplot[0]*subplot[1]:
                ax.legend(names, loc='upper center', frameon=False, ncol=len(names), bbox_to_anchor=(0.,1.2))
                plt.subplots_adjust( top=0.85)
            else:
                ax.legend(names, bbox_to_anchor=(1.05, 0), loc='lower left', borderaxespad=0.)
                plt.subplots_adjust( left=0.1, right=0.85)
Matthieu Tristram's avatar
Matthieu Tristram committed
466
#            ax.legend(names, loc='upper right')
467 468


Matthieu Tristram's avatar
Matthieu Tristram committed
469
def posterior2d( chains, par1, par2, *args, **kwargs):
Matthieu Tristram's avatar
Matthieu Tristram committed
470 471 472 473 474 475 476 477 478
    """
    Plot a 2-D histogram of samples.

    """
    import scipy.stats
    import matplotlib.cm as cm

    x = chain[par1]
    y = chain[par2]
479

Matthieu Tristram's avatar
Matthieu Tristram committed
480 481 482 483 484 485 486 487
#    ax = kwargs.pop("ax", plt.gca())
    extent     = kwargs.pop("extent", [[x.min(), x.max()], [y.min(), y.max()]])
    bins       = kwargs.pop("bins",          50)
    colors     = kwargs.pop("colors",      None)
    alpha      = kwargs.pop("alpha",        0.8)
    linewidths = kwargs.pop("linewidths",  None)
    fill       = kwargs.get("fill",        True)
    datapoints = kwargs.get("datapoints", False)
488
    contours   = kwargs.get("contours",   False)
Matthieu Tristram's avatar
Matthieu Tristram committed
489
    levels     = kwargs.get("levels",   [0.68,0.95])
Matthieu Tristram's avatar
Matthieu Tristram committed
490 491 492
    cmap       = kwargs.get("cmap",        None)

#    cmap = cm.get_cmap(kwargs.get("cmap", None))
493

Matthieu Tristram's avatar
Matthieu Tristram committed
494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509
    X = np.linspace(extent[0][0], extent[0][1], bins + 1)
    Y = np.linspace(extent[1][0], extent[1][1], bins + 1)
    try:
        H, X, Y = np.histogram2d(x.flatten(), y.flatten(), bins=(X, Y),
                                 weights=kwargs.get('weights', None))
    except ValueError:
        raise ValueError("It looks like at least one of your sample columns "
                         "have no dynamic range. You could try using the "
                         "`extent` argument.")

    if datapoints:
        color = 'k' if colors==None else colors[0]
        plt.plot(x, y, "o", color=color, ms=1.5, zorder=-1, alpha=0.1, rasterized=True)

    gkde = scipy.stats.gaussian_kde([x, y])
    dx = (extent[0][1]-extent[0][0])/(bins+1.)
510 511
    dy = (extent[1][1]-extent[1][0])/(bins+1.)

Matthieu Tristram's avatar
Matthieu Tristram committed
512 513
    x2,y2 = np.mgrid[extent[0][0]:extent[0][1]:dx, extent[1][0]:extent[1][1]:dy]
    z   = np.array(gkde.evaluate([x2.flatten(),y2.flatten()])).reshape(x2.shape)
514

Matthieu Tristram's avatar
Matthieu Tristram committed
515
    W = np.append(ctr_level( z, levels),z.max())
516

Matthieu Tristram's avatar
Matthieu Tristram committed
517 518 519 520 521
    if fill:
        plt.contourf(x2, y2, z, W, alpha=alpha, cmap=cm.get_cmap(cmap), antialiased=True, colors=colors)

    if contours:
        plt.contour(x2, y2, z, W,linewidths=linewidths, colors=colors, antialiased=False)
522

Matthieu Tristram's avatar
Matthieu Tristram committed
523 524
    plt.xlim(extent[0])
    plt.ylim(extent[1])
525

Matthieu Tristram's avatar
Matthieu Tristram committed
526 527 528 529 530 531 532 533 534 535 536 537 538 539 540
    plt.xlabel(parname.get(par1,par1))
    plt.ylabel(parname.get(par2,par2))


#backward compatibility
def hist2d( chain, par1, par2, *args, **kwargs):
    """
    Plot a 2-D histogram of samples.

    """
    import scipy.stats
    import matplotlib.cm as cm

    x = chain[par1]
    y = chain[par2]
541

Matthieu Tristram's avatar
Matthieu Tristram committed
542 543 544 545
    extent     = kwargs.pop("extent", [[x.min(), x.max()], [y.min(), y.max()]])
    bins       = kwargs.pop("bins",          50)
    nsmooth    = kwargs.pop("nsmooth",       0.)
    cmap       = kwargs.get("cmap",        None)
546

Matthieu Tristram's avatar
Matthieu Tristram committed
547 548 549 550 551 552 553 554 555 556 557 558
    X = np.linspace(extent[0][0], extent[0][1], bins + 1)
    Y = np.linspace(extent[1][0], extent[1][1], bins + 1)
    try:
        H, X, Y = np.histogram2d(x.flatten(), y.flatten(), bins=(X, Y),
                                 weights=kwargs.get('weights', None))
    except ValueError:
        raise ValueError("It looks like at least one of your sample columns "
                         "have no dynamic range. You could try using the "
                         "`extent` argument.")

    if nsmooth != 0.:
        H = nd.gaussian_filter( H, nsmooth)
559

Matthieu Tristram's avatar
Matthieu Tristram committed
560
    npts = 5.
561 562 563 564
#    plt.imshow( H.T, origin='bottom', **kwargs)
#    plt.xticks(np.arange(npts+1)/npts*bins,["%4.2f" % i for i in np.linspace(extent[0][0],extent[0][1],npts+1)])
#    plt.yticks(np.arange(npts+1)/npts*bins,["%4.2f" % i for i in np.linspace(extent[1][0],extent[1][1],npts+1)])
    plt.pcolormesh( X, Y, H.T, **kwargs)
565

Matthieu Tristram's avatar
Matthieu Tristram committed
566 567 568
    plt.xlabel(parname.get(par1,par1))
    plt.ylabel(parname.get(par2,par2))

569 570


571
def triangle( chains, params, nbin=50, parnames=None, names=[], colors=[], linestyles=[], smooth=1., fontsize=12,
572
              contour=True, fill=False, cmaps=[], bestfit=None, Norm=True, weights=None):
573
    import matplotlib.ticker as mtick
574

Matthieu Tristram's avatar
Matthieu Tristram committed
575 576 577
    fig=plt.figure(figsize=(14,12))
    plt.subplots_adjust(hspace=0.001,wspace=0.001)

578 579 580 581
    #force tuple
    if not( isinstance(chains, list) or isinstance(chains, tuple)):
        chains = (chains,)

Matthieu Tristram's avatar
Matthieu Tristram committed
582
    nchain = len(chains)
583
    npar   = len(params)
Matthieu Tristram's avatar
Matthieu Tristram committed
584

585
    if len(colors) < nchain:
586 587 588
#        ax = plt.gca()
#        color_cycle = ax._get_lines.prop_cycler  #ax._get_lines.color_cycle
#        colors = [next(color_cycle) for i in range(nchain)]
Matthieu Tristram's avatar
Matthieu Tristram committed
589
        colors = plt.rcParams['axes.color_cycle']
590

Matthieu Tristram's avatar
Matthieu Tristram committed
591 592 593 594 595 596 597 598 599
    if len(cmaps) < nchain:
        cmaps = ['Blues','Greens','Greys','Oranges','Purples']
#        cmaps = [m for m in cm.datad if not m.endswith("_r")]

    if type(fill) == bool:
        fill = [fill]*nchain
    if type(contour) == bool:
        contour = [contour]*nchain

600 601
    if len(linestyles) < nchain:
        linestyles = ['-']*nchain
602

603 604 605
    if weights == None:
        weights = [np.ones(len(c['chi2'])) for c in chains]

606
    if parnames == None:
607
        parnames = [parname.get(p,p) for p in params]
608 609 610 611 612 613

    rge = {}
    for par in params:
        tmp = []
        for chain in chains:
            tmp = tmp+list(chain[par])
Matthieu Tristram's avatar
Matthieu Tristram committed
614 615
        rge[par] = list([min(tmp),max(tmp)])
#        rge[par] = list(np.mean(tmp) + np.dot(np.std(tmp), [-4,4]))
Matthieu Tristram's avatar
Matthieu Tristram committed
616 617


618 619 620 621
    for par in params:
        xpar = params.index(par)
        ax=plt.subplot( npar, npar, xpar*npar+xpar+1)

Matthieu Tristram's avatar
Matthieu Tristram committed
622 623
        for c in range(nchain):
            chain = chains[c]
624 625
            weight = weights[c]
            h,X=np.histogram( chain[par], bins=nbin, range=rge[par], normed=True, weights=weight)
Matthieu Tristram's avatar
Matthieu Tristram committed
626
            Y = nd.gaussian_filter1d( h, smooth,mode='nearest')
627 628
            Y = Y/np.max(Y) if Norm else Y
            plt.plot( (X[1:]+X[:-1])/2., Y, color=colors[c], linestyle=linestyles[c])
629

630 631
            if bestfit:
                plt.axvline( (bestfit[c])[par], ls='--', color=colors[c])
632

633
        ax.set_xlim( rge[par])
Matthieu Tristram's avatar
Matthieu Tristram committed
634
        ax.yaxis.set_visible(False)
635
        ax.xaxis.set_visible(xpar==(npar-1))
636
        ax.locator_params(tight=True, nbins=5)
637
        if 'PS' in parname.get(parnames[xpar],parnames[xpar]):
638 639 640
            ax.xaxis.set_major_formatter(mtick.FormatStrFormatter('%.0e'))

        #ax.xaxis.set_major_formatter(mtick.FormatStrFormatter('%.2e'))
Matthieu Tristram's avatar
Matthieu Tristram committed
641

642 643 644
        for ypar in range(xpar)[::-1]:
            par2 = params[ypar]
            ax=plt.subplot( npar, npar, xpar*npar+ypar+1)
Matthieu Tristram's avatar
Matthieu Tristram committed
645 646 647

            for c in range(nchain):
                chain = chains[c]
648 649
                weight = weights[c]
                hmap,binX,binY = np.histogram2d( chain[par], chain[par2], bins=nbin, range=[rge[par],rge[par2]], weights=weight)
650
                Y,X = np.meshgrid((binX[1:]+binX[:-1])/2.,(binY[1:]+binY[:-1])/2.)
Matthieu Tristram's avatar
Matthieu Tristram committed
651
                hmap = nd.gaussian_filter( hmap.T, 2)
Matthieu Tristram's avatar
Matthieu Tristram committed
652
                if fill[c]:
Matthieu Tristram's avatar
Matthieu Tristram committed
653
                    plt.contourf( X,Y, hmap, levels=ctr_level( hmap, [1e-2,0.68,0.95]),
Matthieu Tristram's avatar
Matthieu Tristram committed
654
                                  colors=None, extent=tuple(rge[par]+rge[par2]), cmap=plt.cm.get_cmap(cmaps[c]))
Matthieu Tristram's avatar
Matthieu Tristram committed
655
                if contour[c]:
Matthieu Tristram's avatar
Matthieu Tristram committed
656
                    plt.contour( X,Y, hmap, levels=ctr_level( hmap, [0.68,0.95]),
Matthieu Tristram's avatar
Matthieu Tristram committed
657
                                 colors=colors[c], linestyles=linestyles[c], extent=tuple(rge[par]+rge[par2]))
658

659 660 661 662
                if bestfit:
                    plt.plot( (bestfit[c])[par2], (bestfit[c])[par], '+', color=colors[c])


Matthieu Tristram's avatar
Matthieu Tristram committed
663
#            ax.locator_params( nbins=npar-1)
Matthieu Tristram's avatar
Matthieu Tristram committed
664
            ax.ticklabel_format(useOffset=False)
665 666
#            ax.tick_params(axis='both', which='major', labelsize=10)
#            ax.tick_params(axis='both', which='minor', labelsize=5)
667 668
            ax.yaxis.set_visible(ypar==0)
            ax.xaxis.set_visible(xpar==(npar-1))
669
            ax.locator_params(tight=True, nbins=5)
670
            if 'PS' in parname.get(parnames[ypar],parnames[ypar]):
671
                ax.xaxis.set_major_formatter(mtick.FormatStrFormatter('%.0e'))
Matthieu Tristram's avatar
Matthieu Tristram committed
672 673

    #plot parameter names
674
    for xpar in range(npar):
675 676 677 678 679 680 681
        ax=plt.subplot( npar, npar, xpar*npar+1)
        ax.set_ylabel(parname.get(parnames[xpar],parnames[xpar]))
        ax.tick_params(axis='both', which='major', labelsize=fontsize)
        ax=plt.subplot( npar, npar, (npar-1)*npar+xpar+1)
        ax.set_xlabel(parname.get(parnames[xpar],parnames[xpar]))
        ax.tick_params(axis='both', which='major', labelsize=fontsize)
#        plt.locator_params(tight=True, nbins=5)
682

Matthieu Tristram's avatar
Matthieu Tristram committed
683
    #write legend
684
    if len(names) == nchain:
Matthieu Tristram's avatar
Matthieu Tristram committed
685
        ax=plt.subplot( npar, npar, int(0.25*npar)*npar+0.85*npar-1)
Matthieu Tristram's avatar
Matthieu Tristram committed
686
        for c in range(nchain):
687
            plt.text(0.5,1.-0.2*(c+1), names[c], color=colors[c])
Matthieu Tristram's avatar
Matthieu Tristram committed
688
        ax.axis('off')
689

Matthieu Tristram's avatar
Matthieu Tristram committed
690
    return(fig)
Matthieu Tristram's avatar
Matthieu Tristram committed
691 692 693



Matthieu Tristram's avatar
Matthieu Tristram committed
694 695
def rectangle( chains, par0, params, nbin=50, parnames=None, names=[], colors=[], linestyles=[], smooth=1., fontsize=12,
               contour=True, fill=False, cmaps=[], Norm=True, weights=None, ylim=None):
696
    import matplotlib.ticker as mtick
Matthieu Tristram's avatar
Matthieu Tristram committed
697 698 699 700 701 702 703 704

    #force tuple
    if not( isinstance(chains, list) or isinstance(chains, tuple)):
        chains = (chains,)

    nchain = len(chains)
    npar   = len(params)

Matthieu Tristram's avatar
Matthieu Tristram committed
705 706 707
    fig=plt.figure( figsize=(5*npar,5))
    plt.subplots_adjust(hspace=0.001,wspace=0.001)

Matthieu Tristram's avatar
Matthieu Tristram committed
708 709 710 711 712 713 714 715 716 717 718 719 720
    if len(colors) < nchain:
        colors = plt.rcParams['axes.color_cycle']

    if len(cmaps) < nchain:
        cmaps = ['Blues','Greens','Greys','Oranges','Purples']

    if type(fill) == bool:
        fill = [fill]*nchain
    if type(contour) == bool:
        contour = [contour]*nchain

    if len(linestyles) < nchain:
        linestyles = ['-']*nchain
721

Matthieu Tristram's avatar
Matthieu Tristram committed
722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766
    if weights == None:
        weights = [np.ones(len(c['chi2'])) for c in chains]

    if parnames == None:
        parnames = [parname.get(p,p) for p in params]

    if ylim == None:
        tmp = []
        for chain in chains:
            tmp = tmp+list(chain[par0])
        ylim = list(np.mean(tmp) + np.dot(np.std(tmp), [-4,4]))

    rge = {}
    for par in params:
        tmp = []
        for chain in chains:
            tmp = tmp+list(chain[par])
        rge[par] = list(np.mean(tmp) + np.dot(np.std(tmp), [-4,4]))


    for par in params:
        xpar = params.index(par)

        ax=plt.subplot( 1, npar, xpar+1)

        for c in range(nchain):
            chain = chains[c]
            weight = weights[c]
            hmap,binX,binY = np.histogram2d( chain[par], chain[par0], bins=nbin, range=[rge[par],ylim], weights=weight)
            X,Y = np.meshgrid((binX[1:]+binX[:-1])/2.,(binY[1:]+binY[:-1])/2.)
            hmap = nd.gaussian_filter( hmap.T, 2)
            if fill[c]:
                plt.contourf( X,Y, hmap, levels=np.append(list(ctr_level( hmap, [0.68,0.95])[::-1]),hmap.max()),
                              colors=None, extent=tuple(rge[par]+ylim), cmap=plt.cm.get_cmap(cmaps[c]))
            if contour[c]:
                plt.contour( X,Y, hmap, levels=ctr_level( hmap, [0.68,0.95]),
                             colors=colors[c], linestyles=linestyles[c], extent=tuple(rge[par]+ylim))

        ax.ticklabel_format(useOffset=False)
        ax.locator_params(tight=True, nbins=5)
        ax.set_xlabel(parname.get(parnames[xpar],parnames[xpar]))
        ax.tick_params(axis='both', which='major', labelsize=fontsize)
        ax.yaxis.set_visible(xpar==0)
        if xpar == 0:
            ax.set_ylabel(parname.get(par0,par0))
767

Matthieu Tristram's avatar
Matthieu Tristram committed
768 769
#        if 'PS' in parname.get(parnames[par],parnames[par]):
#            ax.xaxis.set_major_formatter(mtick.FormatStrFormatter('%.0e'))
770

Matthieu Tristram's avatar
Matthieu Tristram committed
771
    #write legend
Matthieu Tristram's avatar
Matthieu Tristram committed
772 773 774
    if len(names) == nchain:
        for c in range(nchain):
            ax.text( 0.9, 1.-0.1*(c+1), names[c], color=colors[c], transform=ax.transAxes, horizontalalignment='right')
Matthieu Tristram's avatar
Matthieu Tristram committed
775 776 777 778 779

    return(fig)



780 781 782 783
def hist1d(x, bins=50, weights=None, smooth=0., normmax=True, **kwargs):
    extents = kwargs.get("extents",[x.min(),x.max()])

    n, b = np.histogram(x, weights=weights, bins=bins,normed=True)
784
    ns = nd.gaussian_filter1d(n,smooth,mode='nearest')
785 786
    if normmax:
        ns /= ns.max()
Matthieu Tristram's avatar
Matthieu Tristram committed
787 788 789

    xx = (b[:-1]+b[1:])/2.
    yy = ns
790 791 792
    plt.plot( xx, yy, **kwargs)
    plt.xlim( extents)

793 794


Matthieu Tristram's avatar
Matthieu Tristram committed
795 796 797 798
def MCcorrelation( chain, par, **kwargs):
    data = [chain[p] for p in par]

    matC = np.corrcoef( data)
799

Matthieu Tristram's avatar
Matthieu Tristram committed
800 801 802 803
    plot_correlation( matC, par, **kwargs)


def plot_correlation( matC, par, **kwargs):
Matthieu Tristram's avatar
Matthieu Tristram committed
804 805 806 807
    vmin = kwargs.pop("vmin", -1)
    vmax = kwargs.pop("vmax",  1)

    plt.figure( figsize=(14,10))
808
    #covmat = np.triu(np.corrcoef( chain.values()[1:]), k=1)
809
    covmat = np.tril( matC, k=-1)
810
    covmat[covmat==0] = np.NaN
Matthieu Tristram's avatar
Matthieu Tristram committed
811
    plt.imshow( covmat, vmin=vmin, vmax=vmax, **kwargs)
812 813
    plt.yticks( list(range(len(covmat))), [parname.get(p,p) for p in par])
    plt.xticks( list(range(len(covmat))), '')
Matthieu Tristram's avatar
Matthieu Tristram committed
814
    plt.colorbar()
815
    for i in range( len(par[:-1])):
816 817
        #    text( i, i+0.25, parname.get(par[i+1],par[i+1]),rotation=-90)
        plt.text( i-0.25, i+0.25, parname.get(par[i],par[i]),rotation='vertical',verticalalignment='bottom')
818

Matthieu Tristram's avatar
Matthieu Tristram committed
819 820


Matthieu Tristram's avatar
Matthieu Tristram committed
821 822
########################################################################

823 824


Matthieu Tristram's avatar
Matthieu Tristram committed
825

Matthieu Tristram's avatar
Matthieu Tristram committed
826 827
########################################################################

Matthieu Tristram's avatar
Matthieu Tristram committed
828
def ctr_level(histo2d, lvl):
829
    """
Matthieu Tristram's avatar
Matthieu Tristram committed
830
    Extract the contours for the 2d plots
831 832
    """

833 834 835 836
    h = histo2d.flatten()*1.
    h.sort()
    cum_h = np.cumsum(h[::-1])
    cum_h /= cum_h[-1]
837

838 839
    alvl = np.searchsorted(cum_h, lvl)
    clist = h[-alvl]
840

Matthieu Tristram's avatar
Matthieu Tristram committed
841
    return clist[::-1]
842

Matthieu Tristram's avatar
Matthieu Tristram committed
843 844 845 846 847

def FeldmanCousins( xmin, sigma, CL=95):
    import os
    from scipy.interpolate import interp1d

848
    fcdata = np.loadtxt( os.getenv("CAMELROOT")+"/work/tools/FeldmanCousins.data").T
Matthieu Tristram's avatar
Matthieu Tristram committed
849 850 851 852 853 854 855
    listCL = [68,90,95,99]

    if not(CL in listCL):
        raise ValueError('CL not in the list')

    clmin = interp1d( fcdata[0], fcdata[listCL.index(CL)*2+1])
    clmax = interp1d( fcdata[0], fcdata[listCL.index(CL)*2+2])
856

Matthieu Tristram's avatar
Matthieu Tristram committed
857 858
    x0=xmin/sigma
    uplim = clmax(x0)
859

Matthieu Tristram's avatar
Matthieu Tristram committed
860 861
    return( uplim*sigma)

862 863 864 865 866 867 868 869 870
########################################################################







########################################################################
871
def read_bestfit( dirname):
872
    if '.gz' in dirname :
873 874 875 876
        f= gzip.open(dirname      , 'rb')
    else :
        f = open( dirname, 'r')

877
    param = np.array( f.readline().split(), dtype=str)
878

879
    newline =  f.readline()
880
    if  not "nan" in newline and not "inf" in newline and not "Abort" in newline:
881
        value = np.array( newline.split(), dtype=np.double)
Xavier Garrido's avatar