Tue Apr 24 11:22:19 CDT 2018

8 files changed, 160 insertions, 60 deletions

diff --git a/utils/enums.py b/utils/enums.py
index 359e104..8a9e868 100644
--- a/utils/enums.py
+++ b/utils/enums.py
@@ -21,11 +21,6 @@ class EnergyDependance(Enum):
     SPECTRAL = 2
 
 
-class FitPriorsCateg(Enum):
-    DEFAULT = 1
-    XS      = 2
-
-
 class Likelihood(Enum):
     FLAT     = 1
     GAUSSIAN = 2
@@ -35,11 +30,17 @@ class Likelihood(Enum):
 
 class ParamTag(Enum):
     NUISANCE  = 1
-    MMANGLES  = 2
-    SCALE     = 3
-    SRCANGLES = 4
-    BESTFIT   = 5
-    NONE      = 6
+    SM_ANGLES = 2
+    MMANGLES  = 3
+    SCALE     = 4
+    SRCANGLES = 5
+    BESTFIT   = 6
+    NONE      = 7
+
+
+class PriorsCateg(Enum):
+    UNIFORM  = 1
+    GAUSSIAN = 2
 
 
 class MCMCSeedType(Enum):
diff --git a/utils/fr.py b/utils/fr.py
index 6a405a7..342a848 100644
--- a/utils/fr.py
+++ b/utils/fr.py
@@ -198,17 +198,18 @@ def normalise_fr(fr):
     return np.array(fr) / float(np.sum(fr))
 
 
-def estimate_scale(args, mass_eigenvalues=MASS_EIGENVALUES):
+def estimate_scale(args):
     """Estimate the scale at which new physics will enter."""
+    m_eign = args.m3x_2
     if args.energy_dependance is EnergyDependance.MONO:
         scale = np.power(
-            10, np.round(np.log10(MASS_EIGENVALUES[1]/args.energy)) - \
+            10, np.round(np.log10(m_eign/args.energy)) - \
             np.log10(args.energy**(args.dimension-3))
         )
     elif args.energy_dependance is EnergyDependance.SPECTRAL:
         scale = np.power(
             10, np.round(
-                np.log10(MASS_EIGENVALUES[1]/np.power(10, np.average(np.log10(args.binning)))) \
+                np.log10(m_eign/np.power(10, np.average(np.log10(args.binning)))) \
                 - np.log10(np.power(10, np.average(np.log10(args.binning)))**(args.dimension-3))
             )
         )
@@ -297,7 +298,7 @@ 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,
+                   sm_u=NUFIT_U, no_bsm=False, fix_mixing=False,
                    fix_mixing_almost=False, fix_scale=False, scale=None,
                    check_uni=True, epsilon=1e-9):
     """Construct the BSM mixing matrix from the BSM parameters.
@@ -316,8 +317,8 @@ def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
     mass_eigenvalues : list, length = 2
         SM mass eigenvalues
 
-    nufit_u : numpy ndarray, dimension 3
-        SM NuFIT mixing matrix
+    sm_u : numpy ndarray, dimension 3
+        SM mixing matrix
 
     no_bsm : bool
         Turn off BSM behaviour
@@ -350,7 +351,7 @@ def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
            [-0.32561308 -3.95946524e-01j,  0.64294909 -2.23453580e-01j, 0.03700830 +5.22032403e-01j]])
 
     """
-    if np.shape(nufit_u) != (3, 3):
+    if np.shape(sm_u) != (3, 3):
         raise ValueError(
             'Input matrix should be a square and dimension 3, '
             'got\n{0}'.format(ham)
@@ -377,7 +378,7 @@ def params_to_BSMu(theta, dim, energy, mass_eigenvalues=MASS_EIGENVALUES,
     mass_matrix = np.array(
         [[0, 0, 0], [0, mass_eigenvalues[0], 0], [0, 0, mass_eigenvalues[1]]]
     )
-    sm_ham = (1./(2*energy))*np.dot(nufit_u, np.dot(mass_matrix, nufit_u.conj().T))
+    sm_ham = (1./(2*energy))*np.dot(sm_u, np.dot(mass_matrix, sm_u.conj().T))
     if no_bsm:
         eg_vector = cardano_eqn(sm_ham)
     else:
diff --git a/utils/gf.py b/utils/gf.py
index 20afc75..59d1033 100644
--- a/utils/gf.py
+++ b/utils/gf.py
@@ -18,6 +18,37 @@ from utils.enums import DataType, SteeringCateg
 from utils.misc import enum_parse, thread_factors
 
 
+def fit_flags(llh_paramset):
+    default_flags = {
+        # False means it's not fixed in minimization
+        'astroFlavorAngle1'         : True,
+        'astroFlavorAngle2'         : True,
+        'astroENorm'                : True,
+        'astroMuNorm'               : True,
+        'astroTauNorm'              : True,
+        'convNorm'                  : True,
+        'promptNorm'                : True,
+        'muonNorm'                  : True,
+        'astroNorm'                 : True,
+        'astroParticleBalance'      : True,
+        'astroDeltaGamma'           : True,
+        'cutoffEnergy'              : True,
+        'CRDeltaGamma'              : True,
+        'piKRatio'                  : True,
+        'NeutrinoAntineutrinoRatio' : True,
+        'darkNorm'                  : True,
+        'domEfficiency'             : True,
+        'holeiceForward'            : True,
+        'anisotropyScale'           : True,
+        'astroNormSec'              : True,
+        'astroDeltaGammaSec'        : True
+    }
+    flags = gf.FitParametersFlag()
+    for param in llh_paramset:
+        flags.__setattr__(param.name, False)
+    return flags
+
+
 def steering_params(args):
     steering_categ = args.ast
     params = gf.SteeringParams()
diff --git a/utils/likelihood.py b/utils/likelihood.py
index 1981c72..04828e8 100644
--- a/utils/likelihood.py
+++ b/utils/likelihood.py
@@ -18,16 +18,21 @@ import GolemFitPy as gf
 
 from utils import fr as fr_utils
 from utils import gf as gf_utils
-from utils.enums import EnergyDependance, Likelihood, ParamTag
+from utils.enums import EnergyDependance, Likelihood, ParamTag, PriorsCateg
 from utils.misc import enum_parse
 
 
-def gaussian_llh(fr, fr_bf, sigma):
+def multi_gaussian(fr, fr_bf, sigma):
     """Multivariate gaussian likelihood."""
     cov_fr = np.identity(3) * sigma
     return np.log(multivariate_normal.pdf(fr, mean=fr_bf, cov=cov_fr))
 
 
+def log_gaussian(x, mu, sig):
+    """Log gaussian for one dimension."""
+    return np.log((1./np.sqrt(2*np.pi*sig**2))*np.exp(-((x-mu)**2)/(2*sig**2)))
+
+
 def likelihood_argparse(parser):
     parser.add_argument(
         '--likelihood', default='gaussian', type=partial(enum_parse, c=Likelihood),
@@ -41,16 +46,25 @@ def likelihood_argparse(parser):
 
 def lnprior(theta, paramset):
     """Priors on theta."""
+    for idx, param in enumerate(paramset):
+        param.value = theta[idx]
     ranges = paramset.ranges
     for value, range in zip(theta, ranges):
         if range[0] <= value <= range[1]:
             pass
         else: return -np.inf
-    return 0.
+
+    prior = 0
+    for param in paramset:
+        if param.prior is PriorsCateg.GAUSSIAN:
+            prior += log_gaussian(
+                param.nominal_value, param.value, param.std
+            )
+    return prior
 
 
 def triangle_llh(theta, args, asimov_paramset, llh_paramset, fitter):
-    """-Log likelihood function for a given theta."""
+    """Log likelihood function for a given theta."""
     if len(theta) != len(llh_paramset):
         raise AssertionError(
             'Length of MCMC scan is not the same as the input '
@@ -89,11 +103,21 @@ def triangle_llh(theta, args, asimov_paramset, llh_paramset, fitter):
         [ParamTag.SCALE, ParamTag.MMANGLES], values=True
     )
 
+    m_eig_names = ['m21_2', 'm3x_2']
+    mass_eigenvalues = [x.value for x in llh_paramset if x.name in m_eig_names]
+
+    ma_names = ['s_12_2', 'c_13_4', 's_23_2', 'dcp']
+    sm_u = fr_utils.angles_to_u(
+        [x.value for x in llh_paramset if x.name in ma_names]
+    )
+
     if args.energy_dependance is EnergyDependance.MONO:
         u = fr_utils.params_to_BSMu(
             theta             = bsm_angles,
             dim               = args.dimension,
             energy            = args.energy,
+            mass_eigenvalues  = mass_eigenvalues,
+            sm_u              = sm_u,
             no_bsm            = args.no_bsm,
             fix_mixing        = args.fix_mixing,
             fix_mixing_almost = args.fix_mixing_almost,
@@ -108,6 +132,8 @@ def triangle_llh(theta, args, asimov_paramset, llh_paramset, fitter):
                 theta             = bsm_angles,
                 dim               = args.dimension,
                 energy            = args.energy,
+                mass_eigenvalues  = mass_eigenvalues,
+                sm_u              = sm_u,
                 no_bsm            = args.no_bsm,
                 fix_mixing        = args.fix_mixing,
                 fix_mixing_almost = args.fix_mixing_almost,
@@ -130,14 +156,14 @@ def triangle_llh(theta, args, asimov_paramset, llh_paramset, fitter):
         return 1.
     elif args.likelihood is Likelihood.GAUSSIAN:
         fr_bf = args.measured_ratio
-        return gaussian_llh(fr, fr_bf, args.sigma_ratio)
+        return multi_gaussian(fr, fr_bf, args.sigma_ratio)
     elif args.likelihood is Likelihood.GOLEMFIT:
         return gf_utils.get_llh(fitter, hypo_paramset)
     elif args.likelihood is Likelihood.GF_FREQ:
         return gf_utils.get_llh_freq(fitter, hypo_paramset)
 
 
-def ln_prob(theta, args, fitter, asimov_paramset, llh_paramset):
+def ln_prob(theta, args, asimov_paramset, llh_paramset, fitter):
     lp = lnprior(theta, paramset=llh_paramset)
     if not np.isfinite(lp):
         return -np.inf
diff --git a/utils/mcmc.py b/utils/mcmc.py
index 97b77dd..fa044ea 100644
--- a/utils/mcmc.py
+++ b/utils/mcmc.py
@@ -72,11 +72,11 @@ def mcmc_argparse(parser):
         help='Type of distrbution to make the initial MCMC seed'
     )
     parser.add_argument(
-        '--plot-angles', type=misc_utils.parse_bool, default='True',
+        '--plot-angles', type=parse_bool, default='True',
         help='Plot MCMC triangle in the angles space'
     )
     parser.add_argument(
-        '--plot-elements', type=misc_utils.parse_bool, default='False',
+        '--plot-elements', type=parse_bool, default='False',
         help='Plot MCMC triangle in the mixing elements space'
     )
 
diff --git a/utils/multinest.py b/utils/multinest.py
index 3a7ab2d..426a951 100644
--- a/utils/multinest.py
+++ b/utils/multinest.py
@@ -15,13 +15,27 @@ import numpy as np
 
 from pymultinest import analyse, run
 
-from utils.likelihood import triangle_llh
+from utils import likelihood
 from utils.misc import make_dir, parse_bool
 
 
-def CubePrior(cube, ndim, n_params):
-    # default are uniform priors
-    return ;
+def CubePrior(cube, ndim, n_params, mn_paramset):
+    if ndim != len(mn_paramset):
+        raise AssertionError(
+            'Length of MultiNest scan paramset is not the same as the input '
+            'params\ncube={0}\nmn_paramset]{1}'.format(cube, mn_paramset)
+        )
+    pranges = mn_paramset.seeds
+    for i in xrange(ndim):
+        mn_paramset[i].value = (pranges[i][1]-pranges[i][0])*cube[i] + pranges[i][0]
+    prior = 0
+    for parm in mn_paramset:
+        if parm.prior is PriorsCateg.GAUSSIAN:
+            prior_penatly += multivariate_normal(
+                parm.nominal_value, mean=parm.value, cov=parm.std
+            )
+    print 'prior', prior
+    return prior
 
 
 def lnProb(cube, ndim, n_params, mn_paramset, llh_paramset, asimov_paramset,
@@ -38,7 +52,7 @@ def lnProb(cube, ndim, n_params, mn_paramset, llh_paramset, asimov_paramset,
         llh_paramset[pm].value = mn_paramset[pm].value
     theta = llh_paramset.values
     # print 'llh_paramset', llh_paramset
-    return triangle_llh(
+    return likelihood.ln_prob(
         theta=theta,
         args=args,
         asimov_paramset=asimov_paramset,
@@ -72,10 +86,6 @@ def mn_argparse(parser):
         '--mn-output', type=str, default='./mnrun/',
         help='Folder to store MultiNest evaluations'
     )
-    parser.add_argument(
-        '--plot-multinest', type=parse_bool, default='False',
-        help='Plot MultiNest evidence'
-    )
 
 
 def mn_evidence(mn_paramset, llh_paramset, asimov_paramset, args, fitter):
diff --git a/utils/param.py b/utils/param.py
index bf230df..13f1a63 100644
--- a/utils/param.py
+++ b/utils/param.py
@@ -12,17 +12,20 @@ from __future__ import absolute_import, division
 import sys
 
 from collections import Sequence
+from copy import deepcopy
 
 import numpy as np
 
 from utils.plot import get_units
 from utils.fr import fr_to_angles
-from utils.enums import Likelihood, ParamTag
+from utils.enums import Likelihood, ParamTag, PriorsCateg
 
 
 class Param(object):
     """Parameter class to store parameters."""
-    def __init__(self, name, value, ranges, seed=None, std=None, tex=None, tag=None):
+    def __init__(self, name, value, ranges, prior=None, seed=None, std=None,
+                 tex=None, tag=None):
+        self._prior = None
         self._seed = None
         self._ranges = None
         self._tex = None
@@ -30,8 +33,10 @@ class Param(object):
 
         self.name = name
         self.value = value
-        self.seed = seed
+        self.nominal_value = deepcopy(value)
+        self.prior = prior
         self.ranges = ranges
+        self.seed = seed
         self.std = std
         self.tex = tex
         self.tag = tag
@@ -45,6 +50,18 @@ class Param(object):
         self._ranges = [val for val in values]
 
     @property
+    def prior(self):
+        return self._prior
+
+    @prior.setter
+    def prior(self, value):
+        if value is None:
+            self._prior = PriorsCateg.UNIFORM
+        else:
+            assert value in PriorsCateg
+            self._prior = value
+
+    @property
     def seed(self):
         if self._seed is None: return self.ranges
         return tuple(self._seed)
@@ -146,6 +163,10 @@ class ParamSet(Sequence):
         return tuple([obj.value for obj in self._params])
 
     @property
+    def nominal_values(self):
+        return tuple([obj.nominal_value for obj in self._params])
+
+    @property
     def seeds(self):
         return tuple([obj.seed for obj in self._params])
 
@@ -185,18 +206,22 @@ class ParamSet(Sequence):
             return ParamSet([io[1] for io in ps])
 
 
-def get_paramsets(args):
+def get_paramsets(args, nuisance_paramset):
     """Make the paramsets for generating the Asmimov MC sample and also running
     the MCMC.
     """
     asimov_paramset = []
     llh_paramset = []
-    if args.likelihood == Likelihood.GOLEMFIT:
-        nuisance_paramset = define_nuisance()
-        asimov_paramset.extend(nuisance_paramset.params)
-        llh_paramset.extend(nuisance_paramset.params)
-        for parm in nuisance_paramset:
-            parm.value = args.__getattribute__(parm.name)
+
+    llh_paramset.extend(
+        [x for x in nuisance_paramset.from_tag(ParamTag.SM_ANGLES)]
+    )
+    if args.likelihood is Likelihood.GOLEMFIT:
+        gf_nuisance = [x for x in nuisance_paramset.from_tag(ParamTag.NUISANCE)]
+        asimov_paramset.extend(gf_nuisance)
+        llh_paramset.extend(gf_nuisance)
+    for parm in llh_paramset:
+        parm.value = args.__getattribute__(parm.name)
     tag = ParamTag.BESTFIT
     flavour_angles = fr_to_angles(args.measured_ratio)
     asimov_paramset.extend([
@@ -207,17 +232,16 @@ def get_paramsets(args):
 
     if not args.fix_mixing and not args.fix_mixing_almost:
         tag = ParamTag.MMANGLES
-        e = 1e-10
         llh_paramset.extend([
-            Param(name='s_12^2', value=0.5, ranges=[0.+e, 1.-e], std=0.2, tex=r'\tilde{s}_{12}^2', tag=tag),
-            Param(name='c_13^4', value=0.5, ranges=[0.+e, 1.-e], std=0.2, tex=r'\tilde{c}_{13}^4', tag=tag),
-            Param(name='s_23^2', value=0.5, ranges=[0.+e, 1.-e], std=0.2, tex=r'\tilde{s}_{23}^4', tag=tag),
-            Param(name='dcp', value=np.pi, ranges=[0.+e, 2*np.pi-e], std=0.2, tex=r'\tilde{\delta_{CP}}', tag=tag)
+            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}^4', tag=tag),
+            Param(name='np_dcp', value=np.pi, ranges=[0., 2*np.pi], std=0.2, tex=r'\tilde{\delta_{CP}}', tag=tag)
         ])
     if args.fix_mixing_almost:
         tag = ParamTag.MMANGLES
         llh_paramset.extend([
-            Param(name='s_23^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{23}^4', tag=tag)
+            Param(name='np_s_23^2', value=0.5, ranges=[0., 1.], std=0.2, tex=r'\tilde{s}_{23}^4', tag=tag)
         ])
     if not args.fix_scale:
         tag = ParamTag.SCALE
diff --git a/utils/plot.py b/utils/plot.py
index 66c888c..6ae8dc2 100644
--- a/utils/plot.py
+++ b/utils/plot.py
@@ -24,7 +24,7 @@ from getdist import plots
 from getdist import mcsamples
 
 from utils import misc as misc_utils
-from utils.enums import EnergyDependance, Likelihood, ParamTag
+from utils.enums import EnergyDependance, Likelihood, ParamTag, StatCateg
 from utils.fr import angles_to_u, angles_to_fr
 
 rc('text', usetex=False)
@@ -231,25 +231,32 @@ def chainer_plot(infile, outfile, outformat, args, llh_paramset):
             g.export(outfile+'_elements.'+of)
 
 
-def plot_multinest(data, outfile, outformat, args, scale_param, label=None):
-    """Make MultiNest factor plot."""
+def plot_statistic(data, outfile, outformat, args, scale_param, label=None):
+    """Make MultiNest factor or LLH value plot."""
     print "Making MultiNest Factor plot"
     fig_text = gen_figtext(args)
     if label is not None: fig_text += '\n' + label
 
     print 'data.shape', data.shape
-    scales, evidences = data.T
-    min_idx = np.argmin(scales)
-    null = evidences[min_idx]
-
-    reduced_ev = -(evidences - null)
+    scales, statistic = data.T
+    if args.stat_method is StatCateg.BAYESIAN:
+        min_idx = np.argmin(scales)
+        null = statistic[min_idx]
+        reduced_ev = -(statistic - null)
+    elif args.stat_method is StatCateg.FREQUENTIST:
+        min_idx = np.argmin(statistic)
+        null = statistic[min_idx]
+        reduced_ev = 2*(statistic - null)
 
     fig = plt.figure(figsize=(7, 5))
     ax = fig.add_subplot(111)
 
     ax.set_xlim(scale_param.ranges)
     ax.set_xlabel('$'+scale_param.tex+'$')
-    ax.set_ylabel(r'Bayes Factor')
+    if args.stat_method is StatCateg.BAYESIAN:
+        ax.set_ylabel(r'Bayes Factor')
+    elif args.stat_method is StatCateg.FREQUENTIST:
+        ax.set_ylabel(r'$\Delta {\rm LLH}$')
 
     ax.plot(scales, reduced_ev)