intergrate GolemFit into scripts

7 files changed, 226 insertions, 112 deletions

diff --git a/fr.py b/fr.py
index 4c95f81..c5dffaf 100755
--- a/fr.py
+++ b/fr.py
@@ -15,11 +15,14 @@ from functools import partial
 
 import numpy as np
 
+import GolemFitPy as gf
+
+from utils import gf as gf_utils
+from utils import likelihood as llh_utils
 from utils import mcmc as mcmc_utils
 from utils import misc as misc_utils
 from utils.enums import Likelihood, ParamTag
 from utils.fr import MASS_EIGENVALUES, normalise_fr
-from utils.gf import gf_argparse
 from utils.misc import Param, ParamSet
 from utils.plot import plot_argparse, chainer_plot
 
@@ -56,14 +59,15 @@ def get_paramsets(args):
     mcmc_paramset = []
     if args.likelihood == Likelihood.GOLEMFIT:
         nuisance_paramset = define_nuisance()
-        asimov_paramset.extend([nuisance_paramset.params])
-        mcmc_paramset.extend([nuisance_paramset.params])
+        asimov_paramset.extend(nuisance_paramset.params)
+        mcmc_paramset.extend(nuisance_paramset.params)
         for parm in nuisance_paramset:
             parm.value = args.__getattribute__(parm.name)
+    tag = ParamTag.BESTFIT
     asimov_paramset.extend([
-        Param(name='astroENorm'  , value=args.measured_ratio[0], ranges=[0., 1.], std=0.2),
-        Param(name='astroMuNorm' , value=args.measured_ratio[1], ranges=[0., 1.], std=0.2),
-        Param(name='astroTauNorm', value=args.measured_ratio[2], ranges=[0., 1.], std=0.2)
+        Param(name='astroENorm'  , value=args.measured_ratio[0], ranges=[0., 1.], std=0.2, tag=tag),
+        Param(name='astroMuNorm' , value=args.measured_ratio[1], ranges=[0., 1.], std=0.2, tag=tag),
+        Param(name='astroTauNorm', value=args.measured_ratio[2], ranges=[0., 1.], std=0.2, tag=tag)
     ])
     asimov_paramset = ParamSet(asimov_paramset)
 
@@ -89,6 +93,7 @@ def get_paramsets(args):
             Param(name='c_2psi', value=0.5, ranges=[0., 1.], std=0.2, tex=r'cos(2\psi)', tag=tag)
         ])
     mcmc_paramset = ParamSet(mcmc_paramset)
+    # TODO(shivesh): unify
     return asimov_paramset, mcmc_paramset
 
 
@@ -173,9 +178,10 @@ def parse_args():
         '--outfile', type=str, default='./untitled',
         help='Path to output chains'
     )
+    llh_utils.likelihood_argparse(parser)
     mcmc_utils.mcmc_argparse(parser)
     nuisance_argparse(parser)
-    gf_argparse(parser)
+    gf_utils.gf_argparse(parser)
     plot_argparse(parser)
     return parser.parse_args()
 
@@ -187,21 +193,34 @@ def main():
 
     np.random.seed(args.seed)
 
-    asmimov_paramset, mcmc_paramset = get_paramsets(args)
+    asimov_paramset, mcmc_paramset = get_paramsets(args)
     outfile = misc_utils.gen_outfile_name(args)
     print '== {0:<25} = {1}'.format('outfile', outfile)
 
     if args.run_mcmc:
+        if args.likelihood is Likelihood.GOLEMFIT:
+            datapaths = gf.DataPaths()
+            sparams = gf_utils.steering_params(args)
+            npp = gf.NewPhysicsParams()
+
+            fitter = gf.GolemFit(datapaths, sparams, npp)
+            gf_utils.set_up_as(fitter, asimov_paramset)
+            # fitter.WriteCompact()
+        else:
+            fitter = None
+
         triangle_llh = partial(
-            mcmc_utils.triangle_llh, args=args
+            llh_utils.triangle_llh, args=args, mcmc_paramset=mcmc_paramset,
+            asimov_paramset=asimov_paramset, fitter=fitter
         )
         lnprior = partial(
-            mcmc_utils.lnprior, paramset=mcmc_paramset
+            llh_utils.lnprior, paramset=mcmc_paramset
         )
 
         ndim = len(mcmc_paramset)
-        ntemps=1
-        p0 = mcmc_utils.gaussian_seed(
+        ntemps = 1
+        # p0 = mcmc_utils.gaussian_seed(
+        p0 = mcmc_utils.flat_seed(
             mcmc_paramset, ntemps=ntemps, nwalkers=args.nwalkers
         )
 
@@ -214,7 +233,9 @@ def main():
             burnin       = args.burnin,
             nsteps       = args.nsteps,
             ntemps       = ntemps,
-            threads      = args.threads
+            threads      = 1
+            # TODO(shivesh): broken because you cannot pickle a GolemFitPy object
+            # threads      = misc_utils.thread_factors(args.threads)[0]
         )
         mcmc_utils.save_chains(samples, outfile)
 
diff --git a/utils/enums.py b/utils/enums.py
index 7798cca..bce3da2 100644
--- a/utils/enums.py
+++ b/utils/enums.py
@@ -56,17 +56,18 @@ class Priors(Enum):
 
 
 class SteeringCateg(Enum):
-    BASELINE              = 1
-    HOLEICE               = 2
-    ABSORPTION            = 3
-    SCATTERING            = 4
-    SCATTERING_ABSORPTION = 5
-    STD                   = 6
-    STD_HALF1             = 7
-    STD_HALF2             = 8
-    LONGLIFE              = 9
-    DPL                   = 10
-
+    P2_0                  = 1
+    P2_1                  = 2
+    # TODO(shivesh): fix this "P2_-1"
+    P2__1                 = 3
+    P2__3                 = 4
+    P2_0_HALF1            = 5
+    P2_0_HALF2            = 6
+    ABSORPTION            = 7
+    SCATTERING            = 8
+    SCATTERING_ABSORPTION = 9
+    LONGLIFE              = 10
+    DPL                   = 11
 
 class XSCateg(Enum):
     HALF    = 1
diff --git a/utils/gf.py b/utils/gf.py
index 99b1f24..eaa9450 100644
--- a/utils/gf.py
+++ b/utils/gf.py
@@ -10,12 +10,58 @@ Useful GolemFit wrappers for the BSM flavour ratio analysis
 from __future__ import absolute_import, division
 
 import argparse
+import socket
 from functools import partial
 
 import GolemFitPy as gf
 
 from utils.enums import *
-from utils.misc import enum_parse
+from utils.misc import enum_parse, thread_factors
+
+
+def steering_params(args):
+    steering_categ = args.ast
+    params = gf.SteeringParams()
+    if 'cobalt0' in socket.gethostname().split('.')[0]:
+        params.quiet = False
+    else:
+        params.quiet = True
+    # TODO(shivesh): figure out right number for missid
+    params.track_to_shower_missid = 0.3
+    params.fastmode = True
+    # params.fastmode = False
+    # params.readCompact = True
+    params.readCompact = False
+    params.do_HESE_reshuffle = False
+    params.sampleToLoad = gf.sampleTag.HESE
+    params.simToLoad= steering_categ.name.lower()
+    params.evalThreads = args.threads
+    # params.evalThreads = thread_factors(args.threads)[1]
+    params.baseline_astro_spectral_index = -2.
+    params.spline_hole_ice = True
+    params.spline_dom_efficiency = True
+    if steering_categ == SteeringCateg.LONGLIFE:
+        params.years = [999]
+        params.simToLoad= 'p2_0'
+    elif steering_categ == SteeringCateg.DPL:
+        params.diffuse_fit_type = gf.DiffuseFitType.DoublePowerLaw
+        params.simToLoad= 'p2_0'
+    return params
+
+
+def set_up_as(fitter, params):
+    print 'Injecting the model', params
+    asimov_params = gf.FitParameters()
+    for parm in params:
+        asimov_params.__setattr__(parm.name, parm.value)
+    fitter.SetupAsimov(asimov_params)
+
+
+def get_llh(fitter, params):
+    fitparams = gf.FitParameters()
+    for parm in params:
+        fitparams.__setattr__(parm.name, parm.value)
+    return fitter.EvalLLH(fitparams)
 
 
 def data_distributions(fitter):
@@ -84,7 +130,7 @@ def gf_argparse(parser):
         choices=DataType, help='select datatype'
     )
     parser.add_argument(
-        '--ast', default='baseline', type=partial(enum_parse, c=SteeringCateg),
+        '--ast', default='p2_0', type=partial(enum_parse, c=SteeringCateg),
         choices=SteeringCateg,
         help='use asimov/fake dataset with specific steering'
     )
diff --git a/utils/likelihood.py b/utils/likelihood.py
new file mode 100644
index 0000000..5c9af51
--- /dev/null
+++ b/utils/likelihood.py
@@ -0,0 +1,91 @@
+# author : S. Mandalia
+#          s.p.mandalia@qmul.ac.uk
+#
+# date   : April 04, 2018
+
+"""
+Likelihood functions for the BSM flavour ratio analysis
+"""
+
+from __future__ import absolute_import, division
+
+import argparse
+from functools import partial
+
+import numpy as np
+from scipy.stats import multivariate_normal
+
+import GolemFitPy as gf
+
+from utils import fr as fr_utils
+from utils import gf as gf_utils
+from utils.enums import Likelihood, ParamTag
+from utils.misc import enum_parse
+
+
+def gaussian_llh(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 likelihood_argparse(parser):
+    parser.add_argument(
+        '--likelihood', default='gaussian', type=partial(enum_parse, c=Likelihood),
+        choices=Likelihood, help='likelihood contour'
+    )
+
+
+def lnprior(theta, paramset):
+    """Priors on theta."""
+    ranges = paramset.ranges
+    for value, range in zip(theta, ranges):
+        if range[0] <= value <= range[1]:
+            pass
+        else: return -np.inf
+    return 0.
+
+
+def triangle_llh(theta, args, asimov_paramset, mcmc_paramset, fitter):
+    """-Log likelihood function for a given theta."""
+    if len(theta) != len(mcmc_paramset):
+        raise AssertionError(
+            'Length of MCMC scan is not the same as the input '
+            'params\ntheta={0}\nmcmc_paramset]{1}'.format(theta, mcmc_paramset)
+        )
+    for idx, param in enumerate(mcmc_paramset):
+        param.value = theta[idx]
+    hypo_paramset = asimov_paramset
+    for param in mcmc_paramset.from_tag(ParamTag.NUISANCE):
+        hypo_paramset[param.name].value = param.value
+
+    if args.fix_source_ratio:
+        fr1, fr2, fr3 = args.source_ratio
+    else:
+        fr1, fr2, fr3 = fr_utils.angles_to_fr(
+            mcmc_paramset.from_tag(ParamTag.SRCANGLES, values=True)
+        )
+    bsm_angles = mcmc_paramset.from_tag(
+        [ParamTag.SCALE, ParamTag.MMANGLES], values=True
+    )
+
+    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
+    )
+    fr = fr_utils.u_to_fr((fr1, fr2, fr3), u)
+    for idx, param in enumerate(hypo_paramset.from_tag(ParamTag.BESTFIT)):
+        param.value = fr[idx]
+
+    if args.likelihood is Likelihood.FLAT:
+        return 1.
+    elif args.likelihood is Likelihood.GAUSSIAN:
+        fr_bf = args.measured_ratio
+        return gaussian_llh(fr, fr_bf, args.sigma_ratio)
+    elif args.likelihood is Likelihood.GOLEMFIT:
+        return gf_utils.get_llh(fitter, hypo_paramset)
diff --git a/utils/mcmc.py b/utils/mcmc.py
index 995a338..d2036da 100644
--- a/utils/mcmc.py
+++ b/utils/mcmc.py
@@ -10,29 +10,13 @@ Useful functions to use an MCMC for the BSM flavour ratio analysis
 from __future__ import absolute_import, division
 
 import argparse
-from functools import partial
 
 import emcee
 import tqdm
 
 import numpy as np
-from scipy.stats import multivariate_normal
 
-import GolemFitPy as gf
-
-from utils import fr as fr_utils
-from utils.enums import Likelihood
-from utils.misc import enum_parse, make_dir, parse_bool
-
-
-def lnprior(theta, paramset):
-    """Priors on theta."""
-    ranges = paramset.ranges
-    for value, range in zip(theta, ranges):
-        if range[0] <= value <= range[1]:
-            pass
-        else: return -np.inf
-    return 0.
+from utils.misc import make_dir, parse_bool
 
 
 def mcmc(p0, triangle_llh, lnprior, ndim, nwalkers, burnin, nsteps, ntemps=1, threads=1):
@@ -66,10 +50,6 @@ def mcmc(p0, triangle_llh, lnprior, ndim, nwalkers, burnin, nsteps, ntemps=1, th
 
 def mcmc_argparse(parser):
     parser.add_argument(
-        '--likelihood', default='gaussian', type=partial(enum_parse, c=Likelihood),
-        choices=Likelihood, help='likelihood contour'
-    )
-    parser.add_argument(
         '--run-mcmc', type=parse_bool, default='True',
         help='Run the MCMC'
     )
@@ -87,10 +67,15 @@ def mcmc_argparse(parser):
     )
 
 
-def gaussian_llh(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 flat_seed(paramset, ntemps, nwalkers):
+    """Get gaussian seed values for the MCMC."""
+    ndim = len(paramset)
+    low = np.array(paramset.ranges).T[0]
+    high = np.array(paramset.ranges).T[1]
+    p0 = np.random.uniform(
+        low=low, high=high, size=[ntemps, nwalkers, ndim]
+    )
+    return p0
 
 
 def gaussian_seed(paramset, ntemps, nwalkers):
@@ -118,42 +103,3 @@ def save_chains(chains, outfile):
     print 'Saving chains to location {0}'.format(outfile+'.npy')
     np.save(outfile+'.npy', chains)
 
-
-def triangle_llh(theta, args):
-    """-Log likelihood function for a given theta."""
-    if args.fix_source_ratio:
-        fr1, fr2, fr3 = args.source_ratio
-        bsm_angles = theta[-5:]
-    else:
-        fr1, fr2, fr3 = fr_utils.angles_to_fr(theta[-2:])
-        bsm_angles = theta[-7:-2]
-
-    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
-    )
-
-    fr = fr_utils.u_to_fr((fr1, fr2, fr3), u)
-    if args.likelihood is Likelihood.FLAT:
-        return 1.
-    elif args.likelihood is Likelihood.GAUSSIAN:
-        fr_bf = args.measured_ratio
-        return gaussian_llh(fr, fr_bf, args.sigma_ratio)
-    elif args.likelihood is Likelihood.GOLEMFIT:
-        raise NotImplementedError('TODO')
-        from collections import namedtuple
-        datapaths = gf.DataPaths()
-        IceModels = namedtuple('IceModels', 'std_half2')
-        fitters = IceModels(*[
-            gf.GolemFit(datapaths,
-                        gf.gen_steering_params(SteeringCateg.__members__[ice]),
-                        xs_params(XSCateg.nom)) for ice in IceModels._fields])
-        for fitter in fitters:
-            fitter.SetFitParametersFlag(fit_flags(FitFlagCateg.xs))
-            fitter.SetFitPriors(fit_priors(FitPriorsCateg.xs))
-
diff --git a/utils/misc.py b/utils/misc.py
index ebc66e2..735f95b 100644
--- a/utils/misc.py
+++ b/utils/misc.py
@@ -9,7 +9,7 @@ Misc functions for the BSM flavour ratio analysis
 
 from __future__ import absolute_import, division
 
-import os
+import os, sys
 import errno
 import multiprocessing
 
@@ -112,6 +112,14 @@ class ParamSet(Sequence):
     def __iter__(self):
         return iter(self._params)
 
+    def __str__(self):
+        o = '\n'
+        for obj in self._params:
+            o += '== {0:<15} = {1:<15}, tag={2:<15}\n'.format(
+                obj.name, obj.value, obj.tag
+            )
+        return o
+
     @property
     def _by_name(self):
         return {obj.name: obj for obj in self._params}
@@ -147,26 +155,21 @@ class ParamSet(Sequence):
     def to_dict(self):
         return {obj.name: obj.value for obj in self._params}
 
-    def from_tag(self, tag):
-        return tuple([obj for obj in self._params if obj.tag is tag])
-
-    def remove_tag(self, tag):
-        return tuple([obj for obj in self._params if obj.tag is not tag])
-
-    def idx_from_tag(self, tag):
-        return tuple([idx for idx, obj in enumerate(self._params)
-                      if obj.tag is tag])
-
-    def not_idx_from_tag(self, tag):
-        return tuple([idx for idx, obj in enumerate(self._params)
-                      if obj.tag is not tag])
-
-    def slice_from_tag(self, array, tags):
-        tags = [tags]
-        idxs = []
-        for tag in tags:
-            idxs.extend(self.idx_from_tag(tag))
-        return array[idxs,]
+    def from_tag(self, tag, values=False, index=False, invert=False):
+        if values and index: assert 0
+        tag = np.atleast_1d(tag)
+        if not invert:
+            ps = [(idx, obj) for idx, obj in enumerate(self._params)
+                  if obj.tag in tag]
+        else:
+            ps = [(idx, obj) for idx, obj in enumerate(self._params)
+                  if obj.tag not in tag]
+        if values:
+            return tuple([io[1].value for io in ps])
+        elif index:
+            return tuple([io[0] for io in ps])
+        else:
+            return ParamSet([io[1] for io in ps])
 
 
 class SortingHelpFormatter(argparse.HelpFormatter):
@@ -286,3 +289,9 @@ def thread_type(t):
     else:
         return int(t)
 
+
+def thread_factors(t):
+    for x in reversed(range(int(np.ceil(np.sqrt(t)))+1)):
+        if t%x == 0:
+            return (x, int(t/x))
+
diff --git a/utils/plot.py b/utils/plot.py
index fb90b3e..a24c69c 100644
--- a/utils/plot.py
+++ b/utils/plot.py
@@ -84,7 +84,7 @@ def gen_figtext(args):
                     'observed flavour ratio = [{3:.2f}, {4:.2f}, ' \
                     '{5:.2f}]\nSigma = {6:.3f}\nDimension = {7}\nEnergy = {8} ' \
                     'GeV'.format(
-                        sr1, sr3, sr2, mr1, mr2, mr3, args.sigma_ratio,
+                        sr1, sr2, sr3, mr1, mr2, mr3, args.sigma_ratio,
                         args.dimension, int(args.energy)
                     )
     else: