plotting

1 files changed, 81 insertions, 1 deletions

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__