intergrate GolemFit into scripts

1 files changed, 34 insertions, 13 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)