refactor fr script

1 files changed, 112 insertions, 91 deletions

diff --git a/fr.py b/fr.py
index 5bc4c82..a6d3b1a 100755
--- a/fr.py
+++ b/fr.py
@@ -12,19 +12,18 @@ from __future__ import absolute_import, division
 
 import os
 import argparse
+from copy import deepcopy
 from functools import partial
 
 import numpy as np
 
 from utils import fr as fr_utils
-from utils import gf as gf_utils
 from utils import llh as llh_utils
 from utils import mcmc as mcmc_utils
 from utils import misc as misc_utils
 from utils import plot as plot_utils
-from utils.enums import EnergyDependance, Likelihood, MixingScenario
 from utils.enums import MCMCSeedType, ParamTag, PriorsCateg
-from utils.param import Param, ParamSet, get_paramsets
+from utils.param import Param, ParamSet
 
 
 def define_nuisance():
@@ -32,45 +31,45 @@ def define_nuisance():
     tag = ParamTag.SM_ANGLES
     nuisance = []
     g_prior = PriorsCateg.GAUSSIAN
-    hg_prior = PriorsCateg.LIMITEDGAUSS
+    lg_prior = PriorsCateg.LIMITEDGAUSS
     e = 1e-9
     nuisance.extend([
-        Param(name='s_12_2', value=0.307,            seed=[0.26, 0.35],     ranges=[0., 1.],      std=0.013,   tex=r's_{12}^2', prior=g_prior,  tag=tag),
-        Param(name='c_13_4', value=(1-(0.02206))**2, seed=[0.950, 0.961],   ranges=[0., 1.],      std=0.00147, tex=r'c_{13}^4', prior=g_prior, tag=tag),
-        Param(name='s_23_2', value=0.538,            seed=[0.31, 0.75],     ranges=[0., 1.],      std=0.069,   tex=r's_{23}^2', prior=g_prior,  tag=tag),
+        Param(name='s_12_2', value=0.307,            seed=[0.26, 0.35],     ranges=[0., 1.],      std=0.013,   tex=r's_{12}^2',    tag=tag),
+        Param(name='c_13_4', value=(1-(0.02206))**2, seed=[0.950, 0.961],   ranges=[0., 1.],      std=0.00147, tex=r'c_{13}^4',    tag=tag),
+        Param(name='s_23_2', value=0.538,            seed=[0.31, 0.75],     ranges=[0., 1.],      std=0.069,   tex=r's_{23}^2',    tag=tag),
         Param(name='dcp',    value=4.08404,          seed=[0+e, 2*np.pi-e], ranges=[0., 2*np.pi], std=2.0,     tex=r'\delta_{CP}', tag=tag),
-        Param(
-            name='m21_2', value=7.40E-23, seed=[7.2E-23, 7.6E-23], ranges=[6.80E-23, 8.02E-23],
-            std=2.1E-24, tex=r'\Delta m_{21}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag
-        ),
-        Param(
-            name='m3x_2', value=2.494E-21, seed=[2.46E-21, 2.53E-21], ranges=[2.399E-21, 2.593E-21],
-            std=3.3E-23, tex=r'\Delta m_{3x}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag
-        )
-        # Param(name='s_12_2', value=0.307,            ranges=[0., 1.],      std=20,   tex=r's_{12}^2', prior=g_prior,  tag=tag),
-        # Param(name='c_13_4', value=(1-(0.02206))**2, ranges=[0., 1.],      std=20, tex=r'c_{13}^4', prior=g_prior, tag=tag),
-        # Param(name='s_23_2', value=0.538,            ranges=[0., 1.],      std=20,  tex=r's_{23}^2', prior=g_prior,  tag=tag),
-        # Param(name='dcp',    value=4.08404,          ranges=[0., 2*np.pi], std=20,     tex=r'\delta_{CP}', tag=tag),
-        # Param(
-        #     name='m21_2', value=7.40E-23, ranges=[6.80E-23, 8.02E-23],
-        #     std=2.1E-24, tex=r'\Delta m_{21}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag
-        # ),
-        # Param(
-        #     name='m3x_2', value=2.494E-21, ranges=[2.399E-21, 2.593E-21],
-        #     std=3.3E-23, tex=r'\Delta m_{3x}^2{\rm GeV}^{-2}', prior=g_prior, tag=tag
-        # )
-    ])
-    tag = ParamTag.NUISANCE
-    nuisance.extend([
-        Param(name='convNorm',        value=1.,  seed=[0.5, 2. ], ranges=[0.1, 10.], std=0.4, prior=lg_prior, tag=tag),
-        Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 20.], std=2.4, prior=lg_prior, tag=tag),
-        Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0. , 10.], std=0.1, tag=tag),
-        Param(name='astroNorm',       value=6.9, seed=[0.,  5. ], ranges=[0. , 20.], std=1.5, tag=tag),
-        Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[-5., 5. ], std=0.1, tag=tag)
     ])
     return ParamSet(nuisance)
 
 
+def get_paramsets(args, nuisance_paramset):
+    """Make the paramsets for generating the Asmimov MC sample and also running
+    the MCMC.
+    """
+    asimov_paramset = []
+    hypo_paramset = []
+
+    hypo_paramset.extend(
+        [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)]
+    )
+
+    for parm in hypo_paramset:
+        parm.value = args.__getattribute__(parm.name)
+
+    hypo_paramset = ParamSet(hypo_paramset)
+
+    tag = ParamTag.BESTFIT
+    flavour_angles = fr_utils.fr_to_angles(args.injected_ratio)
+
+    asimov_paramset.extend([
+        Param(name='astroFlavorAngle1', value=flavour_angles[0], ranges=[ 0., 1.], std=0.2, tag=tag),
+        Param(name='astroFlavorAngle2', value=flavour_angles[1], ranges=[-1., 1.], std=0.2, tag=tag),
+    ])
+    asimov_paramset = ParamSet(asimov_paramset)
+
+    return asimov_paramset, hypo_paramset
+
+
 def nuisance_argparse(parser):
     nuisance = define_nuisance()
     for parm in nuisance:
@@ -82,24 +81,7 @@ def nuisance_argparse(parser):
 
 def process_args(args):
     """Process the input args."""
-    # if args.fix_mixing is not MixingScenario.NONE and args.fix_scale:
-    #     raise NotImplementedError('Fixed mixing and scale not implemented')
-    if args.fix_mixing is not MixingScenario.NONE and args.fix_mixing_almost:
-        raise NotImplementedError(
-            '--fix-mixing and --fix-mixing-almost cannot be used together'
-        )
-
-    args.measured_ratio = fr_utils.normalise_fr(args.measured_ratio)
-    if args.fix_source_ratio:
-        args.source_ratio = fr_utils.normalise_fr(args.source_ratio)
-
-    if args.energy_dependance is EnergyDependance.SPECTRAL:
-        args.binning = np.logspace(
-            np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+1
-        )
-
-    if not args.fix_scale:
-        args.scale, args.scale_region = fr_utils.estimate_scale(args)
+    args.injected_ratio = fr_utils.normalise_fr(args.injected_ratio)
 
 
 def parse_args(args=None):
@@ -109,7 +91,11 @@ def parse_args(args=None):
         formatter_class=misc_utils.SortingHelpFormatter,
     )
     parser.add_argument(
-        '--seed', type=misc_utils.seed_parse, default='25',
+        '--injected-ratio', type=float, nargs=3, default=[1, 2, 0],
+        help='Set the central value for the injected flavour ratio at source'
+    )
+    parser.add_argument(
+        '--seed', type=misc_utils.seed_parse, default='26',
         help='Set the random seed value'
     )
     parser.add_argument(
@@ -117,20 +103,57 @@ def parse_args(args=None):
         help='Set the number of threads to use (int or "max")'
     )
     parser.add_argument(
-        '--outfile', type=str, default='./untitled',
-        help='Path to output results'
+        '--datadir', type=str, default='./untitled',
+        help='Path to store chains'
     )
-    fr_utils.fr_argparse(parser)
-    try:
-        gf_utils.gf_argparse(parser)
-    except: pass
-    llh_utils.likelihood_argparse(parser)
     mcmc_utils.mcmc_argparse(parser)
     nuisance_argparse(parser)
     if args is None: return parser.parse_args()
     else: return parser.parse_args(args.split())
 
 
+def gen_identifier(args):
+    f = '_INJ_{0}'.format(misc_utils.solve_ratio(args.injected_ratio))
+    return f
+
+
+def gen_figtext(args, asimov_paramset):
+    f = ''
+    f += 'Injected ratio = {0}'.format(
+        misc_utils.solve_ratio(args.injected_ratio)
+    )
+    for param in asimov_paramset:
+        f += '\nInjected {0:20s} = {1:.3f}'.format(
+            param.name, param.nominal_value
+        )
+    return f
+
+
+def triangle_llh(theta, args, hypo_paramset):
+    """Log likelihood function for a given theta."""
+    if len(theta) != len(hypo_paramset):
+        raise AssertionError(
+            'Dimensions of scan is not the same as the input '
+            'params\ntheta={0}\nparamset]{1}'.format(theta, hypo_paramset)
+        )
+    for idx, param in enumerate(hypo_paramset):
+        param.value = theta[idx]
+
+    return 1. # Flat LLH
+
+
+def ln_prob(theta, args, hypo_paramset):
+    dc_hypo_paramset = deepcopy(hypo_paramset)
+    lp = llh_utils.lnprior(theta, paramset=dc_hypo_paramset)
+    if not np.isfinite(lp):
+        return -np.inf
+    return lp + triangle_llh(
+        theta,
+        args           = args,
+        hypo_paramset  = dc_hypo_paramset,
+    )
+
+
 def main():
     args = parse_args()
     process_args(args)
@@ -139,60 +162,58 @@ def main():
     if args.seed is not None:
         np.random.seed(args.seed)
 
-    asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance())
-    outfile = misc_utils.gen_outfile_name(args)
-    print '== {0:<25} = {1}'.format('outfile', outfile)
+    asimov_paramset, hypo_paramset = get_paramsets(args, define_nuisance())
 
-    if args.run_mcmc:
-        if args.likelihood is Likelihood.GOLEMFIT:
-            fitter = gf_utils.setup_fitter(args, asimov_paramset)
-        else: fitter = None
+    prefix = ''
+    outfile = args.datadir + '/fr' + prefix + gen_identifier(args)
+    print '== {0:<25} = {1}'.format('outfile', outfile)
 
-        print 'asimov_paramset', asimov_paramset
-        print 'llh_paramset', llh_paramset
+    print 'asimov_paramset', asimov_paramset
+    print 'hypo_paramset', hypo_paramset
 
-        ln_prob = partial(
-            llh_utils.ln_prob, args=args, fitter=fitter,
-            asimov_paramset=asimov_paramset, llh_paramset=llh_paramset
+    if args.run_mcmc:
+        ln_prob_eval = partial(
+            ln_prob,
+            hypo_paramset  = hypo_paramset,
+            args           = args,
         )
 
-        ndim = len(llh_paramset)
         if args.mcmc_seed_type == MCMCSeedType.UNIFORM:
             p0 = mcmc_utils.flat_seed(
-                llh_paramset, nwalkers=args.nwalkers
+                hypo_paramset, nwalkers=args.nwalkers
             )
         elif args.mcmc_seed_type == MCMCSeedType.GAUSSIAN:
             p0 = mcmc_utils.gaussian_seed(
-                llh_paramset, nwalkers=args.nwalkers
+                hypo_paramset, nwalkers=args.nwalkers
             )
 
-        if args.save_measured_fr:
-            n = misc_utils.gen_identifier(args) + '.txt'
-            f = args.output_measured_fr + n
-            if os.path.isfile(f):
-                open(f, 'w').close()
-
         samples = mcmc_utils.mcmc(
             p0       = p0,
-            ln_prob  = ln_prob,
-            ndim     = ndim,
+            ln_prob  = ln_prob_eval,
+            ndim     = len(hypo_paramset),
             nwalkers = args.nwalkers,
             burnin   = args.burnin,
             nsteps   = args.nsteps,
             args     = args,
-            threads  = 1
-            # TODO(shivesh): broken because you cannot pickle a GolemFitPy object
-            # threads      = misc_utils.thread_factors(args.threads)[0]
+            threads  = args.threads
         )
-        mcmc_utils.save_chains(samples, outfile)
 
+        mmxs = map(fr_utils.angles_to_u, samples)
+        frs = np.array(
+            [fr_utils.u_to_fr(args.injected_ratio, x) for x in mmxs]
+        )
+        mcmc_utils.save_chains(frs, outfile)
+
+    of = outfile[:5]+outfile[5:].replace('data', 'plots')+'_posterior'
     plot_utils.chainer_plot(
         infile       = outfile+'.npy',
-        outfile      = outfile[:5]+outfile[5:].replace('data', 'plots'),
-        outformat    = ['pdf'],
+        outfile      = of,
+        outformat    = ['png'],
         args         = args,
-        llh_paramset = llh_paramset
+        llh_paramset = hypo_paramset,
+        fig_text     = gen_figtext(args, hypo_paramset)
     )
+
     print "DONE!"