add option to fix all but a single BSM mixing angle

4 files changed, 42 insertions, 16 deletions

diff --git a/utils/fr.py b/utils/fr.py
index 4242e97..7f9d855 100644
--- a/utils/fr.py
+++ b/utils/fr.py
@@ -200,7 +200,8 @@ NUFIT_U = angles_to_u((0.307, (1-0.02195)**2, 0.565, 3.97935))
 
 def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
                    nufit_u=NUFIT_U, no_bsm=False, fix_mixing=False,
-                   fix_scale=False, scale=None, check_uni=True):
+                   fix_mixing_almost=False, fix_scale=False, scale=None,
+                   check_uni=True):
     """Construct the BSM mixing matrix from the BSM parameters.
 
     Parameters
@@ -226,6 +227,9 @@ def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
     fix_mixing : bool
         Fix the BSM mixing angles
 
+    fix_mixing_almost : bool
+        Fix the BSM mixing angles except one
+
     fix_scale : bool
         Fix the BSM scale
 
@@ -254,8 +258,16 @@ def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
             'got\n{0}'.format(ham)
         )
 
+    if fix_mixing and fix_mixing_almost:
+        raise NotImplementedError(
+            '--fix-mixing and --fix-mixing-almost cannot be used together'
+        )
+
     if fix_mixing:
         s12_2, c13_4, s23_2, dcp, sc2 = 0.5, 1.0-1E-6, 0.5, 0., theta
+    elif fix_mixing_almost:
+        s12_2, c13_4, dcp = 0.5, 1.0-1E-6, 0.
+        s23_2, sc2 = theta
     elif fix_scale:
         s12_2, c13_4, s23_2, dcp = theta
         sc2 = np.log10(scale)
diff --git a/utils/likelihood.py b/utils/likelihood.py
index ec64823..fe4301d 100644
--- a/utils/likelihood.py
+++ b/utils/likelihood.py
@@ -88,26 +88,28 @@ def triangle_llh(theta, args, asimov_paramset, mcmc_paramset, fitter):
 
     if args.energy_dependance is EnergyDependance.MONO:
         u = fr_utils.params_to_BSMu(
-            theta      = bsm_angles,
-            dim        = args.dimension,
-            energy     = args.energy,
-            no_bsm     = args.no_bsm,
-            fix_mixing = args.fix_mixing,
-            fix_scale  = args.fix_scale,
-            scale      = args.scale
+            theta             = bsm_angles,
+            dim               = args.dimension,
+            energy            = args.energy,
+            no_bsm            = args.no_bsm,
+            fix_mixing        = args.fix_mixing,
+            fix_mixing_almost = args.fix_mixing_almost,
+            fix_scale         = args.fix_scale,
+            scale             = args.scale
         )
         fr = fr_utils.u_to_fr(source_flux, u)
     elif args.energy_dependance is EnergyDependance.SPECTRAL:
         mf_perbin = []
         for i_sf, sf_perbin in enumerate(source_flux):
-            u = fr_utils.params_to_BSMu(
-                theta      = bsm_angles,
-                dim        = args.dimension,
-                energy     = args.energy,
-                no_bsm     = args.no_bsm,
-                fix_mixing = args.fix_mixing,
-                fix_scale  = args.fix_scale,
-                scale      = args.scale
+            u                     = fr_utils.params_to_BSMu(
+                theta             = bsm_angles,
+                dim               = args.dimension,
+                energy            = args.energy,
+                no_bsm            = args.no_bsm,
+                fix_mixing        = args.fix_mixing,
+                fix_mixing_almost = args.fix_mixing_almost,
+                fix_scale         = args.fix_scale,
+                scale             = args.scale
             )
             fr = fr_utils.u_to_fr(sf_perbin, u)
             mf_perbin.append(fr)
diff --git a/utils/misc.py b/utils/misc.py
index a189f07..331e78a 100644
--- a/utils/misc.py
+++ b/utils/misc.py
@@ -214,6 +214,11 @@ def gen_outfile_name(args):
                 int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000),
                 int(sr[0]*100), int(sr[1]*100), int(sr[2]*100), args.dimension
             )
+        elif args.fix_mixing_almost:
+            outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fix_mixing_almost'.format(
+                int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000),
+                int(sr[0]*100), int(sr[1]*100), int(sr[2]*100), args.dimension
+            )
         elif args.fix_scale:
             outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_sfr_{4:03d}_{5:03d}_{6:03d}_DIM{7}_fixed_scale_{8}'.format(
                 int(mr[0]*100), int(mr[1]*100), int(mr[2]*100), int(si*1000),
@@ -231,6 +236,11 @@ def gen_outfile_name(args):
                 int(mr[0]*100), int(mr[1]*100), int(mr[2]*100),
                 int(si*1000), args.dimension
             )
+        elif args.fix_mixing_almost:
+            outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fix_mixing_almost'.format(
+                int(mr[0]*100), int(mr[1]*100), int(mr[2]*100),
+                int(si*1000), args.dimension
+            )
         elif args.fix_scale:
             outfile = args.outfile+'_{0:03d}_{1:03d}_{2:03d}_{3:04d}_DIM{4}_fixed_scale_{5}'.format(
                 int(mr[0]*100), int(mr[1]*100), int(mr[2]*100),
diff --git a/utils/plot.py b/utils/plot.py
index c83181c..c70ffec 100644
--- a/utils/plot.py
+++ b/utils/plot.py
@@ -198,6 +198,8 @@ def chainer_plot(infile, outfile, outformat, args, mcmc_paramset):
             g.export(outfile+'_angles.'+of)
 
     if args.plot_elements:
+        if args.fix_mixing_almost:
+            raise NotImplementedError
         nu_index = mcmc_paramset.from_tag(ParamTag.NUISANCE, index=True)
         fr_index = mcmc_paramset.from_tag(ParamTag.MMANGLES, index=True)
         sc_index = mcmc_paramset.from_tag(ParamTag.SCALE, index=True)