Wed 26 Jun 17:34:04 BST 2019

1 files changed, 19 insertions, 18 deletions

diff --git a/mc_unitary.py b/mc_unitary.py
index 40ec61e..12a2db2 100755
--- a/mc_unitary.py
+++ b/mc_unitary.py
@@ -5,7 +5,7 @@
 # date   : March 17, 2018
 
 """
-Sample SM points for (x, 1-x, 0) only assuming unitarity
+Sample points only assuming unitarity
 """
 
 from __future__ import absolute_import, division
@@ -38,10 +38,6 @@ def define_nuisance():
         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),
     ])
-    tag = ParamTag.SRCANGLES
-    nuisance.extend([
-        Param(name='astroX', value=0.5, seed=[0., 1.], ranges=[0., 1.], std=0.1, tex=r'x', tag=tag)
-    ])
     return ParamSet(nuisance)
 
 
@@ -53,9 +49,7 @@ def get_paramsets(args, nuisance_paramset):
     hypo_paramset = []
 
     hypo_paramset.extend(
-        [x for x in nuisance_paramset.from_tag((
-            ParamTag.SM_ANGLES, ParamTag.SRCANGLES
-        ))]
+        [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)]
     )
 
     for parm in hypo_paramset:
@@ -64,9 +58,11 @@ def get_paramsets(args, nuisance_paramset):
     hypo_paramset = ParamSet(hypo_paramset)
 
     tag = ParamTag.BESTFIT
+    flavour_angles = fr_utils.fr_to_angles(args.source_ratio)
+
     asimov_paramset.extend([
-        Param(name='astroFlavorAngle1', value=0., ranges=[ 0., 1.], std=0.2, tag=tag),
-        Param(name='astroFlavorAngle2', value=0., ranges=[-1., 1.], std=0.2, tag=tag),
+        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)
 
@@ -84,7 +80,7 @@ def nuisance_argparse(parser):
 
 def process_args(args):
     """Process the input args."""
-    pass
+    args.source_ratio = fr_utils.normalise_fr(args.source_ratio)
 
 
 def parse_args(args=None):
@@ -94,6 +90,10 @@ def parse_args(args=None):
         formatter_class=misc_utils.SortingHelpFormatter,
     )
     parser.add_argument(
+        '--source-ratio', type=float, nargs=3, default=[1, 2, 0],
+        help='Set the source flavour ratio'
+    )
+    parser.add_argument(
         '--seed', type=misc_utils.seed_parse, default='26',
         help='Set the random seed value'
     )
@@ -107,13 +107,17 @@ def parse_args(args=None):
     )
     mcmc_utils.mcmc_argparse(parser)
     nuisance_argparse(parser)
-    misc_utils.remove_option(parser, 'injected_ratio')
     misc_utils.remove_option(parser, 'plot_angles')
     misc_utils.remove_option(parser, 'plot_elements')
     if args is None: return parser.parse_args()
     else: return parser.parse_args(args.split())
 
 
+def gen_identifier(args):
+    f = '_SRC_{0}'.format(misc_utils.solve_ratio(args.source_ratio))
+    return f
+
+
 def triangle_llh(theta, args, hypo_paramset):
     """Log likelihood function for a given theta."""
     if len(theta) != len(hypo_paramset):
@@ -150,7 +154,7 @@ def main():
     asimov_paramset, hypo_paramset = get_paramsets(args, define_nuisance())
 
     prefix = ''
-    outfile = args.datadir + '/mc_unitarity' + prefix
+    outfile = args.datadir + '/mc_unitary' + prefix + gen_identifier(args)
     print '== {0:<25} = {1}'.format('outfile', outfile)
 
     print 'asimov_paramset', asimov_paramset
@@ -183,11 +187,9 @@ def main():
             threads  = args.threads
         )
 
-        nsamples = len(samples)
-        srcs = [fr_utils.normalise_fr((x, 1-x, 0)) for x in samples.T[-1]]
-        mmxs = map(fr_utils.angles_to_u, samples.T[:-1].T)
+        mmxs = map(fr_utils.angles_to_u, samples)
         frs = np.array(
-            [fr_utils.u_to_fr(srcs[i], mmxs[i]) for i in xrange(nsamples)]
+            [fr_utils.u_to_fr(args.source_ratio, x) for x in mmxs]
         )
         mcmc_utils.save_chains(frs, outfile)
 
@@ -199,4 +201,3 @@ main.__doc__ = __doc__
 
 if __name__ == '__main__':
     main()
-