diff options
| -rwxr-xr-x | chainer_plot.py | 10 | ||||
| -rwxr-xr-x | fr.py | 249 | ||||
| -rwxr-xr-x | mcmc_scan.py | 671 | ||||
| -rw-r--r-- | utils/__init__.py | 0 | ||||
| -rw-r--r-- | utils/enums.py | 68 | ||||
| -rw-r--r-- | utils/fr.py | 353 | ||||
| -rw-r--r-- | utils/gf.py | 105 | ||||
| -rw-r--r-- | utils/mcmc.py | 167 | ||||
| -rw-r--r-- | utils/misc.py | 230 |
9 files changed, 1177 insertions, 676 deletions
diff --git a/chainer_plot.py b/chainer_plot.py index 8964bf5..3facb30 100755 --- a/chainer_plot.py +++ b/chainer_plot.py @@ -17,7 +17,7 @@ import getdist from getdist import plots from getdist import mcsamples -import mcmc_scan +from utils.fr import angles_to_u, angles_to_fr rc('text', usetex=False) rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18}) @@ -116,14 +116,14 @@ def plot(infile, angles, outfile, measured_ratio, sigma_ratio, fix_sfr, print 'ranges', ranges def flat_angles_to_u(x): - return abs(mcmc_scan.angles_to_u(x)).astype(np.float32).flatten().tolist() + return abs(angles_to_u(x)).astype(np.float32).flatten().tolist() raw = np.load(infile) print 'raw.shape', raw.shape if not angles: nuisance, raw = raw[:,5:], raw[:,-5:] if fix_mixing: - fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + fr_elements = np.array(map(angles_to_fr, raw[:,-2:])) sc_elements = raw[:,:-2] Tchain = np.column_stack([sc_elements, fr_elements]) elif fix_sfr: @@ -135,11 +135,11 @@ def plot(infile, angles, outfile, measured_ratio, sigma_ratio, fix_sfr, Tchain = np.column_stack([m_elements, sc_elements]) else: if fix_scale: - fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + fr_elements = np.array(map(angles_to_fr, raw[:,-2:])) m_elements = np.array(map(flat_angles_to_u, raw[:,:-2])) Tchain = np.column_stack([m_elements, fr_elements]) else: - fr_elements = np.array(map(mcmc_scan.angles_to_fr, raw[:,-2:])) + fr_elements = np.array(map(angles_to_fr, raw[:,-2:])) sc_elements = raw[:,-3:-2] m_elements = np.array(map(flat_angles_to_u, raw[:,:-3])) Tchain = np.column_stack([m_elements, sc_elements, fr_elements]) @@ -0,0 +1,249 @@ +#! /usr/bin/env python +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +HESE BSM flavour ratio analysis script +""" + +from __future__ import absolute_import, division + +import argparse +from functools import partial + +import numpy as np + +from utils import mcmc as mcmc_utils +from utils import misc as misc_utils +from utils.fr import MASS_EIGENVALUES, normalise_fr +from utils.gf import gf_argparse +from utils.misc import Param, ParamSet + +import chainer_plot + + +def define_nuisance(): + """Define the nuisance parameters to scan over with default values, + ranges and sigma. + """ + nuisance = ParamSet( + Param(name='convNorm' , value=1. , ranges=[0., 5.], std=0.3), + Param(name='promptNorm' , value=0. , ranges=[0., 5.], std=0.05), + Param(name='muonNorm' , value=1. , ranges=[0., 5.], std=0.1), + Param(name='astroNorm' , value=1. , ranges=[0., 5.], std=0.1), + Param(name='astroDeltaGamma', value=2. , ranges=[0., 5.], std=0.1) + ) + return nuisance + + +def nuisance_argparse(parser): + nuisance_paramset = define_nuisance() + for parm in nuisance_paramset: + parser.add_argument( + '--'+parm.name, type=float, default=parm.value, + help=parm.name+' to inject' + ) + + +def get_paramsets(args): + """Make the paramsets for generating the Asmimov MC sample and also running + the MCMC. + """ + nuisance_paramset = define_nuisance() + for parm in nuisance_paramset: + parm.value = args.__getattribute__(parm.name) + asimov_paramset = ParamSet( + nuisance_paramset.params + + [Param(name='astroENorm' , value=args.measured_ratio[0], ranges=[0., 1.], std=0.2), + Param(name='astroMuNorm' , value=args.measured_ratio[1], ranges=[0., 1.], std=0.2), + Param(name='astroTauNorm' , value=args.measured_ratio[2], ranges=[0., 1.], std=0.2)] + ) + + mcmc_paramset = [] + if not args.fix_mixing: + mcmc_paramset.extend([ + Param(name='s_12^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{12}^2'), + Param(name='c_13^4', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{c}_{13}^4'), + Param(name='s_23^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{23}^4'), + Param(name='dcp', value=np.pi, ranges=[0., 2*np.pi], std=0.2, tex=r'\tilde{\delta_{CP}}') + ]) + if not args.fix_scale: + logLam, scale_region = np.log10(args.scale), np.log10(args.scale_region) + lL_range = (logLam-scale_region, logLam+scale_region) + mcmc_paramset.append( + Param(name='logLam', value=logLam, ranges=lL_range, std=3, tex=r'{\rm log}_{10}\Lambda') + ) + if not args.fix_source_ratio: + mcmc_paramset.extend([ + Param(name='s_phi4', value=0.5, ranges=[0., 1.], std=0.2, tex=r'sin^4(\phi)'), + Param(name='c_2psi', value=0.5, ranges=[0., 1.], std=0.2, tex=r'cos(2\psi)') + ]) + mcmc_paramset = ParamSet(nuisance_paramset.params + mcmc_paramset) + return nuisance_paramset, mcmc_paramset + + +def process_args(args): + """Process the input args.""" + if args.fix_mixing and args.fix_source_ratio: + raise NotImplementedError('Fixed mixing and sfr not implemented') + if args.fix_mixing and args.fix_scale: + raise NotImplementedError('Fixed mixing and scale not implemented') + + args.measured_ratio = normalise_fr(args.measured_ratio) + if args.fix_source_ratio: + args.source_ratio = normalise_fr(args.source_ratio) + + if not args.fix_scale: + args.scale = np.power( + 10, np.round(np.log10(MASS_EIGENVALUES[1]/args.energy)) - \ + np.log10(args.energy**(args.dimension-3)) + ) + """Estimate the scale of when to expect the BSM term to contribute.""" + + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description="BSM flavour ratio analysis") + parser.add_argument( + '--measured-ratio', type=int, nargs=3, default=[1, 1, 1], + help='Set the central value for the measured flavour ratio at IceCube' + ) + parser.add_argument( + '--sigma-ratio', type=float, default=0.01, + help='Set the 1 sigma for the measured flavour ratio for a gaussian LLH' + ) + parser.add_argument( + '--fix-source-ratio', type=misc_utils.parse_bool, default='False', + help='Fix the source flavour ratio' + ) + parser.add_argument( + '--source-ratio', type=int, nargs=3, default=[2, 1, 0], + help='Set the source flavour ratio for the case when you want to fix it' + ) + parser.add_argument( + '--no-bsm', type=misc_utils.parse_bool, default='False', + help='Turn off BSM terms' + ) + parser.add_argument( + '--fix-mixing', type=misc_utils.parse_bool, default='False', + help='Fix all mixing parameters except one' + ) + parser.add_argument( + '--fix-scale', type=misc_utils.parse_bool, default='False', + help='Fix the new physics scale' + ) + parser.add_argument( + '--scale', type=float, default=1e-30, + help='Set the new physics scale' + ) + parser.add_argument( + '--scale-region', type=float, default=1e10, + help='Set the size of the box to scan for the scale' + ) + parser.add_argument( + '--dimension', type=int, default=3, + help='Set the new physics dimension to consider' + ) + parser.add_argument( + '--energy', type=float, default=1000, + help='Set the energy scale' + ) + parser.add_argument( + '--seed', type=int, default=99, + help='Set the random seed value' + ) + parser.add_argument( + '--threads', type=misc_utils.thread_type, default='1', + help='Set the number of threads to use (int or "max")' + ) + parser.add_argument( + '--outfile', type=str, default='./untitled', + help='Path to output chains' + ) + nuisance_argparse(parser) + gf_argparse(parser) + mcmc_utils.mcmc_argparse(parser) + return parser.parse_args() + + +def main(): + args = parse_args() + process_args(args) + misc_utils.print_args(args) + + np.random.seed(args.seed) + + asmimov_paramset, mcmc_paramset = get_paramsets(args) + outfile = misc_utils.gen_outfile_name(args) + print '== {0:<25} = {1}'.format('outfile', outfile) + + if args.run_mcmc: + triangle_llh = partial( + mcmc_utils.triangle_llh, args=args + ) + lnprior = partial( + mcmc_utils.lnprior, paramset=mcmc_paramset + ) + + ndim = len(mcmc_paramset) + ntemps=1 + p0 = mcmc_utils.gaussian_seed( + mcmc_paramset, ntemps=ntemps, nwalkers=args.nwalkers + ) + + samples = mcmc_utils.mcmc( + p0 = p0, + triangle_llh = triangle_llh, + lnprior = lnprior, + ndim = ndim, + nwalkers = args.nwalkers, + burnin = args.burnin, + nsteps = args.nsteps, + ntemps = ntemps, + threads = args.threads + ) + mcmc_utils.save_chains(samples, outfile) + + scale_bounds = (args.scale/args.scale_region, args.scale*args.scale_region) + print "Making triangle plots" + chainer_plot.plot( + infile = outfile+'.npy', + angles = True, + outfile = outfile[:5]+outfile[5:].replace('data', 'plots')+'_angles.pdf', + measured_ratio = args.measured_ratio, + sigma_ratio = args.sigma_ratio, + fix_sfr = args.fix_source_ratio, + fix_mixing = args.fix_mixing, + fix_scale = args.fix_scale, + source_ratio = args.source_ratio, + scale = args.scale, + dimension = args.dimension, + energy = args.energy, + scale_bounds = scale_bounds, + ) + # chainer_plot.plot( + # infile=outfile+'.npy', + # angles=False, + # outfile=outfile[:5]+outfile[5:].replace('data', 'plots')+'.pdf', + # measured_ratio=args.measured_ratio, + # sigma_ratio=args.sigma_ratio, + # fix_sfr=args.fix_source_ratio, + # fix_mixing=args.fix_mixing, + # fix_scale=FIX_SCALE, + # source_ratio=args.source_ratio, + # scale=args.scale, + # dimension=args.dimension, + # energy=args.energy, + # scale_bounds=scale_bounds, + # ) + print "DONE!" + + +main.__doc__ = __doc__ + + +if __name__ == '__main__': + main() + diff --git a/mcmc_scan.py b/mcmc_scan.py deleted file mode 100755 index 90003a9..0000000 --- a/mcmc_scan.py +++ /dev/null @@ -1,671 +0,0 @@ -#! /usr/bin/env python -""" -From a gaussian likelihood run an MCMC scan to find the posteriors -""" - -from __future__ import absolute_import, division - -import os, sys -import errno - -import argparse -import multiprocessing -from copy import deepcopy - -import numpy as np -from scipy import linalg - -import emcee -import tqdm - -import GolemFitPy as gf -import chainer_plot - - -RUN_SCAN = True - -GOLEMFIT = None -FIX_MIXING = False -FIX_SFR = True -SOURCE_FR = [1, 2, 0] -FIX_SCALE = True -NO_BSM = False - -DIMENSION = 3 -ENERGY = 1000000 # GeV -MEASURED_FR = [1, 1, 1] -SIGMA = 0.001 -SCALE_RATIO = 100. -MASS_EIGENVALUES = [7.40E-23, 2.515E-21] -# MASS_EIGENVALUES = [7.40E-100, 2.515E-100] -FLAT = False - -CHANGE_MASS = False -if CHANGE_MASS: - SCALE = 1E-27 -else: - SCALE = np.power(10, np.round(np.log10(MASS_EIGENVALUES[1]/ENERGY)) - np.log10(ENERGY**(DIMENSION-3))) -SCALE2_BOUNDS = (SCALE*1E-10, SCALE*1E10) - - -def set_up_gf(): - steering_params = gf.SteeringParams() - datapaths = gf.DataPaths() - npp = gf.NewPhysicsParams() - - steering_params.logEbinWidth=(7.0-4.78)/20. - steering_params.cosThbinWidth=0.2 - steering_params.correct_prompt_knee=True - steering_params.fastmode=True - # TODO(shivesh): figure out right number for missid - steering_params.track_to_shower_missid = 0.3 - golemfitsourcepath = os.environ['GOLEMSOURCEPATH'] + '/GolemFit' - datapaths.compact_file_path = golemfitsourcepath + '/compact_data/' # only used for FastMode; not currently working - datapaths.data_path = golemfitsourcepath + '/data/' - datapaths.mc_path = golemfitsourcepath + '/monte_carlo/' - - golemfit = gf.GolemFit(datapaths, steering_params, npp) - return golemfit - - -def set_up_as(golemfit, parms): - asimov_parameters = gf.FitParameters() - asimov_parameters.convNorm = parms['convNorm'] - asimov_parameters.promptNorm = parms['promptNorm'] - asimov_parameters.muonNorm = parms['muonNorm'] - asimov_parameters.astroNorm = parms['astroNorm'] - asimov_parameters.astroDeltaGamma = parms['astroDeltaGamma'] - asimov_parameters.astroENorm = parms['astroENorm'] - asimov_parameters.astroMuNorm = parms['astroMuNorm'] - asimov_parameters.astroTauNorm = parms['astroTauNorm'] - golemfit.SetupAsimov(asimov_parameters) - - -def get_llh(golemfit, parms): - fitparams = gf.FitParameters() - fitparams.convNorm = parms['convNorm'] - fitparams.promptNorm = parms['promptNorm'] - fitparams.muonNorm = parms['muonNorm'] - fitparams.astroNorm = parms['astroNorm'] - fitparams.astroDeltaGamma = parms['astroDeltaGamma'] - fitparams.astroENorm = parms['astroENorm'] - fitparams.astroMuNorm = parms['astroMuNorm'] - fitparams.astroTauNorm = parms['astroTauNorm'] - return golemfit.EvalLLH(fitparams) - - -def test_uni(x): - """Test the unitarity of a matrix.""" - # print 'Unitarity test:\n{0}'.format(abs(np.dot(x, x.conj().T))) - return abs(np.dot(x, x.conj().T)) - - -def angles_to_fr(angles): - sphi4, c2psi = angles - - psi = (0.5)*np.arccos(c2psi) - - sphi2 = np.sqrt(sphi4) - cphi2 = 1. - sphi2 - spsi2 = np.sin(psi)**2 - cspi2 = 1. - spsi2 - - x = sphi2*cspi2 - y = sphi2*spsi2 - z = cphi2 - return x, y, z - - -def angles_to_u(angles): - s12_2, c13_4, s23_2, dcp = angles - dcp = np.complex128(dcp) - - c12_2 = 1. - s12_2 - c13_2 = np.sqrt(c13_4) - s13_2 = 1. - c13_2 - c23_2 = 1. - s23_2 - - t12 = np.arcsin(np.sqrt(s12_2)) - t13 = np.arccos(np.sqrt(c13_2)) - t23 = np.arcsin(np.sqrt(s23_2)) - - c12 = np.cos(t12) - s12 = np.sin(t12) - c13 = np.cos(t13) - s13 = np.sin(t13) - c23 = np.cos(t23) - s23 = np.sin(t23) - - p1 = np.array([[1 , 0 , 0] , [0 , c23 , s23] , [0 , -s23 , c23]] , dtype=np.complex128) - p2 = np.array([[c13 , 0 , s13*np.exp(-1j*dcp)] , [0 , 1 , 0] , [-s13*np.exp(1j*dcp) , 0 , c13]] , dtype=np.complex128) - p3 = np.array([[c12 , s12 , 0] , [-s12 , c12 , 0] , [0 , 0 , 1]] , dtype=np.complex128) - - u = np.dot(np.dot(p1, p2), p3) - return u - - -def cardano_eqn(ham): - """Diagonalise the effective Hamiltonian 3x3 matrix into the form - h_{eff} = UE_{eff}U^{dagger} using the procedure in PRD91, 052003 (2015) - """ - a = -np.trace(ham) - b = (0.5) * ((np.trace(ham))**2 - np.trace(np.dot(ham, ham))) - c = -linalg.det(ham) - - Q = (1/9.) * (a**2 - 3*b) - R = (1/54.) * (2*a**3 - 9*a*b + 27*c) - theta = np.arccos(R / np.sqrt(Q**3)) - - E1 = -2 * np.sqrt(Q) * np.cos(theta/3.) - (1/3.)*a - E2 = -2 * np.sqrt(Q) * np.cos((theta - 2.*np.pi)/3.) - (1/3.)*a - E3 = -2 * np.sqrt(Q) * np.cos((theta + 2.*np.pi)/3.) - (1/3.)*a - - A1 = ham[1][2] * (ham[0][0] - E1) - ham[1][0]*ham[0][2] - A2 = ham[1][2] * (ham[0][0] - E2) - ham[1][0]*ham[0][2] - A3 = ham[1][2] * (ham[0][0] - E3) - ham[1][0]*ham[0][2] - - B1 = ham[2][0] * (ham[1][1] - E1) - ham[2][1]*ham[1][0] - B2 = ham[2][0] * (ham[1][1] - E2) - ham[2][1]*ham[1][0] - B3 = ham[2][0] * (ham[1][1] - E3) - ham[2][1]*ham[1][0] - - C1 = ham[1][0] * (ham[2][2] - E1) - ham[1][2]*ham[2][0] - C2 = ham[1][0] * (ham[2][2] - E2) - ham[1][2]*ham[2][0] - C3 = ham[1][0] * (ham[2][2] - E3) - ham[1][2]*ham[2][0] - - N1 = np.sqrt(abs(A1*B1)**2 + abs(A1*C1)**2 + abs(B1*C1)**2) - N2 = np.sqrt(abs(A2*B2)**2 + abs(A2*C2)**2 + abs(B2*C2)**2) - N3 = np.sqrt(abs(A3*B3)**2 + abs(A3*C3)**2 + abs(B3*C3)**2) - - mm = np.array([ - [np.conjugate(B1)*C1 / N1, np.conjugate(B2)*C2 / N2, np.conjugate(B3)*C3 / N3], - [A1*C1 / N1, A2*C2 / N2, A3*C3 / N3], - [A1*B1 / N1, A2*B2 / N2, A3*B3 / N3] - ]) - return mm - - -# s_12^2, c_13^4, s_23^2, dcp -NUFIT_U = angles_to_u((0.307, (1-0.02195)**2, 0.565, 3.97935)) - - -def params_to_BSMu(theta): - if FIX_MIXING: - s12_2, c13_4, s23_2, dcp, sc2 = 0.5, 1.0-1E-6, 0.5, 0., theta - elif FIX_SCALE: - s12_2, c13_4, s23_2, dcp = theta - sc2 = np.log10(SCALE) - else: - s12_2, c13_4, s23_2, dcp, sc2 = theta - sc2 = np.power(10., sc2) - sc1 = sc2 / SCALE_RATIO - - mass_matrix = np.array( - [[0, 0, 0], [0, MASS_EIGENVALUES[0], 0], [0, 0, MASS_EIGENVALUES[1]]] - ) - sm_ham = (1./(2*ENERGY))*np.dot(NUFIT_U, np.dot(mass_matrix, NUFIT_U.conj().T)) - if NO_BSM: - eg_vector = cardano_eqn(sm_ham) - tu = test_uni(eg_vector) - if not abs(np.trace(tu) - 3.) < 1e-5 or not abs(np.sum(tu) - 3.) < 1e-5: - raise AssertionError('Matrix is not unitary!\neg_vector\n{0}\ntest u\n{1}'.format(eg_vector, tu)) - return eg_vector - - new_physics_u = angles_to_u((s12_2, c13_4, s23_2, dcp)) - scale_matrix = np.array( - [[0, 0, 0], [0, sc1, 0], [0, 0, sc2]] - ) - bsm_term = (ENERGY**(DIMENSION-3)) * np.dot(new_physics_u, np.dot(scale_matrix, new_physics_u.conj().T)) - - bsm_ham = sm_ham + bsm_term - - eg_vector = cardano_eqn(bsm_ham) - tu = test_uni(eg_vector) - if not abs(np.trace(tu) - 3.) < 1e-5 or not abs(np.sum(tu) - 3.) < 1e-5: - raise AssertionError('Matrix is not unitary!\neg_vector\n{0}\ntest u\n{1}'.format(eg_vector, tu)) - return eg_vector - - -def u_to_fr(initial_fr, matrix): - """Compute the observed flavour ratio assuming decoherence.""" - # TODO(shivesh): energy dependence - composition = np.einsum( - 'ai, bi, a -> b', abs(matrix)**2, abs(matrix)**2, initial_fr - ) - ratio = composition / np.sum(initial_fr) - return ratio - - -def triangle_llh(theta): - """-Log likelihood function for a given theta.""" - convNorm, promptNorm, muonNorm, astroNorm, astroDeltaGamma = theta[:5] - theta = deepcopy(theta[5:]) - - if FIX_SFR: - fr1, fr2, fr3 = SOURCE_FR - u = params_to_BSMu(theta) - else: - fr1, fr2, fr3 = angles_to_fr(theta[-2:]) - u = params_to_BSMu(theta[:-2]) - - fr = u_to_fr((fr1, fr2, fr3), u) - - fit_values = { - 'convNorm': convNorm, - 'promptNorm': promptNorm, - 'muonNorm': muonNorm, - 'astroNorm': astroNorm, - 'astroDeltaGamma': astroDeltaGamma, - 'astroENorm': fr[0], - 'astroMuNorm': fr[1], - 'astroTauNorm': fr[2] - } - if FLAT: - return 10. - else: - return get_llh(GOLEMFIT, fit_values) - - -def lnprior(theta): - """Priors on theta.""" - convNorm, promptNorm, muonNorm, astroNorm, astroDeltaGamma = theta[:5] - theta = deepcopy(theta[5:]) - - # Nuisance parameter bounds - if 0. < convNorm < 5. and 0. < promptNorm < 5. and 0. < muonNorm < 5. and \ - 0. < astroNorm < 5. and 0. < astroDeltaGamma < 5.: - pass - else: return -np.inf - - if FIX_SFR: - if FIX_MIXING: - s12_2, sc2 = theta - elif FIX_SCALE: - s12_2, c13_4, s23_2, dcp = theta - else: - s12_2, c13_4, s23_2, dcp, sc2 = theta - else: - if FIX_MIXING: - sc2, sphi4, c2psi = theta - elif FIX_SCALE: - s12_2, c13_4, s23_2, dcp, sphi4, c2psi = theta - else: - s12_2, c13_4, s23_2, dcp, sc2, sphi4, c2psi = theta - if not FIX_SCALE: - sc2 = np.power(10., sc2) - - if not FIX_SFR: - # Flavour ratio bounds - if 0. <= sphi4 <= 1.0 and -1.0 <= c2psi <= 1.0: - pass - else: return -np.inf - - # Mixing angle bounds - if not FIX_MIXING: - if 0. <= s12_2 <= 1. and 0. <= c13_4 <= 1. and 0. <= s23_2 <= 1. \ - and 0. <= dcp <= 2*np.pi: - pass - else: return -np.inf - - # Scale bounds - if not FIX_SCALE: - if SCALE2_BOUNDS[0] <= sc2 <= SCALE2_BOUNDS[1]: - pass - else: return -np.inf - - return 0. - - -def lnprob(theta): - """Prob function for mcmc.""" - lp = lnprior(theta) - if not np.isfinite(lp): - return -np.inf - return lp + triangle_llh(theta) - - -def parse_args(): - """Parse command line arguments""" - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - '--measured-ratio', type=int, nargs=3, default=[1, 1, 1], - help='Set the central value for the measured flavour ratio at IceCube' - ) - parser.add_argument( - '--sigma-ratio', type=float, default=0.01, - help='Set the 1 sigma for the measured flavour ratio' - ) - parser.add_argument( - '--fix-source-ratio', type=str, default='False', - help='Fix the source flavour ratio' - ) - parser.add_argument( - '--source-ratio', type=int, nargs=3, default=[2, 1, 0], - help='Set the source flavour ratio for the case when you want to fix it' - ) - parser.add_argument( - '--no-bsm', type=str, default='False', - help='Turn off BSM terms' - ) - parser.add_argument( - '--fix-mixing', type=str, default='False', - help='Fix all mixing parameters except one' - ) - parser.add_argument( - '--fix-scale', type=str, default='False', - help='Fix the new physics scale' - ) - parser.add_argument( - '--scale', type=float, default=1e-30, - help='Set the new physics scale' - ) - parser.add_argument( - '--dimension', type=int, default=3, - help='Set the new physics dimension to consider' - ) - parser.add_argument( - '--energy', type=float, default=1000, - help='Set the energy scale' - ) - parser.add_argument( - '--flat-llh', type=str, default='False', - help='Run with a flat likelihood' - ) - parser.add_argument( - '--burnin', type=int, default=100, - help='Amount to burnin' - ) - parser.add_argument( - '--nwalkers', type=int, default=100, - help='Number of walkers' - ) - parser.add_argument( - '--nsteps', type=int, default=2000, - help='Number of steps to run' - ) - parser.add_argument( - '--seed', type=int, default=99, - help='Set the random seed value' - ) - parser.add_argument( - '--outfile', type=str, default='./untitled', - help='Path to output chains' - ) - args = parser.parse_args() - return args - - -def main(): - args = parse_args() - - np.random.seed(args.seed) - - global DIMENSION - global ENERGY - global MEASURED_FR - global SIGMA - global FLAT - global FIX_SFR - global SOURCE_FR - global FIX_MIXING - global FIX_SCALE - global SCALE - global SCALE2_BOUNDS - global NO_BSM - - DIMENSION = args.dimension - ENERGY = args.energy - - MEASURED_FR = np.array(args.measured_ratio) / float(np.sum(args.measured_ratio)) - SIGMA = args.sigma_ratio - - if args.flat_llh.lower() == 'true': - FLAT = True - elif args.flat_llh.lower() == 'false': - FLAT = False - else: - raise ValueError - - if args.fix_source_ratio.lower() == 'true': - FIX_SFR = True - elif args.fix_source_ratio.lower() == 'false': - FIX_SFR = False - else: - raise ValueError - - if args.fix_mixing.lower() == 'true': - FIX_MIXING = True - elif args.fix_mixing.lower() == 'false': - FIX_MIXING = False - else: - raise ValueError - - if args.fix_scale.lower() == 'true': - FIX_SCALE = True - elif args.fix_scale.lower() == 'false': - FIX_SCALE = False - else: - raise ValueError - - if args.no_bsm.lower() == 'true': - NO_BSM = True - elif args.no_bsm.lower() == 'false': - NO_BSM = False - else: - raise ValueError - - if FIX_SFR: - SOURCE_FR = np.array(args.source_ratio) / float(np.sum(args.source_ratio)) - - if FIX_SCALE: - SCALE = args.scale - else: - if CHANGE_MASS: - SCALE = 1E-27 - else: - SCALE = np.power(10, np.round(np.log10(MASS_EIGENVALUES[1]/ENERGY)) - np.log10(ENERGY**(DIMENSION-3))) - - if FIX_MIXING and FIX_SFR: - raise NotImplementedError('Fixed mixing and sfr not implemented') - if FIX_MIXING and FIX_SCALE: - raise NotImplementedError('Fixed mixing and scale not implemented') - - SCALE2_BOUNDS = (SCALE*1E-10, SCALE*1E10) - - print 'RUN_SCAN = {0}'.format(RUN_SCAN) - print 'MEASURED_FR = {0}'.format(MEASURED_FR) - print 'SIGMA = {0}'.format(SIGMA) - print 'FLAT = {0}'.format(FLAT) - print 'NO_BSM = {0}'.format(NO_BSM) - print 'ENERGY = {0}'.format(ENERGY) - print 'DIMENSION = {0}'.format(DIMENSION) - print 'SCALE2_BOUNDS = {0}'.format(SCALE2_BOUNDS) - print 'FIX_SFR = {0}'.format(FIX_SFR) - if FIX_SFR: - print 'SOURCE_FR = {0}'.format(SOURCE_FR) - print 'FIX_MIXING = {0}'.format(FIX_MIXING) - print 'FIX_SCALE = {0}'.format(FIX_SCALE) - if FIX_SCALE: - print 'SCALE = {0}'.format(SCALE) - - global GOLEMFIT - if RUN_SCAN: - GOLEMFIT = set_up_gf() - - # TODO(shivesh): make all these into args - inject = { - 'convNorm': 1., - 'promptNorm': 0., - 'muonNorm': 1., - 'astroNorm': 1., - 'astroDeltaGamma': 2., - 'astroENorm': MEASURED_FR[0], - 'astroMuNorm': MEASURED_FR[1], - 'astroTauNorm': MEASURED_FR[2], - } - - print "Injecting this model: ", inject - set_up_as(GOLEMFIT, inject) - - if FIX_SFR: - if FIX_MIXING: - ndim = 7 - elif FIX_SCALE: - ndim = 9 - else: - ndim = 10 - else: - if FIX_MIXING: - ndim = 8 - elif FIX_SCALE: - ndim = 11 - else: - ndim = 12 - nwalkers = args.nwalkers - ntemps = 1 - burnin = args.burnin - betas = np.array([1e0, 1e-1, 1e-2, 1e-3, 1e-4]) - s2 = np.average(np.log10(SCALE2_BOUNDS)) - if FIX_SFR: - if FIX_MIXING: - p0_base = [1., 0., 1., 1., 2., 0.5, s2] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 0.2, 3] - elif FIX_SCALE: - p0_base = [1., 0., 1., 1., 2., 0.5, 0.5, 0.5, np.pi] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2] - else: - p0_base = [1., 0., 1., 1., 2., 0.5, 0.5, 0.5, np.pi, s2] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2, 3] - else: - if FIX_MIXING: - p0_base = [1., 0., 1., 1., 2., s2, 0.5, 0.0] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 3, 0.2, 0.2] - elif FIX_SCALE: - p0_base = [1., 0., 1., 1., 2., 0.5, 0.5, 0.5, np.pi, 0.5, 0.0] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2] - else: - p0_base = [1., 0., 1., 1., 2., 0.5, 0.5, 0.5, np.pi, s2, 0.5, 0.0] - p0_std = [0.3, 0.05, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2, 3, 0.2, 0.2] - - print 'p0_base', p0_base - print 'p0_std', p0_std - p0 = np.random.normal(p0_base, p0_std, size=[ntemps, nwalkers, ndim]) - print map(lnprior, p0[0]) - - if RUN_SCAN: - # threads = multiprocessing.cpu_count() - threads = 1 - sampler = emcee.PTSampler( - ntemps, nwalkers, ndim, triangle_llh, lnprior, threads=threads - ) - - if RUN_SCAN: - print "Running burn-in" - for result in tqdm.tqdm(sampler.sample(p0, iterations=burnin), total=burnin): - pos, prob, state = result - sampler.reset() - print "Finished burn-in" - - nsteps = args.nsteps - - print "Running" - for _ in tqdm.tqdm(sampler.sample(pos, iterations=nsteps), total=nsteps): - pass - print "Finished" - - if FIX_SFR: - if FIX_MIXING: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fix_mixing'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), int(SIGMA*1000), - int(SOURCE_FR[0]*100), int(SOURCE_FR[1]*100), int(SOURCE_FR[2]*100), DIMENSION - ) - elif FIX_SCALE: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fixed_scale_{8}'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), int(SIGMA*1000), - int(SOURCE_FR[0]*100), int(SOURCE_FR[1]*100), int(SOURCE_FR[2]*100), DIMENSION, SCALE - ) - else: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_single_scale'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), int(SIGMA*1000), - int(SOURCE_FR[0]*100), int(SOURCE_FR[1]*100), int(SOURCE_FR[2]*100), DIMENSION - ) - else: - if FIX_MIXING: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fix_mixing'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), - int(SIGMA*1000), DIMENSION - ) - elif FIX_SCALE: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fixed_scale_{5}'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), - int(SIGMA*1000), DIMENSION, SCALE - ) - else: - outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}'.format( - int(MEASURED_FR[0]*100), int(MEASURED_FR[1]*100), int(MEASURED_FR[2]*100), - int(SIGMA*1000), DIMENSION - ) - - if RUN_SCAN: - samples = sampler.chain[0, :, :, :].reshape((-1, ndim)) - print 'acceptance fraction', sampler.acceptance_fraction - print 'sum of acceptance fraction', np.sum(sampler.acceptance_fraction) - print 'np.unique(samples[:,0]).shape', np.unique(samples[:,0]).shape - - try: - print 'autocorrelation', sampler.acor - except: - print 'WARNING : NEED TO RUN MORE SAMPLES FOR FILE ' + outfile - - if FLAT: - outfile += '_flat' - - if RUN_SCAN: - try: - os.makedirs(outfile[:-len(os.path.basename(outfile))]) - except OSError as exc: # Python >2.5 - if exc.errno == errno.EEXIST and os.path.isdir(outfile[:-len(os.path.basename(outfile))]): - pass - else: - raise - np.save(outfile+'.npy', samples) - - print "Making triangle plots" - chainer_plot.plot( - infile=outfile+'.npy', - angles=True, - outfile=outfile[:5]+outfile[5:].replace('data', 'plots')+'_angles.pgf', - measured_ratio=MEASURED_FR, - sigma_ratio=SIGMA, - fix_sfr=FIX_SFR, - fix_mixing=FIX_MIXING, - fix_scale=FIX_SCALE, - source_ratio=SOURCE_FR, - scale=SCALE, - dimension=DIMENSION, - energy=ENERGY, - scale_bounds=SCALE2_BOUNDS, - ) - # chainer_plot.plot( - # infile=outfile+'.npy', - # angles=False, - # outfile=outfile[:5]+outfile[5:].replace('data', 'plots')+'.pdf', - # measured_ratio=MEASURED_FR, - # sigma_ratio=SIGMA, - # fix_sfr=FIX_SFR, - # fix_mixing=FIX_MIXING, - # fix_scale=FIX_SCALE, - # source_ratio=SOURCE_FR, - # scale=SCALE, - # dimension=DIMENSION, - # energy=ENERGY, - # scale_bounds=SCALE2_BOUNDS, - # ) - print "DONE!" - - -main.__doc__ = __doc__ - - -if __name__ == '__main__': - main() - diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/utils/__init__.py diff --git a/utils/enums.py b/utils/enums.py new file mode 100644 index 0000000..31885de --- /dev/null +++ b/utils/enums.py @@ -0,0 +1,68 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Define Enums for the BSM flavour ratio analysis +""" + +from enum import Enum + + +class DataType(Enum): + REAL = 1 + FAKE = 2 + ASMIMOV = 3 + + +class FitCateg(Enum): + HESESPL = 1 + HESEDPL = 2 + ZPSPL = 3 + ZPDPL = 4 + NUNUBAR2 = 5 + + +class FitFlagCateg(Enum): + DEFAULT = 1 + XS = 2 + + +class FitPriorsCateg(Enum): + DEFAULT = 1 + XS = 2 + + +class Likelihood(Enum): + FLAT = 1 + GAUSSIAN = 2 + GOLEMFIT = 3 + + +class Priors(Enum): + UNIFORM = 1 + LOGUNIFORM = 2 + JEFFREYS = 3 + + +class SteeringCateg(Enum): + BASELINE = 1 + HOLEICE = 2 + ABSORPTION = 3 + SCATTERING = 4 + SCATTERING_ABSORPTION = 5 + STD = 6 + STD_HALF1 = 7 + STD_HALF2 = 8 + LONGLIFE = 9 + DPL = 10 + + +class XSCateg(Enum): + HALF = 1 + NOM = 2 + TWICE = 3 + TWICE01 = 4 + TWICE02 = 5 + diff --git a/utils/fr.py b/utils/fr.py new file mode 100644 index 0000000..ddcb5d2 --- /dev/null +++ b/utils/fr.py @@ -0,0 +1,353 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Useful functions for the BSM flavour ratio analysis +""" + +from __future__ import absolute_import, division + +import sys + +import numpy as np +from scipy import linalg + + +MASS_EIGENVALUES = [7.40E-23, 2.515E-21] +"""SM mass eigenvalues""" + + +def angles_to_fr(src_angles): + """Convert angular projection of the source flavour ratio back into the + flavour ratio. + + Parameters + ---------- + src_angles : list, length = 2 + sin(phi)^4 and cos(psi)^2 + + Returns + ---------- + flavour ratios (nue, numu, nutau) + + Examples + ---------- + >>> print angles_to_fr((0.3, 0.4)) + (0.38340579025361626, 0.16431676725154978, 0.45227744249483393) + + """ + sphi4, c2psi = src_angles + + psi = (0.5)*np.arccos(c2psi) + + sphi2 = np.sqrt(sphi4) + cphi2 = 1. - sphi2 + spsi2 = np.sin(psi)**2 + cspi2 = 1. - spsi2 + + x = sphi2*cspi2 + y = sphi2*spsi2 + z = cphi2 + return x, y, z + + +def angles_to_u(bsm_angles): + """Convert angular projection of the mixing matrix elements back into the + mixing matrix elements. + + Parameters + ---------- + bsm_angles : list, length = 4 + sin(12)^2, cos(13)^4, sin(23)^2 and deltacp + + Returns + ---------- + unitary numpy ndarray of shape (3, 3) + + Examples + ---------- + >>> from fr import angles_to_u + >>> print angles_to_u((0.2, 0.3, 0.5, 1.5)) + array([[ 0.66195018+0.j , 0.33097509+0.j , 0.04757188-0.6708311j ], + [-0.34631487-0.42427084j, 0.61741198-0.21213542j, 0.52331757+0.j ], + [ 0.28614067-0.42427084j, -0.64749908-0.21213542j, 0.52331757+0.j ]]) + + """ + s12_2, c13_4, s23_2, dcp = bsm_angles + dcp = np.complex128(dcp) + + c12_2 = 1. - s12_2 + c13_2 = np.sqrt(c13_4) + s13_2 = 1. - c13_2 + c23_2 = 1. - s23_2 + + t12 = np.arcsin(np.sqrt(s12_2)) + t13 = np.arccos(np.sqrt(c13_2)) + t23 = np.arcsin(np.sqrt(s23_2)) + + c12 = np.cos(t12) + s12 = np.sin(t12) + c13 = np.cos(t13) + s13 = np.sin(t13) + c23 = np.cos(t23) + s23 = np.sin(t23) + + p1 = np.array([[1 , 0 , 0] , [0 , c23 , s23] , [0 , -s23 , c23]] , dtype=np.complex128) + p2 = np.array([[c13 , 0 , s13*np.exp(-1j*dcp)] , [0 , 1 , 0] , [-s13*np.exp(1j*dcp) , 0 , c13]] , dtype=np.complex128) + p3 = np.array([[c12 , s12 , 0] , [-s12 , c12 , 0] , [0 , 0 , 1]] , dtype=np.complex128) + + u = np.dot(np.dot(p1, p2), p3) + return u + + +def cardano_eqn(ham): + """Diagonalise the effective Hamiltonian 3x3 matrix into the form + h_{eff} = UE_{eff}U^{dagger} using the procedure in PRD91, 052003 (2015). + + Parameters + ---------- + ham : numpy ndarray of shape (3, 3) + sin(12)^2, cos(13)^4, sin(23)^2 and deltacp + + Returns + ---------- + unitary numpy ndarray of shape (3, 3) + + Examples + ---------- + >>> import numpy as np + >>> from fr import cardano_eqn + >>> ham = np.array( + >>> [[ 0.66195018+0.j , 0.33097509+0.j , 0.04757188-0.6708311j ], + >>> [-0.34631487-0.42427084j, 0.61741198-0.21213542j, 0.52331757+0.j ], + >>> [ 0.28614067-0.42427084j, -0.64749908-0.21213542j, 0.52331757+0.j ]] + >>> ) + >>> print cardano_eqn(ham) + array([[-0.11143379-0.58863683j, -0.09067747-0.48219068j, 0.34276625-0.08686465j], + [ 0.14835519+0.47511473j, -0.18299305+0.40777481j, 0.31906300+0.82514223j], + [-0.62298966+0.07231745j, -0.61407815-0.42709603j, 0.03660313+0.30160428j]]) + + """ + if np.shape(ham) != (3, 3): + raise ValueError( + 'Input matrix should be a square and dimension 3, ' + 'got\n{0}'.format(ham) + ) + + a = -np.trace(ham) + b = (0.5) * ((np.trace(ham))**2 - np.trace(np.dot(ham, ham))) + c = -linalg.det(ham) + + Q = (1/9.) * (a**2 - 3*b) + R = (1/54.) * (2*a**3 - 9*a*b + 27*c) + theta = np.arccos(R / np.sqrt(Q**3)) + + E1 = -2 * np.sqrt(Q) * np.cos(theta/3.) - (1/3.)*a + E2 = -2 * np.sqrt(Q) * np.cos((theta - 2.*np.pi)/3.) - (1/3.)*a + E3 = -2 * np.sqrt(Q) * np.cos((theta + 2.*np.pi)/3.) - (1/3.)*a + + A1 = ham[1][2] * (ham[0][0] - E1) - ham[1][0]*ham[0][2] + A2 = ham[1][2] * (ham[0][0] - E2) - ham[1][0]*ham[0][2] + A3 = ham[1][2] * (ham[0][0] - E3) - ham[1][0]*ham[0][2] + + B1 = ham[2][0] * (ham[1][1] - E1) - ham[2][1]*ham[1][0] + B2 = ham[2][0] * (ham[1][1] - E2) - ham[2][1]*ham[1][0] + B3 = ham[2][0] * (ham[1][1] - E3) - ham[2][1]*ham[1][0] + + C1 = ham[1][0] * (ham[2][2] - E1) - ham[1][2]*ham[2][0] + C2 = ham[1][0] * (ham[2][2] - E2) - ham[1][2]*ham[2][0] + C3 = ham[1][0] * (ham[2][2] - E3) - ham[1][2]*ham[2][0] + + N1 = np.sqrt(abs(A1*B1)**2 + abs(A1*C1)**2 + abs(B1*C1)**2) + N2 = np.sqrt(abs(A2*B2)**2 + abs(A2*C2)**2 + abs(B2*C2)**2) + N3 = np.sqrt(abs(A3*B3)**2 + abs(A3*C3)**2 + abs(B3*C3)**2) + + mm = np.array([ + [np.conjugate(B1)*C1 / N1, np.conjugate(B2)*C2 / N2, np.conjugate(B3)*C3 / N3], + [A1*C1 / N1, A2*C2 / N2, A3*C3 / N3], + [A1*B1 / N1, A2*B2 / N2, A3*B3 / N3] + ]) + return mm + + +def normalise_fr(fr): + """Normalise an input flavour combination to a flavour ratio. + + Parameters + ---------- + fr : list, length = 3 + flavour combination + + Returns + ---------- + numpy ndarray flavour ratio + + Examples + ---------- + >>> from fr import normalise_fr + >>> print normalise_fr((1, 2, 3)) + array([ 0.16666667, 0.33333333, 0.5 ]) + + """ + return np.array(fr) / float(np.sum(fr)) + + +NUFIT_U = angles_to_u((0.307, (1-0.02195)**2, 0.565, 3.97935)) +"""NuFIT mixing matrix (s_12^2, c_13^4, s_23^2, dcp)""" + + +def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES, + nufit_u=NUFIT_U, no_bsm=False, fix_mixing=False, + fix_scale=False, scale=None, check_uni=True): + """Construct the BSM mixing matrix from the BSM parameters. + + Parameters + ---------- + theta : list, length > 3 + BSM parameters + + dim : int + Dimension of BSM physics + + energy : float + Energy in GeV + + mass_eigenvalues : list, length = 2 + SM mass eigenvalues + + nufit_u : numpy ndarray, dimension 3 + SM NuFIT mixing matrix + + no_bsm : bool + Turn off BSM behaviour + + fix_mixing : bool + Fix the BSM mixing angles + + fix_scale : bool + Fix the BSM scale + + scale : float + Used with fix_scale - scale at which to fix + + check_uni : bool + Check the resulting BSM mixing matrix is unitary + + Returns + ---------- + unitary numpy ndarray of shape (3, 3) + + Examples + ---------- + >>> from fr import params_to_BSMu + >>> print params_to_BSMu((0.2, 0.3, 0.5, 1.5, -20), dim=3, energy=1000) + array([[ 0.18658169 -6.34190523e-01j, -0.26460391 +2.01884200e-01j, 0.67247096 -9.86808417e-07j], + [-0.50419832 +2.14420570e-01j, -0.36013768 +5.44254868e-01j, 0.03700961 +5.22039894e-01j], + [-0.32561308 -3.95946524e-01j, 0.64294909 -2.23453580e-01j, 0.03700830 +5.22032403e-01j]]) + + """ + if np.shape(nufit_u) != (3, 3): + raise ValueError( + 'Input matrix should be a square and dimension 3, ' + 'got\n{0}'.format(ham) + ) + + if fix_mixing: + s12_2, c13_4, s23_2, dcp, sc2 = 0.5, 1.0-1E-6, 0.5, 0., theta + elif fix_scale: + s12_2, c13_4, s23_2, dcp = theta + sc2 = np.log10(scale) + else: + s12_2, c13_4, s23_2, dcp, sc2 = theta + sc2 = np.power(10., sc2) + sc1 = sc2 / 100. + + mass_matrix = np.array( + [[0, 0, 0], [0, mass_eigenvalues[0], 0], [0, 0, mass_eigenvalues[1]]] + ) + sm_ham = (1./(2*energy))*np.dot(nufit_u, np.dot(mass_matrix, nufit_u.conj().T)) + if no_bsm: + eg_vector = cardano_eqn(sm_ham) + else: + new_physics_u = angles_to_u((s12_2, c13_4, s23_2, dcp)) + scale_matrix = np.array( + [[0, 0, 0], [0, sc1, 0], [0, 0, sc2]] + ) + bsm_term = (energy**(dim-3)) * np.dot(new_physics_u, np.dot(scale_matrix, new_physics_u.conj().T)) + + bsm_ham = sm_ham + bsm_term + eg_vector = cardano_eqn(bsm_ham) + + if check_uni: + tu = test_unitarity(eg_vector) + if not abs(np.trace(tu) - 3.) < 1e-5 or \ + not abs(np.sum(tu) - 3.) < 1e-5: + raise AssertionError( + 'Matrix is not unitary!\neg_vector\n{0}\ntest ' + 'u\n{1}'.format(eg_vector, tu) + ) + return eg_vector + + +def test_unitarity(x, prnt=False): + """Test the unitarity of a matrix. + + Parameters + ---------- + x : numpy ndarray + Matrix to evaluate + + prnt : bool + Print the result + + Returns + ---------- + numpy ndarray + + Examples + ---------- + >>> from fr import test_unitarity + >>> x = np.identity(3) + >>> print test_unitarity(x) + array([[ 1., 0., 0.], + [ 0., 1., 0.], + [ 0., 0., 1.]]) + + """ + f = abs(np.dot(x, x.conj().T)) + if prnt: + print 'Unitarity test:\n{0}'.format(f) + return f + + +def u_to_fr(source_fr, matrix): + """Compute the observed flavour ratio assuming decoherence. + + Parameters + ---------- + source_fr : list, length = 3 + Source flavour ratio components + + matrix : numpy ndarray, dimension 3 + Mixing matrix + + Returns + ---------- + Measured flavour ratio + + Examples + ---------- + >>> from fr import params_to_BSMu, u_to_fr + >>> print u_to_fr((1, 2, 0), params_to_BSMu((0.2, 0.3, 0.5, 1.5, -20), 3, 1000)) + array([ 0.33740075, 0.33176584, 0.33083341]) + + """ + # TODO(shivesh): energy dependence + composition = np.einsum( + 'ai, bi, a -> b', abs(matrix)**2, abs(matrix)**2, source_fr + ) + ratio = composition / np.sum(source_fr) + return ratio + diff --git a/utils/gf.py b/utils/gf.py new file mode 100644 index 0000000..766c161 --- /dev/null +++ b/utils/gf.py @@ -0,0 +1,105 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Useful GolemFit wrappers for the BSM flavour ratio analysis +""" + +from __future__ import absolute_import, division + +import argparse +from functools import partial + +import GolemFitPy as gf + +from utils.enums import * +from utils.misc import enum_keys, enum_parse + + +def data_distributions(fitter): + hdat = fitter.GetDataDistribution() + binedges = np.asarray([fitter.GetZenithBinsData(), fitter.GetEnergyBinsData()]) + return hdat, binedges + + +def fit_flags(fitflag_categ): + flags = gf.FitParametersFlag() + if fitflag_categ is FitFlagCateg.xs: + # False means it's not fixed in minimization + flags.NeutrinoAntineutrinoRatio = False + return flags + + +def fit_params(fit_categ): + params = gf.FitParameters() + params.astroNorm = 7.5 + params.astroDeltaGamma = 0.9 + if fit_categ is FitCateg.hesespl: + params.astroNormSec = 0 + elif fit_categ is FitCateg.hesedpl: + params.astroNormSec=7.0 + elif fit_categ is FitCateg.zpspl: + # zero prompt, single powerlaw + params.promptNorm = 0 + params.astroNormSec = 0 + elif fit_categ is FitCateg.zpdpl: + # zero prompt, double powerlaw + params.promptNorm = 0 + params.astroNormSec=7.0 + elif fit_categ is FitCateg.nunubar2: + params.NeutrinoAntineutrinoRatio = 2 + + +def fit_priors(fitpriors_categ): + priors = gf.Priors() + if fitpriors_categ == FitPriorsCateg.xs: + priors.promptNormCenter = 1 + priors.promptNormWidth = 3 + priors.astroDeltaGammaCenter = 0 + priors.astroDeltaGammaWidth = 1 + return priors + + +def gen_steering_params(steering_categ, quiet=False): + params = gf.SteeringParams() + if quiet: params.quiet = True + params.fastmode = False + params.do_HESE_reshuffle = False + params.numc_tag = steering_categ.name + params.baseline_astro_spectral_index = -2. + if steering_categ is SteeringCateg.LONGLIFE: + params.years = [999] + params.numc_tag = 'std_half1' + if steering_categ is SteeringCateg.DPL: + params.diffuse_fit_type = gf.DiffuseFitType.DoublePowerLaw + params.numc_tag = 'std_half1' + return params + + +def gf_argparse(parser): + parser.add_argument( + '--data', default='real', type=partial(enum_parse, c=DataType), + choices=enum_keys(DataType), help='select datatype' + ) + parser.add_argument( + '--ast', default='baseline', type=partial(enum_parse, c=SteeringCateg), + choices=enum_keys(SteeringCateg), + help='use asimov/fake dataset with specific steering' + ) + parser.add_argument( + '--aft', default='hesespl', type=partial(enum_parse, c=FitCateg), + choices=enum_keys(FitCateg), + help='use asimov/fake dataset with specific Fit' + ) + parser.add_argument( + '--axs', default='nom', type=partial(enum_parse, c=XSCateg), + choices=enum_keys(XSCateg), + help='use asimov/fake dataset with xs scaling' + ) + parser.add_argument( + '--priors', default='uniform', type=partial(enum_parse, c=Priors), + choices=enum_keys(Priors), help='Bayesian priors' + ) + diff --git a/utils/mcmc.py b/utils/mcmc.py new file mode 100644 index 0000000..d91764f --- /dev/null +++ b/utils/mcmc.py @@ -0,0 +1,167 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Useful functions to use an MCMC for the BSM flavour ratio analysis +""" + +from __future__ import absolute_import, division + +import os +import errno + +import argparse +from functools import partial + +import emcee +import tqdm + +import numpy as np +from scipy.stats import multivariate_normal + +from utils import fr as fr_utils +from utils.enums import Likelihood +from utils.misc import enum_keys, enum_parse, parse_bool + + +def lnprior(theta, paramset): + """Priors on theta.""" + ranges = paramset.ranges + for value, range in zip(theta, ranges): + if range[0] <= value <= range[1]: + pass + else: return -np.inf + return 0. + + +def mcmc(p0, triangle_llh, lnprior, ndim, nwalkers, burnin, nsteps, ntemps=1, threads=1): + """Run the MCMC.""" + sampler = emcee.PTSampler( + ntemps, nwalkers, ndim, triangle_llh, lnprior, threads=threads + ) + + print "Running burn-in" + for result in tqdm.tqdm(sampler.sample(p0, iterations=burnin), total=burnin): + pos, prob, state = result + sampler.reset() + print "Finished burn-in" + + print "Running" + for _ in tqdm.tqdm(sampler.sample(pos, iterations=nsteps), total=nsteps): + pass + print "Finished" + + samples = sampler.chain[0, :, :, :].reshape((-1, ndim)) + print 'acceptance fraction', sampler.acceptance_fraction + print 'sum of acceptance fraction', np.sum(sampler.acceptance_fraction) + print 'np.unique(samples[:,0]).shape', np.unique(samples[:,0]).shape + try: + print 'autocorrelation', sampler.acor + except: + print 'WARNING : NEED TO RUN MORE SAMPLES' + + return samples + + +def mcmc_argparse(parser): + parser.add_argument( + '--likelihood', default='gaussian', type=partial(enum_parse, c=Likelihood), + choices=enum_keys(Likelihood), help='likelihood contour' + ) + parser.add_argument( + '--run-mcmc', type=parse_bool, default='True', + help='Run the MCMC' + ) + parser.add_argument( + '--burnin', type=int, default=100, + help='Amount to burnin' + ) + parser.add_argument( + '--nwalkers', type=int, default=100, + help='Number of walkers' + ) + parser.add_argument( + '--nsteps', type=int, default=2000, + help='Number of steps to run' + ) + + +def gaussian_llh(fr, fr_bf, sigma): + """Multivariate gaussian likelihood.""" + cov_fr = np.identity(3) * sigma + return np.log(multivariate_normal.pdf(fr, mean=fr_bf, cov=cov_fr)) + + +def gaussian_seed(paramset, ntemps, nwalkers): + """Get gaussian seed values for the MCMC.""" + ndim = len(paramset) + p0 = np.random.normal( + paramset.values, paramset.stds, size=[ntemps, nwalkers, ndim] + ) + return p0 + + +def save_chains(chains, outfile): + """Save the chains. + + Parameters + ---------- + chains : numpy ndarray + MCMC chains to save + + outfile : str + Output file location of chains + + """ + try: + os.makedirs(outfile[:-len(os.path.basename(outfile))]) + except OSError as exc: # Python >2.5 + if exc.errno == errno.EEXIST and os.path.isdir(outfile[:-len(os.path.basename(outfile))]): + pass + else: + raise + print 'Saving chains to location {0}'.format(outfile+'.npy') + np.save(outfile+'.npy', chains) + + +def triangle_llh(theta, args): + """-Log likelihood function for a given theta.""" + if args.fix_source_ratio: + fr1, fr2, fr3 = args.source_ratio + bsm_angles = theta[-5:] + else: + fr1, fr2, fr3 = fr_utils.angles_to_fr(theta[-2:]) + bsm_angles = theta[-7:-2] + + u = fr_utils.params_to_BSMu( + theta = bsm_angles, + dim = args.dimension, + energy = args.energy, + no_bsm = args.no_bsm, + fix_mixing = args.fix_mixing, + fix_scale = args.fix_scale, + scale = args.scale + ) + + fr = fr_utils.u_to_fr((fr1, fr2, fr3), u) + fr_bf = args.measured_ratio + if args.likelihood is Likelihood.FLAT: + return 1. + elif args.likelihood is Likelihood.GAUSSIAN: + return gaussian_llh(fr, fr_bf, args.sigma_ratio) + elif args.likelihood is Likelihood.GOLEMFIT: + raise NotImplementedError('TODO') + import GolemFitPy as gf + from collections import namedtuple + datapaths = gf.DataPaths() + IceModels = namedtuple('IceModels', 'std_half2') + fitters = IceModels(*[ + gf.GolemFit(datapaths, + gf.gen_steering_params(SteeringCateg.__members__[ice]), + xs_params(XSCateg.nom)) for ice in IceModels._fields]) + for fitter in fitters: + fitter.SetFitParametersFlag(fit_flags(FitFlagCateg.xs)) + fitter.SetFitPriors(fit_priors(FitPriorsCateg.xs)) + diff --git a/utils/misc.py b/utils/misc.py new file mode 100644 index 0000000..5c3eb2e --- /dev/null +++ b/utils/misc.py @@ -0,0 +1,230 @@ +# author : S. Mandalia +# s.p.mandalia@qmul.ac.uk +# +# date : March 17, 2018 + +""" +Misc functions for the BSM flavour ratio analysis +""" + +from __future__ import absolute_import, division + +import os +import errno +from collections import Sequence +import multiprocessing + +import numpy as np + +from utils.enums import Likelihood + + +class Param(object): + """Parameter class to store parameters. + """ + def __init__(self, name, value, ranges, std=None, tex=None): + self._ranges = None + self._tex = None + + self.name = name + self.value = value + self.ranges = ranges + self.std = std + self.tex = tex + + @property + def ranges(self): + return tuple(self._ranges) + + @ranges.setter + def ranges(self, values): + self._ranges = [val for val in values] + + @property + def tex(self): + return r'{0}'.format(self.tex) + + @tex.setter + def tex(self, t): + self._tex = t if t is not None else r'{\rm %s}' % self.name + + +class ParamSet(Sequence): + """Container class for a set of parameters. + """ + def __init__(self, *args): + param_sequence = [] + for arg in args: + try: + param_sequence.extend(arg) + except TypeError: + param_sequence.append(arg) + + # Disallow duplicated params + all_names = [p.name for p in param_sequence] + unique_names = set(all_names) + if len(unique_names) != len(all_names): + duplicates = set([x for x in all_names if all_names.count(x) > 1]) + raise ValueError('Duplicate definitions found for param(s): ' + + ', '.join(str(e) for e in duplicates)) + + # Elements of list must be Param type + assert all([isinstance(x, Param) for x in param_sequence]), \ + 'All params must be of type "Param"' + + self._params = param_sequence + + def __len__(self): + return len(self._params) + + def __getitem__(self, i): + if isinstance(i, int): + return self._params[i] + elif isinstance(i, basestring): + return self._by_name[i] + + def __getattr__(self, attr): + try: + return super(ParamSet, self).__getattribute__(attr) + except AttributeError: + t, v, tb = sys.exc_info() + try: + return self[attr] + except KeyError: + raise t, v, tb + + def __iter__(self): + return iter(self._params) + + @property + def _by_name(self): + return {obj.name: obj for obj in self._params} + + @property + def names(self): + return tuple([obj.name for obj in self._params]) + + @property + def values(self): + return tuple([obj.value for obj in self._params]) + + @property + def ranges(self): + return tuple([obj.ranges for obj in self._params]) + + @property + def stds(self): + return tuple([obj.std for obj in self._params]) + + @property + def params(self): + return self._params + + def to_dict(self): + return {obj.name: obj.value for obj in self._params} + + +def gen_outfile_name(args): + """Generate a name for the output file based on the input args. + + Parameters + ---------- + args : argparse + argparse object to print + + """ + mr = args.measured_ratio + si = args.sigma_ratio + if args.fix_source_ratio: + sr = args.source_ratio + if args.fix_mixing: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fix_mixing'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000), + int(sr[0]*100), int(sr[1]*100), int(sr[2]*100), args.dimension + ) + elif args.fix_scale: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fixed_scale_{8}'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000), + int(sr[0]*100), int(sr[1]*100), int(sr[2]*100), args.dimension, + args.scale + ) + else: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_single_scale'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000), + int(sr[0]*100), int(sr[1]*100), int(sr[2]*100), args.dimension + ) + else: + if args.fix_mixing: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fix_mixing'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), + int(si*1000), args.dimension + ) + elif args.fix_scale: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fixed_scale_{5}'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), + int(si*1000), args.dimension, args.scale + ) + else: + outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}'.format( + int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), + int(si*1000), args.dimension + ) + if args.likelihood is Likelihood.FLAT: outfile += '_flat' + return outfile + + +def parse_bool(s): + """Parse a string to a boolean. + + Parameters + ---------- + s : str + String to parse + + Returns + ---------- + bool + + Examples + ---------- + >>> from misc import parse_bool + >>> print parse_bool('true') + True + + """ + if s.lower() == 'true': + return True + elif s.lower() == 'false': + return False + else: + raise ValueError + + +def print_args(args): + """Print the input arguments. + + Parameters + ---------- + args : argparse + argparse object to print + + """ + arg_vars = vars(args) + for key in arg_vars.iterkeys(): + print '== {0:<25} = {1}'.format(key, arg_vars[key]) + + +def enum_keys(e): + return e.__members__.keys() + + +def enum_parse(s, c): + return c[s.upper()] + + +def thread_type(t): + if t.lower() == 'max': + return multiprocessing.cpu_count() + else: + return int(t) + |