From eb98a187b1bc4ee8e8c962d0e5d1a3d0ae77bc35 Mon Sep 17 00:00:00 2001 From: shivesh Date: Tue, 6 Nov 2018 17:49:17 -0600 Subject: plotting --- plot_llh/sample.py | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 81 insertions(+), 1 deletion(-) (limited to 'plot_llh/sample.py') diff --git a/plot_llh/sample.py b/plot_llh/sample.py index 469a538..0a896e7 100644 --- a/plot_llh/sample.py +++ b/plot_llh/sample.py @@ -12,10 +12,64 @@ import argparse from functools import partial import numpy as np +from scipy.stats import multivariate_normal from utils import fr as fr_utils +from utils import likelihood as llh_utils +from utils import mcmc as mcmc_utils from utils import misc as misc_utils from utils.param import Param, ParamSet, get_paramsets +from utils.enums import MixingScenario, MCMCSeedType + + +def triangle_llh(theta, args): + """-Log likelihood function for a given theta.""" + fr = fr_utils.angles_to_fr(theta) + fr_bf = args.measured_ratio + cov_fr = np.identity(3) * args.sigma_ratio + return np.log(multivariate_normal.pdf(fr, mean=fr_bf, cov=cov_fr)) + + +def ln_prior(theta): + """Priors on theta.""" + sphi4, c2psi = theta + # Flavour ratio bounds + if 0. <= sphi4 <= 1.0 and -1.0 <= c2psi <= 1.0: + pass + else: return -np.inf + return 0. + + +def ln_prob(theta, args): + """Prob function for mcmc.""" + lp = ln_prior(theta) + if not np.isfinite(lp): + return -np.inf + return lp + triangle_llh(theta, args) + + +def process_args(args): + """Process the input args.""" + if args.fix_mixing is MixingScenario.NONE: + raise AssertionError('Must set a mixing scenario using --fix-mixing') + if not args.fix_source_ratio: + raise AssertionError('Must set source ratio using --fix-source-ratio') + + args.source_ratio = fr_utils.normalise_fr(args.source_ratio) + if args.fix_mixing is MixingScenario.T12: + s12_2, c13_4, s23_2, dcp = 0.5, 1.0, 0.0, 0. + elif args.fix_mixing is MixingScenario.T13: + s12_2, c13_4, s23_2, dcp = 0.0, 0.25, 0.0, 0. + elif args.fix_mixing is MixingScenario.T23: + s12_2, c13_4, s23_2, dcp = 0.0, 1.0, 0.5, 0. + + mm = np.array( + fr_utils.angles_to_u((s12_2, c13_4, s23_2, dcp)), dtype=np.complex256 + ) + args.measured_ratio = fr_utils.u_to_fr(args.source_ratio, mm) + + if not args.fix_scale: + args.scale, args.scale_region = fr_utils.estimate_scale(args) def parse_args(args=None): @@ -41,6 +95,8 @@ def parse_args(args=None): help='Plot MultiNest evidence or LLH value' ) fr_utils.fr_argparse(parser) + llh_utils.likelihood_argparse(parser) + mcmc_utils.mcmc_argparse(parser) if args is None: return parser.parse_args() else: return parser.parse_args(args.split()) @@ -53,10 +109,34 @@ def main(): if args.seed is not None: np.random.seed(args.seed) - asimov_paramset, llh_paramset = get_paramsets(args, ParamSet()) + paramset = ParamSet([ + Param(name='sphi4', value=0.5, ranges=[0.0, 1.0], std=0.2), + Param(name='c2psi', value=0.0, ranges=[-1.0, 1.0], std=0.2) + ]) + outfile = misc_utils.gen_outfile_name(args) print '== {0:<25} = {1}'.format('outfile', outfile) + if args.run_mcmc: + ndim = len(paramset) + print paramset + + if args.mcmc_seed_type == MCMCSeedType.UNIFORM: + p0 = mcmc_utils.flat_seed(paramset, nwalkers=args.nwalkers) + + samples = mcmc_utils.mcmc( + p0 = p0, + ln_prob = partial(ln_prob, args=args), + ndim = ndim, + nwalkers = args.nwalkers, + burnin = args.burnin, + nsteps = args.nsteps, + threads = args.threads + ) + fr_chains = np.array(map(fr_utils.angles_to_fr, samples), dtype=float) + mcmc_utils.save_chains(fr_chains, outfile) + print "DONE!" + main.__doc__ = __doc__ -- cgit v1.2.3