refactor sensitivity/limit scripts

1 files changed, 130 insertions, 290 deletions

diff --git a/sens.py b/sens.py
index 7049b83..22d464a 100755
--- a/sens.py
+++ b/sens.py
@@ -22,12 +22,10 @@ from utils import fr as fr_utils
 from utils import gf as gf_utils
 from utils import llh as llh_utils
 from utils import misc as misc_utils
-from utils import plot as plot_utils
-from utils.enums import EnergyDependance, Likelihood, MixingScenario, ParamTag
-from utils.enums import PriorsCateg, SensitivityCateg, StatCateg
-from utils.param import Param, ParamSet, get_paramsets
-
 from utils import mn as mn_utils
+from utils.enums import Likelihood, ParamTag
+from utils.enums import PriorsCateg, SensitivityCateg, StatCateg, Texture
+from utils.param import Param, ParamSet
 
 
 def define_nuisance():
@@ -58,20 +56,62 @@ def define_nuisance():
         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)
-        # Param(name='convNorm',        value=1.,  seed=[0.5, 2. ], ranges=[0. , 50.], std=0.3,  tag=tag),
-        # Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 50.], std=0.05, tag=tag),
-        # Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0. , 50.], std=0.1,  tag=tag),
-        # Param(name='astroNorm',       value=6.9, seed=[0.1, 10.], ranges=[0. , 50.], std=0.1,  tag=tag),
-        # Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[-5., 5. ], std=0.1,  tag=tag)
-        # Param(name='convNorm',        value=1.,  seed=[0.5, 2. ], ranges=[0.5, 2. ], std=0.3,  tag=tag),
-        # Param(name='promptNorm',      value=0.,  seed=[0. , 6. ], ranges=[0. , 6. ], std=0.05, tag=tag),
-        # Param(name='muonNorm',        value=1.,  seed=[0.1, 2. ], ranges=[0.1, 2. ], std=0.1,  tag=tag),
-        # Param(name='astroNorm',       value=6.9, seed=[0.1, 10.], ranges=[0.1, 10.], std=0.1,  tag=tag),
-        # Param(name='astroDeltaGamma', value=2.5, seed=[2.4, 3. ], ranges=[2.4, 3. ], 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 = []
+    llh_paramset = []
+
+    gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)]
+
+    llh_paramset.extend(
+        [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)]
+    )
+    llh_paramset.extend(gf_nuisance)
+
+    for parm in llh_paramset:
+        parm.value = args.__getattribute__(parm.name)
+
+    tag = ParamTag.MMANGLES
+    llh_paramset.extend([
+        Param(name='np_s_12^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{12}^2', tag=tag),
+        Param(name='np_c_13^4', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{c}_{13}^4', tag=tag),
+        Param(name='np_s_23^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{23}^2', tag=tag),
+        Param(name='np_dcp',  value=np.pi, ranges=[0., 2*np.pi], std=0.2, tex=r'\tilde{\delta_{CP}}', tag=tag)
+    ])
+
+    boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension]
+    tag = ParamTag.SCALE
+    llh_paramset.append(
+        Param(
+            name='logLam', value=np.mean(boundaries), ranges=boundaries, std=3,
+            tex=r'{\rm log}_{10}\left (\Lambda^{-1}'+get_units(args.dimension)+r'\right )',
+            tag=tag
+        )
+    )
+    llh_paramset = ParamSet(llh_paramset)
+
+    tag = ParamTag.BESTFIT
+    if args.data in [DataType.ASIMOV, DataType.REALISATION]:
+        flavour_angles = fr_to_angles(args.injected_ratio)
+    else:
+        flavour_angles = fr_to_angles([1, 1, 1])
+
+    asimov_paramset.extend(gf_nuisance)
+    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, llh_paramset
+
+
 def nuisance_argparse(parser):
     nuisance = define_nuisance()
     for parm in nuisance:
@@ -83,43 +123,25 @@ 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'
-        )
-    if args.sens_run and args.fix_scale:
-        raise NotImplementedError(
-            '--sens-run and --fix-scale 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)
+    args.source_ratio = fr_utils.normalise_fr(args.source_ratio)
+    if args.data is not DataType.REAL:
+        args.injected_ratio = fr_utils.normalise_fr(args.injected_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.binning = np.logspace(
+        np.log10(args.binning[0]), np.log10(args.binning[1]), args.binning[2]+1
+    )
 
-    if args.sens_eval_bin.lower() == 'all':
-        args.sens_eval_bin = None
+    if args.eval_segment.lower() == 'all':
+        args.eval_segment = None
     else:
-        args.sens_eval_bin = int(args.sens_eval_bin)
-
-    if args.sens_eval_bin is not None and args.plot_statistic:
-        print 'Cannot make plot when specific scale bin is chosen'
-        args.plot_statistic = False
+        args.eval_segment = int(args.eval_segment)
 
     if args.stat_method is StatCateg.FREQUENTIST and \
        args.likelihood is Likelihood.GOLEMFIT:
         args.likelihood = Likelihood.GF_FREQ
 
-    args.burnin = False
+    if args.texture is Texture.NONE:
+        raise ValueError('Must assume a BSM texture')
 
 
 def parse_args(args=None):
@@ -137,39 +159,20 @@ 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 chains'
-    )
-    parser.add_argument(
-        '--sens-run', type=misc_utils.parse_bool, default='True',
-        help='Generate sensitivities'
-    )
-    parser.add_argument(
-        '--run-method', default='full',
-        type=partial(misc_utils.enum_parse, c=SensitivityCateg),
-        choices=SensitivityCateg,
-        help='Choose which type of sensivity plot to make'
-    )
-    parser.add_argument(
-        '--stat-method', default='bayesian',
-        type=partial(misc_utils.enum_parse, c=StatCateg), choices=StatCateg,
-        help='Statistical method to employ'
+        '--datadir', type=str, default='./untitled',
+        help='Path to store chains'
     )
     parser.add_argument(
-        '--sens-bins', type=int, default=10,
-        help='Number of bins for the Bayes factor plot'
+        '--segments', type=int, default=10,
+        help='Number of new physics scales to evaluate'
     )
     parser.add_argument(
-        '--sens-eval-bin', type=str, default='all',
-        help='Which bin to evalaute for Bayes factor plot'
-    )
-    parser.add_argument(
-        '--plot-statistic', type=misc_utils.parse_bool, default='False',
-        help='Plot MultiNest evidence or LLH value'
+        '--eval-segment', type=str, default='all',
+        help='Which point to evalaute'
     )
     fr_utils.fr_argparse(parser)
     gf_utils.gf_argparse(parser)
-    llh_utils.likelihood_argparse(parser)
+    llh_utils.llh_argparse(parser)
     mn_utils.mn_argparse(parser)
     nuisance_argparse(parser)
     if args is None: return parser.parse_args()
@@ -186,235 +189,72 @@ def main():
 
     asimov_paramset, llh_paramset = get_paramsets(args, define_nuisance())
 
-    scale = llh_paramset.from_tag(ParamTag.SCALE)[0]
-    mmangles = llh_paramset.from_tag(ParamTag.MMANGLES)
-    if args.run_method is SensitivityCateg.FULL:
-        st_paramset_arr = [llh_paramset.from_tag(ParamTag.SCALE, invert=True)]
-    elif args.run_method in [SensitivityCateg.FIXED_ANGLE, SensitivityCateg.CORR_ANGLE]:
-        nscale_pmset = llh_paramset.from_tag([ParamTag.SCALE, ParamTag.MMANGLES], invert=True)
-        st_paramset_arr = [nscale_pmset] * 3
-    elif args.run_method in [SensitivityCateg.FIXED_ONE_ANGLE, SensitivityCateg.CORR_ONE_ANGLE]:
-        nscale_pmset = llh_paramset.from_tag(ParamTag.SCALE, invert=True)
-        st_paramset_arr = []
-        for x in xrange(3):
-            st_paramset_arr.append(
-                ParamSet([prms for prms in nscale_pmset
-                          if mmangles[x].name != prms.name])
-            )
-
-    scan_scales = np.linspace(
-        np.log10(args.scale_region[0]), np.log10(args.scale_region[1]), args.sens_bins
+    # Scale and BSM mixings will be fixed.
+    scale_prm    = llh_paramset.from_tag(ParamTag.SCALE)[0]
+    base_mn_pset = llh_paramset.from_tag(
+        [ParamTag.SCALE, ParamTag.MMANGLES], invert=True
     )
-    scan_scales = np.concatenate([[-100.], scan_scales])
-    # scan_scales = np.array([-100.])
-
-    corr_angles_categ = [SensitivityCateg.CORR_ANGLE, SensitivityCateg.CORR_ONE_ANGLE]
-    fixed_angle_categ = [SensitivityCateg.FIXED_ANGLE, SensitivityCateg.FIXED_ONE_ANGLE]
 
-    if args.run_method in corr_angles_categ:
-        scan_angles = np.linspace(0+1e-9, 1-1e-9, args.sens_bins)
-    else: scan_angles = np.array([0])
-    print 'scan_scales', scan_scales
-    print 'scan_angles', scan_angles
+    # Array of scales to scan over.
+    boundaries = fr_utils.SCALE_BOUNDARIES[args.dimension]
+    eval_scales = np.linspace(boundaries[0], boundaries[1], args.segments-1)
+    eval_scales = np.concatenate([[-100.], eval_scales])
 
-    if args.sens_eval_bin is None:
-        eval_dim = args.sens_bins
+    # Evaluate just one point (job), or all points.
+    if args.eval_segment is None:
+        eval_dim = args.segments
     else: eval_dim = 1
 
-    out = args.outfile+'/{0}/{1}/{2}/fr_stat'.format(
-        *map(misc_utils.parse_enum, [args.stat_method, args.run_method, args.data])
+    outfile = args.datadir + '/{0}/{1}/fr_stat'.format(
+        *map(misc_utils.parse_enum, [args.stat_method, args.data])
     ) + misc_utils.gen_identifier(args)
-    if args.sens_run:
-        if args.likelihood in [Likelihood.GOLEMFIT, Likelihood.GF_FREQ]:
-            fitter = gf_utils.setup_fitter(args, asimov_paramset)
-            if args.stat_method is StatCateg.FREQUENTIST:
-                flags, gf_nuisance = gf_utils.fit_flags(llh_paramset)
-                llh_paramset = llh_paramset.remove_params(gf_nuisance)
-                asimov_paramset = asimov_paramset.remove_params(gf_nuisance)
-                st_paramset_arr = [x.remove_params(gf_nuisance)
-                                   for x in st_paramset_arr]
-                fitter.SetFitParametersFlag(flags)
-        else: fitter = None
-
-        if args.run_method is SensitivityCateg.FULL:
-            statistic_arr = np.full((eval_dim, 2), np.nan)
-        elif args.run_method in fixed_angle_categ:
-            statistic_arr = np.full((len(st_paramset_arr), eval_dim, 2), np.nan)
-        elif args.run_method in corr_angles_categ:
-            statistic_arr = np.full(
-                (len(st_paramset_arr), eval_dim, eval_dim, 3), np.nan
-            )
 
-        for idx_scen, sens_paramset in enumerate(st_paramset_arr):
-            print '|||| SCENARIO = {0}'.format(idx_scen)
-            if args.run_method in fixed_angle_categ:
-                for x in mmangles: x.value = 0.+1e-9
-                if idx_scen == 0 or idx_scen == 2:
-                    mmangles[idx_scen].value = np.sin(np.pi/4.)**2
-                    """s_12^2 or s_23^2"""
-                    mmangles[1].value = 1.
-                    """c_13^4"""
-                elif idx_scen == 1:
-                    mmangles[idx_scen].value = np.cos(np.pi/4.)**4
-                    """c_13^4"""
-
-            for idx_an, an in enumerate(scan_angles):
-                if args.run_method in corr_angles_categ:
-                    for x in mmangles: x.value = 0.+1e-9
-                    angle = np.arcsin(np.sqrt(an))/2.
-                    if idx_scen == 0 or idx_scen == 2:
-                        mmangles[idx_scen].value = np.sin(angle)**2
-                        """s_12^2 or s_23^2"""
-                        mmangles[1].value = 1.
-                        """c_13^4"""
-                    elif idx_scen == 1:
-                        mmangles[idx_scen].value = np.cos(angle)**4
-                        """c_13^4"""
-
-                for idx_sc, sc in enumerate(scan_scales):
-                    if args.sens_eval_bin is not None:
-                        if args.run_method in corr_angles_categ:
-                            index = idx_an*args.sens_bins + idx_sc
-                        else: index = idx_sc
-                        if index == args.sens_eval_bin:
-                            if idx_scen == 0:
-                                out += '_scale_{0:.0E}'.format(np.power(10, sc))
-                                if args.run_method in corr_angles_categ:
-                                    out += '_angle_{0:<04.2}'.format(an)
-                        else: continue
-
-                    sc_param = llh_paramset.from_tag(ParamTag.SCALE)[0]
-                    if sc < sc_param.ranges[0]:
-                        sc_param.ranges = (sc, sc_param.ranges[1])
-
-                    if idx_sc == 0 or args.sens_eval_bin is not None:
-                        print '|||| ANGLE = {0:<04.2}'.format(float(an))
-                    print '|||| SCALE = {0:.0E}'.format(np.power(10, sc))
-                    scale.value = sc
-                    if args.stat_method is StatCateg.BAYESIAN:
-                        identifier = 'b{0}_{1}_sce{2}_sca{3}_an{4}'.format(
-                            args.sens_eval_bin, args.sens_bins, idx_scen, sc, idx_an
-                        )
-                        try:
-                            stat = mn_utils.mn_evidence(
-                                mn_paramset     = sens_paramset,
-                                llh_paramset    = llh_paramset,
-                                asimov_paramset = asimov_paramset,
-                                args            = args,
-                                fitter          = fitter,
-                                identifier      = identifier
-                            )
-                        except:
-                            print 'Failed run, continuing'
-                            # raise
-                            continue
-                        print '## Evidence = {0}'.format(stat)
-                    elif args.stat_method is StatCateg.FREQUENTIST:
-                        raise NotImplementedError('Still needs testing')
-                        def fn(x):
-                            # TODO(shivesh): should be seed or ranges?
-                            # Force prior ranges to be inside "seed"
-                            for el in x:
-                                if el < 0 or el > 1: return np.inf
-                            pranges = sens_paramset.seeds
-                            for i, name in enumerate(sens_paramset.names):
-                                llh_paramset[name].value = \
-                                    (pranges[i][1]-pranges[i][0])*x[i] + pranges[i][0]
-                            theta = llh_paramset.values
-                            llh = llh_utils.ln_prob(
-                                theta=theta, args=args, asimov_paramset=asimov_paramset,
-                                llh_paramset=llh_paramset, fitter=fitter
-                            )
-                            # print 'llh_paramset', llh_paramset
-                            # print '-llh', -llh
-                            return -llh
-
-                        n_params = len(sens_paramset)
-                        bounds = np.dstack([np.zeros(n_params), np.ones(n_params)])[0]
-                        x0_arr = np.linspace(0+1e-3, 1-1e-3, 3)
-                        stat = -np.inf
-                        try:
-                            for x0_v in x0_arr:
-                                x0 = np.full(n_params, x0_v)
-                                # res = minimize(fun=fn, x0=x0, method='Powell', tol=1e-3)
-                                res = minimize(fun=fn, x0=x0, method='Nelder-Mead', tol=1e-3)
-                                # res = minimize(fun=fn, x0=x0, method='L-BFGS-B', tol=1e-3, bounds=bounds)
-                                # res = minimize(fun=fn, x0=x0, method='BFGS', tol=1e-3)
-                                s = -fn(res.x)
-                                if s > stat:
-                                    stat = s
-                        except AssertionError:
-                            # print 'Failed run, continuing'
-                            raise
-                            continue
-                        print '=== final llh', stat
-                    if args.run_method is SensitivityCateg.FULL:
-                        statistic_arr[idx_sc] = np.array([sc, stat])
-                    elif args.run_method in fixed_angle_categ:
-                        if args.sens_eval_bin is not None:
-                            statistic_arr[idx_scen][0] = np.array([sc, stat])
-                        else:
-                            statistic_arr[idx_scen][idx_sc] = np.array([sc, stat])
-                    elif args.run_method in corr_angles_categ:
-                        if args.sens_eval_bin is not None:
-                            statistic_arr[idx_scen][0][0] = np.array([an, sc, stat])
-                        else:
-                            statistic_arr[idx_scen][idx_an][idx_sc] = np.array([an, sc, stat])
-
-        misc_utils.make_dir(out)
-        print 'Saving to {0}'.format(out+'.npy')
-        np.save(out+'.npy', statistic_arr)
-
-    if args.plot_statistic:
-        print 'Plotting statistic'
-        if args.sens_run: raw = statistic_arr
-        else: raw = np.load(out+'.npy')
-        data = ma.masked_invalid(raw)
-
-        basename = os.path.dirname(out) + '/statrun/' + os.path.basename(out)
-        baseoutfile = basename[:5]+basename[5:].replace('data', 'plots')
-        if args.run_method is SensitivityCateg.FULL:
-            plot_utils.plot_statistic(
-                data        = data,
-                outfile     = baseoutfile,
-                outformat   = ['png'],
-                args        = args,
-                scale_param = scale
+    # Setup Golemfit.
+    gf_utils.setup_fitter(args, asimov_paramset)
+
+    # Initialise data structure.
+    stat_arr = np.full((eval_dim, 2), np.nan)
+
+    for idx_sc, scale in enumerate(eval_scales):
+        if args.eval_segment is not None:
+            if idx_sc == args.eval_segment:
+                outfile += '_scale_{0:.0E}'.format(np.power(10, scale))
+            else: continue
+        print '|||| SCALE = {0:.0E}'.format(np.power(10, scale))
+
+        # Lower scale boundary for first (NULL) point and set the scale param.
+        if scale < scale_prm.ranges[0]:
+            scale_prm.ranges = (scale, scale_prm.ranges[1])
+        scale_prm.value = scale
+
+        if args.stat_method is StatCateg.BAYESIAN:
+            identifier = 'b{0}_{1}_sce{2}_sca{3}'.format(
+                args.eval_segment, args.segments, int(args.texture), scale
             )
-        elif args.run_method in fixed_angle_categ:
-            for idx_scen in xrange(len(st_paramset_arr)):
-                print '|||| SCENARIO = {0}'.format(idx_scen)
-                outfile = baseoutfile + '_SCEN{0}'.format(idx_scen)
-                if idx_scen == 0: label = r'$\mathcal{O}_{12}=\pi/4$'
-                elif idx_scen == 1: label = r'$\mathcal{O}_{13}=\pi/4$'
-                elif idx_scen == 2: label = r'$\mathcal{O}_{23}=\pi/4$'
-                plot_utils.plot_statistic(
-                    data        = data[idx_scen],
-                    outfile     = outfile,
-                    outformat   = ['png'],
-                    args        = args,
-                    scale_param = scale,
-                    label       = label
+            try:
+                stat = mn_utils.mn_evidence(
+                    mn_paramset     = base_mn_pset,
+                    llh_paramset    = llh_paramset,
+                    asimov_paramset = asimov_paramset,
+                    args            = args,
+                    identifier      = identifier
                 )
-        elif args.run_method in corr_angles_categ:
-            for idx_scen in xrange(len(st_paramset_arr)):
-                print '|||| SCENARIO = {0}'.format(idx_scen)
-                basescenoutfile = baseoutfile + '_SCEN{0}'.format(idx_scen)
-                if idx_scen == 0: label = r'$\mathcal{O}_{12}='
-                elif idx_scen == 1: label = r'$\mathcal{O}_{13}='
-                elif idx_scen == 2: label = r'$\mathcal{O}_{23}='
-                for idx_an, an in enumerate(scan_angles):
-                    print '|||| ANGLE = {0:<04.2}'.format(float(an))
-                    outfile = basescenoutfile + '_ANGLE{0}'.format(idx_an)
-                    _label = label + r'{0:<04.2}$'.format(an)
-                    plot_utils.plot_statistic(
-                        data        = data[idx_scen][index][:,1:],
-                        outfile     = outfile,
-                        outformat   = ['png'],
-                        args        = args,
-                        scale_param = scale,
-                        label       = _label
-                    )
+            except:
+                print 'Failed run, continuing'
+                # raise
+                continue
+            print '## Evidence = {0}'.format(stat)
+        elif args.stat_method is StatCateg.FREQUENTIST:
+            raise NotImplementedError('Still needs testing')
+
+        if args.eval_segment is not None:
+            stat_arr[0] = np.array([scale, stat])
+        else:
+            stat_arr[idx_sc] = np.array([scale, stat])
+
+    misc_utils.make_dir(outfile)
+    print 'Saving to {0}'.format(outfile+'.npy')
+    np.save(outfile+'.npy', stat_arr)
 
 
 main.__doc__ = __doc__