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
    print(l[index].replace("\n",""))
127

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
        else:
            #Gelman-Runbin stat on the last half samples for each iteration
235 236
#SP fix for python3: add int
            n = int((isamp - length_min)/2)
Matthieu Tristram's avatar
Matthieu Tristram committed
237
            tmp = [data[1:,length_min+n:isamp] for data in chains]
238 239

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

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

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

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


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

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

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

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

273
        #Tscore
Matthieu Tristram's avatar
Matthieu Tristram committed
274 275 276
        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) )
277

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

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





287 288
########################################################################
#PARAMETER NAMES
Matthieu Tristram's avatar
Matthieu Tristram committed
289 290
########################################################################
parname = {
291
    'omega_b'              : '$\\Omega_\mathrm{b}h^2$',
292
    'omega_cdm'            : '$\\Omega_\mathrm{c}h^2$',
293
    'Omega_b'              : '$\\Omega_\mathrm{b}$',
294
    'Omega_cdm'            : '$\\Omega_\mathrm{c}$',
Matthieu Tristram's avatar
Matthieu Tristram committed
295 296 297 298 299
    '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})$',
300 301
    'H0'                   : '$H_0$',
    'sigma8'               : '$\\sigma_8$',
302
    'Alens'                : '$\mathrm{A}_\mathrm{L}$',
303
    'm_ncdm'               : '$\\Sigma m_\\nu$',
304 305
    'N_eff'                : '$N_\mathrm{eff}$',
    'r'                    : '$r$',
Matthieu Tristram's avatar
Matthieu Tristram committed
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 337
    '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}$',
338 339
    'AdustTP'              : '$A_\mathrm{dust}^\mathrm{TE}$',
    'AdustPP'              : '$A_\mathrm{dust}^\mathrm{EE}$',
perdereau's avatar
perdereau committed
340
    'SPT_ADust'            : '$A^\mathrm{dust}_\mathrm{SPT}$',
Matthieu Tristram's avatar
Matthieu Tristram committed
341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365
    '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
366 367


368
to_cosmomcName = {
Matthieu Tristram's avatar
Matthieu Tristram committed
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 403
    '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
404

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


Matthieu Tristram's avatar
Matthieu Tristram committed
408 409

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

Matthieu Tristram's avatar
Matthieu Tristram committed
417
    fig=plt.figure(figsize=figsize)
418
    if len(colors) < nchain:
419 420 421
#        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
422
        colors = plt.rcParams['axes.color_cycle']
423 424 425 426 427 428 429

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

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

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

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

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

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

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

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

455 456
    #legend
    if names != []:
Matthieu Tristram's avatar
Matthieu Tristram committed
457
#        if len(params) >= 2:
Matthieu Tristram's avatar
Matthieu Tristram committed
458 459
        if len(params) ==1:
            ax.legend(names, loc='upper center')
460
        else:
Matthieu Tristram's avatar
Matthieu Tristram committed
461 462 463 464 465 466
            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
467
#            ax.legend(names, loc='upper right')
468 469


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

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

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

Matthieu Tristram's avatar
Matthieu Tristram committed
481 482 483 484 485 486 487 488
#    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)
489
    contours   = kwargs.get("contours",   False)
Matthieu Tristram's avatar
Matthieu Tristram committed
490
    levels     = kwargs.get("levels",   [0.68,0.95])
Matthieu Tristram's avatar
Matthieu Tristram committed
491 492 493
    cmap       = kwargs.get("cmap",        None)

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

Matthieu Tristram's avatar
Matthieu Tristram committed
495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510
    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.)
511 512
    dy = (extent[1][1]-extent[1][0])/(bins+1.)

Matthieu Tristram's avatar
Matthieu Tristram committed
513 514
    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)
515

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

Matthieu Tristram's avatar
Matthieu Tristram committed
518 519 520 521 522
    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)
523

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

Matthieu Tristram's avatar
Matthieu Tristram committed
527 528 529 530 531 532 533 534 535 536 537 538 539 540 541
    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]
542

Matthieu Tristram's avatar
Matthieu Tristram committed
543 544 545 546
    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)
547

Matthieu Tristram's avatar
Matthieu Tristram committed
548 549 550 551 552 553 554 555 556 557 558 559
    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)
560

Matthieu Tristram's avatar
Matthieu Tristram committed
561
    npts = 5.
562 563 564 565
#    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)
566

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

570 571


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

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

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

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

586
    if len(colors) < nchain:
587 588 589
#        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
590
        colors = plt.rcParams['axes.color_cycle']
591

Matthieu Tristram's avatar
Matthieu Tristram committed
592 593 594 595 596 597 598 599 600
    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

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

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

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

    rge = {}
    for par in params:
        tmp = []
        for chain in chains:
            tmp = tmp+list(chain[par])
Matthieu Tristram's avatar
Matthieu Tristram committed
615 616
        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
617 618


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

Matthieu Tristram's avatar
Matthieu Tristram committed
623 624
        for c in range(nchain):
            chain = chains[c]
625 626
            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
627
            Y = nd.gaussian_filter1d( h, smooth,mode='nearest')
628 629
            Y = Y/np.max(Y) if Norm else Y
            plt.plot( (X[1:]+X[:-1])/2., Y, color=colors[c], linestyle=linestyles[c])
630

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

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

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

643 644 645
        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
646 647 648

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

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


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

    #plot parameter names
675
    for xpar in range(npar):
676 677 678 679 680 681 682
        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)
683

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

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



Matthieu Tristram's avatar
Matthieu Tristram committed
695 696
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):
697
    import matplotlib.ticker as mtick
Matthieu Tristram's avatar
Matthieu Tristram committed
698 699 700 701 702 703 704 705

    #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
706 707 708
    fig=plt.figure( figsize=(5*npar,5))
    plt.subplots_adjust(hspace=0.001,wspace=0.001)

Matthieu Tristram's avatar
Matthieu Tristram committed
709 710 711 712 713 714 715 716 717 718 719 720 721
    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
722

Matthieu Tristram's avatar
Matthieu Tristram committed
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 767
    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))
768

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

Matthieu Tristram's avatar
Matthieu Tristram committed
772
    #write legend
Matthieu Tristram's avatar
Matthieu Tristram committed
773 774 775
    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
776 777 778 779 780

    return(fig)



781 782 783 784
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)
785
    ns = nd.gaussian_filter1d(n,smooth,mode='nearest')
786 787
    if normmax:
        ns /= ns.max()
Matthieu Tristram's avatar
Matthieu Tristram committed
788 789 790

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

794 795


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

    matC = np.corrcoef( data)
800

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


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

    plt.figure( figsize=(14,10))
809
    #covmat = np.triu(np.corrcoef( chain.values()[1:]), k=1)
810
    covmat = np.tril( matC, k=-1)
811
    covmat[covmat==0] = np.NaN
Matthieu Tristram's avatar
Matthieu Tristram committed
812
    plt.imshow( covmat, vmin=vmin, vmax=vmax, **kwargs)
813 814
    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
815
    plt.colorbar()
816
    for i in range( len(par[:-1])):
817 818
        #    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')
819

Matthieu Tristram's avatar
Matthieu Tristram committed
820 821


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

824 825


Matthieu Tristram's avatar
Matthieu Tristram committed
826

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

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

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

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

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

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

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

849
    fcdata = np.loadtxt( os.getenv("CAMELROOT")+"/work/tools/FeldmanCousins.data").T
Matthieu Tristram's avatar
Matthieu Tristram committed
850 851 852 853 854 855 856
    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])
857

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

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

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







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

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

880
    newline =  f.readline()
881
    if  not "nan" in newline and not "inf" in newline and not "Abort" in newline: